memoize: make thread-safe

Technologicat · Technologicat · commit f1551f467145 · 2021-06-16T13:37:23.000+03:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -204,6 +204,8 @@ The same applies if you need the macro parts of `unpythonic` (i.e. import anythi
 
 - Fix bug: `fup`/`fupdate`/`ShadowedSequence` now actually accept an infinite-length iterable as a replacement sequence (under the obvious usage limitations), as the documentation has always claimed.
 
+- Fix bug: `memoize` is now thread-safe.
+
 
 ---
 
diff --git a/unpythonic/fun.py b/unpythonic/fun.py
@@ -20,6 +20,7 @@
 from collections import namedtuple
 from functools import wraps, partial as functools_partial
 from inspect import signature
+from threading import RLock
 from typing import get_type_hints
 
 from .arity import (_resolve_bindings, tuplify_bindings, _bind)
@@ -61,18 +62,35 @@ def memoize(f):
 
     **CAUTION**: ``f`` must be pure (no side effects, no internal state
     preserved between invocations) for this to make any sense.
+
+    Beginning with v0.15.0, `memoize` is thread-safe even when the same memoized
+    function instance is called concurrently from multiple threads. Exactly one
+    thread will compute the result. If `f` is recursive, the thread that acquired
+    the lock is the one that is allowed to recurse into the memoized `f`.
     """
+    # One lock per use site of `memoize`. We use an `RLock` to allow recursive calls
+    # to the memoized `f` in the thread that acquired the lock.
+    lock = RLock()
     memo = {}
     @wraps(f)
     def memoized(*args, **kwargs):
         k = tuplify_bindings(_resolve_bindings(f, args, kwargs, _partial=False))
-        if k not in memo:
-            try:
-                result = (_success, maybe_force_args(f, *args, **kwargs))
-            except BaseException as err:
-                result = (_fail, err)
-            memo[k] = result  # should yell separately if k is not a valid key
-        kind, value = memo[k]
+        try:  # EAFP to eliminate TOCTTOU.
+            kind, value = memo[k]
+        except KeyError:
+            # But we still need to be careful to avoid race conditions.
+            with lock:
+                if k not in memo:
+                    # We were the first thread to acquire the lock.
+                    try:
+                        result = (_success, maybe_force_args(f, *args, **kwargs))
+                    except BaseException as err:
+                        result = (_fail, err)
+                    memo[k] = result  # should yell separately if k is not a valid key
+                else:
+                    # Some other thread acquired the lock before us.
+                    pass
+            kind, value = memo[k]
         if kind is _fail:
             raise value
         return value
diff --git a/unpythonic/tests/test_fun.py b/unpythonic/tests/test_fun.py
@@ -5,6 +5,9 @@
 
 from collections import Counter
 import sys
+from queue import Queue
+import threading
+from time import sleep
 
 from ..dispatch import generic
 from ..fun import (memoize, partial, curry, apply,
@@ -16,6 +19,8 @@
                    to1st, to2nd, tokth, tolast, to,
                    withself)
 from ..funutil import Values
+from ..it import allsame
+from ..misc import slurp
 
 from ..dynassign import dyn
 
@@ -135,6 +140,36 @@ def t():
                 fail["memoize should not prevent exception propagation."]  # pragma: no cover
         test[evaluations == 1]
 
+    with testset("@memoize thread-safety"):
+        def threadtest():
+            @memoize
+            def f(x):
+                # Sleep a "long" time to make actual concurrent operation more likely.
+                sleep(0.001)
+
+                # The trick here is that because only one thread will acquire the lock
+                # for the memo, then for the same `x`, all the results should be the same.
+                return (id(threading.current_thread()), x)
+
+            comm = Queue()
+            def worker(que):
+                # The value of `x` doesn't matter, as long as it's the same in all workers.
+                r = f(42)
+                que.put(r)
+
+            n = 1000
+            threads = [threading.Thread(target=worker, args=(comm,), kwargs={}) for _ in range(n)]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join()
+
+            # Test that all threads finished, and that the results from each thread are the same.
+            results = slurp(comm)
+            test[the[len(results)] == the[n]]
+            test[allsame(results)]
+        threadtest()
+
     with testset("partial (type-checking wrapper)"):
         def nottypedfunc(x):
             return "ok"
