Diffracting tree in Python
Binary tree of balancers with prism elimination that spreads tokens to counter leaves with low contention.
- File
diffracting_tree.py, 173 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/diffracting-tree/. The authors are listed in prov/8G5PTB.jsonl.
import random
import threading
import time
from collections.abc import Generator
from typing import TypeVar
R = TypeVar("R")
SPIN = "spin" # what a step generator yields while it waits for a partner
TIMEOUT = object() # the result of an exchange that found no partner
EMPTY, WAITING, BUSY = "EMPTY", "WAITING", "BUSY"
Steps = Generator[str | None, None, R]
def run(steps: "Steps[R]") -> R:
"""Run a step generator to its end in the calling thread and return its result."""
while True:
try:
if next(steps) == SPIN:
time.sleep(0)
except StopIteration as stop:
return stop.value
class Atomic:
"""A word with atomic get, set and compare-and-swap, simulated with a lock."""
def __init__(self, value: object) -> None:
self._value = value
self._lock = threading.Lock()
def get(self) -> object:
with self._lock:
return self._value
def set(self, value: object) -> None:
with self._lock:
self._value = value
def cas(self, expect: object, new: object) -> bool:
with self._lock:
if self._value == expect:
self._value = new
return True
return False
def fetch_and_add(self, k: int) -> int:
with self._lock:
old = self._value
self._value = old + k # type: ignore[operator]
return old # type: ignore[return-value]
def exchange_steps(slot: Atomic, me: int, patience: int) -> "Steps[object]":
"""The lock-free exchanger of Herlihy and Shavit: the partner's id, or TIMEOUT.
The slot holds a pair (item, state). Each access is one step.
"""
for _ in range(patience): # @line L12
item, state = slot.get() # type: ignore[misc] # @line L13
yield None
if state == EMPTY:
ok = slot.cas((item, EMPTY), (me, WAITING)) # @line L14
yield None
if not ok:
continue
for _ in range(patience): # @line L15
item, state = slot.get() # type: ignore[misc]
yield SPIN
if state == BUSY:
slot.set((None, EMPTY))
yield None
return item
ok = slot.cas((me, WAITING), (None, EMPTY)) # @line L16
yield None
if ok:
return TIMEOUT
item, _ = slot.get() # type: ignore[misc] # @line L17
slot.set((None, EMPTY))
yield None
return item
if state == WAITING:
ok = slot.cas((item, WAITING), (me, BUSY)) # @line L18
yield None
if ok:
return item
else:
yield SPIN # BUSY: another pair is finishing
return TIMEOUT # @line L19
class Balancer:
"""A toggle bit with a prism of exchangers in front of it."""
def __init__(self, prism_size: int) -> None:
self.toggle = Atomic(0)
self.prism = [Atomic((None, EMPTY)) for _ in range(prism_size)]
self.toggled = [0, 0] # tokens that left through the toggle, per side
self.diffracted = 0 # pairs that left through the prism
self.stats = threading.Lock() # guards the two statistics above
class DiffractingTree:
"""A diffracting tree counter of Shavit and Zemach, in the form of Herlihy and Shavit.
width = 2^k output wires; a binary tree of width - 1 balancers sends each
token to one wire, and wire i hands out i, i + width, i + 2 width, ...
Two tokens that meet in the prism of a balancer leave on different sides
without touching its toggle. The counter is quiescently consistent and
lock-free, not linearizable. Callers pass distinct ids.
"""
def __init__(self, width: int, prism_size: int = 2, patience: int = 3, seed: int = 0) -> None:
if width < 2 or width & (width - 1):
raise ValueError("width must be a power of two, at least 2")
self.width = width
self.patience = patience
self.balancers: list[Balancer] = []
depth = 0
while len(self.balancers) < width - 1:
size = max(1, prism_size >> depth)
self.balancers.extend(Balancer(size) for _ in range(1 << depth))
depth += 1
self.count = [Atomic(0) for _ in range(width)]
self.rng = random.Random(seed)
self.rng_lock = threading.Lock()
def _balance(self, b: Balancer, me: int) -> "Steps[int]":
with self.rng_lock:
i = self.rng.randrange(len(b.prism)) # @line L7
partner = yield from exchange_steps(b.prism[i], me, self.patience) # @line L8
if partner is not TIMEOUT: # @line L9
d = 0 if me < partner else 1 # type: ignore[operator]
if d == 0:
with b.stats:
b.diffracted += 1
return d
while True: # @line L10
t = b.toggle.get()
yield None
ok = b.toggle.cas(t, 1 - t) # type: ignore[operator]
yield None
if ok:
with b.stats:
b.toggled[t] += 1 # type: ignore[index]
return t # @line L11
def get_and_increment_steps(self, me: int) -> "Steps[int]":
j, wire, bit = 0, 0, 1 # @line L1
while j < self.width - 1: # @line L2
d = yield from self._balance(self.balancers[j], me) # @line L3
wire, bit, j = wire + d * bit, 2 * bit, 2 * j + 1 + d # @line L4
k = self.count[wire].fetch_and_add(1) # @line L5
yield None
return wire + self.width * k # type: ignore[operator] # @line L6
def get_and_increment(self, me: int) -> int:
return run(self.get_and_increment_steps(me))
def wire_counts(self) -> list[int]:
return [c.get() for c in self.count] # type: ignore[misc]
def check_step_property(self) -> None:
"""At a quiet moment: wire counts fall by at most 1 from left to right, and so do balancers."""
y = self.wire_counts()
assert all(0 <= y[i] - y[j] <= 1 for i in range(len(y)) for j in range(i + 1, len(y)))
for b in self.balancers:
assert b.toggled[0] - b.toggled[1] == b.toggle.get()
assert all(s.get() == (None, EMPTY) for s in b.prism)
test_diffracting_tree.py 199 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/diffracting-tree/. The authors are listed in prov/8G5PTB.jsonl.
import random
import threading
from diffracting_tree import EMPTY, TIMEOUT, WAITING, Atomic, DiffractingTree, exchange_steps, run
# ---- deterministic interleaving driver and linearizability check ----
Op = tuple # (thread, name, arg, result, invoke time, response time)
def interleave(threads, schedule, limit=100000):
"""Run the operations of each thread, one atomic step at a time.
threads[t] is a list of (name, arg, make) where make() returns the step
generator of the operation. schedule is a sequence of thread numbers; a
pick of a thread with no work left is skipped. After the schedule, the
remaining steps run round robin. Returns the history as a list of Op.
"""
todo = [list(reversed(ops)) for ops in threads]
cur: list = [None] * len(threads)
history: list = []
clock = 0
def step(t):
nonlocal clock
if cur[t] is None:
if not todo[t]:
return False
name, arg, make = todo[t].pop()
cur[t] = (name, arg, make(), clock)
name, arg, gen, inv = cur[t]
clock += 1
try:
next(gen)
except StopIteration as stop:
history.append((t, name, arg, stop.value, inv, clock))
cur[t] = None
return True
for t in schedule:
step(t)
steps = 0
while any(cur) or any(todo):
for t in range(len(threads)):
step(t)
steps += 1
assert steps < limit, "no progress"
return history
def linearizable(history, init, apply):
"""Wing and Gong search: is there a legal order that respects real time?
apply(state, name, arg) returns (new state, result); states are hashable.
"""
n = len(history)
full = (1 << n) - 1
seen = set()
def dfs(done, state):
if done == full:
return True
if (done, state) in seen:
return False
seen.add((done, state))
for i in range(n):
if done >> i & 1:
continue
inv_i = history[i][4]
if any(not done >> j & 1 and history[j][5] < inv_i for j in range(n)):
continue
s2, r = apply(state, history[i][1], history[i][2])
if r == history[i][3] and dfs(done | 1 << i, s2):
return True
return False
return dfs(0, init)
def counter_spec(n, name, arg):
return n + 1, n
def quiescently_consistent(history, start):
"""For a counter: at each quiescent point the finished calls hold exactly start .. start + m - 1."""
values = [r for _, _, _, r, _, _ in history]
if len(set(values)) != len(values):
return False
times = sorted({h[5] for h in history} | {h[4] for h in history})
for t in times:
if any(h[4] < t < h[5] for h in history):
continue # some call is pending at t
done = sorted(h[3] for h in history if h[5] <= t)
if done != list(range(start, start + len(done))):
return False
return True
def make_ops(tree, me, count):
return [("inc", me, lambda: tree.get_and_increment_steps(me)) for _ in range(count)]
def test_example() -> None:
# Width 2, one balancer. B passes the toggle and stalls; A then calls twice.
tree = DiffractingTree(2, prism_size=1, patience=1)
b = tree.get_and_increment_steps(1)
while tree.balancers[0].toggle.get() == 0:
next(b) # B times out in the prism and takes side 0 of the toggle
a1 = tree.get_and_increment(0)
a2 = tree.get_and_increment(0)
vb = run(b)
assert (a1, a2, vb) == (1, 0, 2) # A got 1 and then 0: not linearizable
h = [(0, "inc", 0, 1, 10, 20), (0, "inc", 0, 0, 30, 40), (1, "inc", 1, 2, 0, 50)]
assert not linearizable(h, 0, counter_spec)
assert quiescently_consistent(h, 0)
tree.check_step_property()
def test_diffraction() -> None:
tree = DiffractingTree(2, prism_size=1, patience=4)
a = tree.get_and_increment_steps(5)
b = tree.get_and_increment_steps(9)
next(a) # A reads the slot: EMPTY
next(a) # A installs itself: WAITING
assert tree.balancers[0].prism[0].get() == (5, WAITING)
next(b) # B reads the slot
next(b) # B swaps in: BUSY, B leaves with partner 5 on side 1
assert run(b) == 1 # side 1, wire 1, first value of the wire
assert run(a) == 0 # A sees BUSY, takes 9, leaves on side 0
assert tree.balancers[0].toggle.get() == 0 # the toggle was not touched
assert tree.balancers[0].diffracted == 1
tree.check_step_property()
def test_exchanger_alone_times_out() -> None:
slot = Atomic((None, EMPTY))
assert run(exchange_steps(slot, 3, 2)) is TIMEOUT
assert slot.get() == (None, EMPTY)
def test_edges() -> None:
tree = DiffractingTree(2)
assert [tree.get_and_increment(0) for _ in range(5)] == [0, 1, 2, 3, 4]
for bad in (1, 3, 12):
try:
DiffractingTree(bad)
except ValueError:
continue
raise AssertionError("bad width accepted")
def test_random_against_model() -> None:
rng = random.Random(4711)
for _ in range(500):
width = rng.choice([2, 4, 8, 16])
tree = DiffractingTree(width, prism_size=rng.randint(1, 4), patience=1, seed=rng.randrange(99))
k = rng.randint(0, 40)
assert [tree.get_and_increment(rng.randrange(8)) for _ in range(k)] == list(range(k))
tree.check_step_property()
def test_random_interleavings_are_quiescently_consistent() -> None:
rng = random.Random(90210)
not_linearizable = 0
for _ in range(500):
width = rng.choice([2, 4, 8])
tree = DiffractingTree(width, prism_size=rng.randint(1, 3), patience=rng.randint(1, 3), seed=rng.randrange(99))
start = rng.randint(0, 3)
for _ in range(start):
tree.get_and_increment(0)
threads = [make_ops(tree, t, rng.randint(1, 3)) for t in range(3)]
schedule = [rng.randrange(3) for _ in range(rng.randint(0, 80))]
h = interleave(threads, schedule)
assert quiescently_consistent(h, start)
if not linearizable(h, start, counter_spec):
not_linearizable += 1
tree.check_step_property()
assert not_linearizable > 0 # the counter is really weaker than linearizable
def test_threads() -> None:
width, per, n = 8, 200, 6
tree = DiffractingTree(width, prism_size=4, patience=8, seed=7)
got: list[list[int]] = [[] for _ in range(n)]
def work(t: int) -> None:
for _ in range(per):
got[t].append(tree.get_and_increment(t))
ts = [threading.Thread(target=work, args=(t,)) for t in range(n)]
for t in ts:
t.start()
for t in ts:
t.join()
assert sorted(x for g in got for x in g) == list(range(n * per))
tree.check_step_property()
What the next level needs
- the tests did not run in this build