164 lines
4.1 KiB
Python
164 lines
4.1 KiB
Python
import threading
|
|
import time
|
|
import unittest
|
|
|
|
from singleflight import SingleFlight
|
|
|
|
|
|
class TestSingleFlight(unittest.TestCase):
|
|
def test_single_call_runs_fn(self):
|
|
sf = SingleFlight()
|
|
calls = {"n": 0}
|
|
|
|
def fn():
|
|
calls["n"] += 1
|
|
return 42
|
|
|
|
self.assertEqual(sf.do("k", fn), 42)
|
|
self.assertEqual(calls["n"], 1)
|
|
|
|
def test_concurrent_coalesces(self):
|
|
sf = SingleFlight()
|
|
calls = {"n": 0}
|
|
lock = threading.Lock()
|
|
barrier = threading.Barrier(8)
|
|
|
|
def fn():
|
|
with lock:
|
|
calls["n"] += 1
|
|
time.sleep(0.2)
|
|
return "result"
|
|
|
|
results = [None] * 8
|
|
errors = []
|
|
|
|
def worker(i):
|
|
barrier.wait()
|
|
try:
|
|
results[i] = sf.do("k", fn)
|
|
except BaseException as e: # noqa
|
|
errors.append(e)
|
|
|
|
ts = [threading.Thread(target=worker, args=(i,)) for i in range(8)]
|
|
for t in ts:
|
|
t.start()
|
|
for t in ts:
|
|
t.join()
|
|
self.assertFalse(errors)
|
|
self.assertTrue(all(r == "result" for r in results))
|
|
self.assertEqual(calls["n"], 1)
|
|
|
|
def test_exception_propagates_to_all(self):
|
|
sf = SingleFlight()
|
|
calls = {"n": 0}
|
|
|
|
def fn():
|
|
calls["n"] += 1
|
|
time.sleep(0.2)
|
|
raise ValueError("boom")
|
|
|
|
results = []
|
|
errors = []
|
|
barrier = threading.Barrier(5)
|
|
|
|
def worker():
|
|
barrier.wait()
|
|
try:
|
|
sf.do("k", fn)
|
|
except Exception as e: # noqa
|
|
errors.append(e)
|
|
else:
|
|
results.append("ok")
|
|
|
|
ts = [threading.Thread(target=worker) for _ in range(5)]
|
|
for t in ts:
|
|
t.start()
|
|
for t in ts:
|
|
t.join()
|
|
self.assertEqual(len(errors), 5)
|
|
self.assertTrue(all(isinstance(e, ValueError) and str(e) == "boom" for e in errors))
|
|
self.assertEqual(calls["n"], 1)
|
|
# Same exception object shared.
|
|
self.assertTrue(len({id(e) for e in errors}) == 1)
|
|
|
|
def test_subsequent_call_runs_again(self):
|
|
sf = SingleFlight()
|
|
calls = {"n": 0}
|
|
|
|
def fn():
|
|
calls["n"] += 1
|
|
return calls["n"]
|
|
|
|
self.assertEqual(sf.do("k", fn), 1)
|
|
self.assertEqual(sf.do("k", fn), 2)
|
|
self.assertEqual(calls["n"], 2)
|
|
|
|
def test_different_keys_independent(self):
|
|
sf = SingleFlight()
|
|
calls = {"a": 0, "b": 0}
|
|
|
|
def make(key):
|
|
def fn():
|
|
calls[key] += 1
|
|
time.sleep(0.1)
|
|
return key
|
|
return fn
|
|
|
|
r = {}
|
|
def wa():
|
|
r["a"] = sf.do("a", make("a"))
|
|
def wb():
|
|
r["b"] = sf.do("b", make("b"))
|
|
ta = threading.Thread(target=wa); tb = threading.Thread(target=wb)
|
|
ta.start(); tb.start(); ta.join(); tb.join()
|
|
self.assertEqual(r["a"], "a")
|
|
self.assertEqual(r["b"], "b")
|
|
self.assertEqual(calls["a"], 1)
|
|
self.assertEqual(calls["b"], 1)
|
|
|
|
def test_inflight_count(self):
|
|
sf = SingleFlight()
|
|
started = threading.Event()
|
|
release = threading.Event()
|
|
|
|
def fn():
|
|
started.set()
|
|
release.wait()
|
|
return 1
|
|
|
|
def leader():
|
|
sf.do("k", fn)
|
|
|
|
t = threading.Thread(target=leader)
|
|
t.start()
|
|
started.wait()
|
|
self.assertEqual(sf.inflight_count, 1)
|
|
release.set()
|
|
t.join()
|
|
self.assertEqual(sf.inflight_count, 0)
|
|
|
|
def test_leader_crash_does_not_leak_key(self):
|
|
sf = SingleFlight()
|
|
|
|
def fn():
|
|
raise RuntimeError("crash")
|
|
|
|
try:
|
|
sf.do("k", fn)
|
|
except RuntimeError:
|
|
pass
|
|
# Key freed: next call runs fn again.
|
|
calls = {"n": 0}
|
|
|
|
def fn2():
|
|
calls["n"] += 1
|
|
return "ok"
|
|
|
|
self.assertEqual(sf.do("k", fn2), "ok")
|
|
self.assertEqual(calls["n"], 1)
|
|
self.assertEqual(sf.inflight_count, 0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|