ops: refuse same-key writes from other threads during an operation; batch failure notes (#5)

cosmodol/ops.py

@@ -51,10 +51,16 @@
 - A layer that performs the write in another thread or later (a write-behind buffer)
   escapes the redirect: that write is the leaf's ordinary upsert or delete, without the
   operation's semantics (no ETag check, outside a batch's transaction). The function
-  raises ``RedirectError``, but cannot undo it.
+  raises ``RedirectError``, but cannot undo it. A write of the *same key* from another
+  thread or task while the operation is pending is refused (``ConcurrentWriteError``, a
+  ``RedirectError``), so a layer that writes from a thread the operation waits for cannot
+  slip past it. What escapes: a write that lands after the operation has finished, and a
+  write through a different leaf object on the same container (``partition()`` child, a
+  second ``CosmosItems``).
 - ``batch_items`` collects its writes through the layers first and commits them as one
   transaction afterwards: a layer that acts after the inner write returns (a write-through
-  cache, an audit log) has acted even if the transaction then fails.
+  cache, an audit log) has acted even if the transaction then fails. The error that reaches
+  the caller carries a note (``__notes__``) naming the keys the layers had acted on.
 - ``batch_items``'s ``patch`` operation has no mapping-write equivalent, so no layer's
   write logic can run for it; on a wrapped (or subclassed) store it is refused unless
   ``unguarded_patch=True``. Its key is still mapped and scope-checked.
@@ -232,7 +238,7 @@ def batch_items(
     opaque to any key mapping; use ``cosmodol.batch(container, triples, partition_key)``.
     """
     leaf = _leaf_of(store, CosmosItems, fn="batch_items")
-    triples, created = [], set()
+    triples, created, acted = [], set(), []
 
     def collect_write(leaf_k, leaf_v):
         return leaf_k, leaf._prepare_body(leaf_k, leaf_v)
@@ -278,6 +284,8 @@ def collect_key(leaf_k):
         elif k not in created:
             _require_member(store, leaf, k, ItemNotFoundError, CosmosItems)
         what = f"batch_items {op!r}"
+        if op in ("create", "upsert", "replace", "delete"):
+            acted.append(k)
         if op in ("create", "upsert", "replace"):
             leaf_k, body = _routed(
                 store, leaf, "__setitem__", collect_write, k, rest[0], what=what
@@ -289,7 +297,21 @@ def collect_key(leaf_k):
             leaf_k = collect_key(_leaf_key(store, leaf, k))
             args = (leaf_k, rest[0]) if op == "patch" else (leaf_k,)
         triples.append((op, args, kwargs))
-    return _batch(leaf.container, triples, leaf.partition_key_value)
+    try:
+        return _batch(leaf.container, triples, leaf.partition_key_value)
+    except Exception as e:
+        if acted and _write_logic_above(store, leaf, CosmosItems):
+            note = (
+                "batch_items: the batch may not have committed, but the wrapping layers had "
+                f"already run their write logic for these keys: {acted!r}. A write-through "
+                "cache, audit log or index they keep may now disagree with the container "
+                "(i2mint/cosmodol#5)."
+            )
+            if hasattr(e, "add_note"):
+                e.add_note(note)
+            else:  # Python 3.10: the same attribute, which only newer tracebacks print
+                e.__notes__ = [*getattr(e, "__notes__", []), note]
+        raise
 
 
 def partition(store, pk_value, *, wrap: Optional[Callable] = None) -> Any:

cosmodol/routing.py

@@ -26,13 +26,40 @@
 operation's semantics applied to it: that write is refused instead.
 
 The redirect is held in a ``ContextVar``, so it is invisible to other threads and tasks, and
-it only applies to the leaf object it names. The flip side: a layer that performs the write
-in another thread, or later (a write-behind buffer), escapes it, and that write happens as
-the leaf's ordinary upsert or delete.
+it only applies to the leaf object it names. The flip side is a layer that performs the
+write in another thread, or later (a write-behind buffer). Two things are done about it
+(i2mint/cosmodol#5):
+
+- While an operation is pending, a write of *the same leaf key* arriving from a context that
+  has no redirect (another thread, or a task that did not copy the context) is refused with
+  :class:`ConcurrentWriteError`, before it acts, instead of running as an ordinary upsert or
+  delete outside the operation's checks. Writes of other keys are not affected, and neither
+  is a write through a *different leaf object* on the same container (a ``partition()``
+  child, a second ``CosmosItems``): the registry matches leaves by identity.
+- A write performed *after* the operation has finished is indistinguishable from any other
+  writer's, so it cannot be refused: the operation raises :class:`RedirectError` when the
+  write it asked for never reached the leaf, which is the only signal there is.
+
+>>> import threading
+>>> started, release = threading.Event(), threading.Event()
+>>> def slow(k, v):
+...     started.set(); release.wait(5); return "done"
+>>> def op():
+...     with redirect(leaf, "__setitem__", "a", slow):
+...         leaf["a"] = 1
+>>> t = threading.Thread(target=op); t.start(); _ = started.wait(5)
+>>> try:
+...     leaf["a"] = 2  # same key, other thread, while the operation is pending
+... except ConcurrentWriteError as e:
+...     print("refused")
+refused
+>>> leaf["b"] = 3  # another key is fine
+>>> release.set(); t.join()
 """
 
 from __future__ import annotations
 
+import threading
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass, field
@@ -43,12 +70,16 @@ class RedirectError(RuntimeError):
     """A wrapping layer did not turn one mapping write into exactly one leaf write."""
 
 
+class ConcurrentWriteError(RedirectError):
+    """A write of a key arrived from another thread or task while an operation on it was pending."""
+
+
 #: The leaf methods a redirect guards: while one is pending on a leaf, any of them called
 #: with anything but the expected method and key is refused.
 WRITE_METHODS = ("__setitem__", "__delitem__")
 
 
-@dataclass
+@dataclass(eq=False)  # compared by identity: the registry removes by it, and a leaf's == reads the store
 class Redirect:
     """One pending redirect of ``leaf.<method>(key, ...)`` to ``handler``.
 
@@ -68,6 +99,11 @@ class Redirect:
 
 _active: ContextVar[Optional[Redirect]] = ContextVar("cosmodol_redirect", default=None)
 
+# Every redirect pending in any thread, so that a write from a context with no redirect can be
+# checked against them (a ContextVar is invisible to other threads).
+_pending: list[Redirect] = []
+_pending_lock = threading.Lock()
+
 
 @contextmanager
 def redirect(leaf, method: str, key, handler: Callable) -> Iterator[Redirect]:
@@ -79,10 +115,37 @@ def redirect(leaf, method: str, key, handler: Callable) -> Iterator[Redirect]:
     """
     r = Redirect(leaf, method, key, handler)
     token = _active.set(r)
+    with _pending_lock:
+        _pending.append(r)
     try:
         yield r
     finally:
-        _active.reset(token)
+        try:
+            with _pending_lock:
+                _pending.remove(r)
+        finally:
+            _active.reset(token)
+
+
+def _refuse_if_pending(leaf, method: str, key) -> None:
+    """Refuse a write of ``key`` on ``leaf`` while another context has an operation pending on it."""
+    with _pending_lock:
+        pending = next(
+            (p for p in _pending if p.leaf is leaf and _same_key(p.key, key)), None
+        )
+    if pending is not None:
+        raise ConcurrentWriteError(
+            f"{method}({key!r}) on the cosmodol store was refused: another thread, task or "
+            f"redirect has an operation ({pending.method} of {pending.key!r}) pending on the same "
+            "key. Retry once it has finished (i2mint/cosmodol#5)."
+        )
+
+
+def _same_key(a, b) -> bool:
+    try:
+        return bool(a == b)
+    except Exception:  # an exotic key type: treat as different
+        return False
 
 
 def claim(leaf, method: str, key) -> Optional[Callable]:
@@ -93,6 +156,7 @@ def claim(leaf, method: str, key) -> Optional[Callable]:
     """
     r = _active.get()
     if r is None or r.leaf is not leaf:
+        _refuse_if_pending(leaf, method, key)
         return None
     if r.claimed:
         outcome = "failed" if r.error is not None else "was performed"

cosmodol/tests/test_ops_layers.py

@@ -266,3 +266,123 @@ def __contains__(self, k):
     with pytest.raises(KeyError):
         replace_item(store, "doc1", {"payload": "X"})
     assert _state(c)["doc1"] == "SOMEONE ELSE'S"
+
+
+# ---------------------------------------------------------------------------
+# Limits of the redirect (i2mint/cosmodol#5)
+# ---------------------------------------------------------------------------
+
+
+class _ThreadedWrite(Store):
+    """A write-behind-ish layer: it writes through a worker thread (the redirect's ContextVar
+    does not follow it) and waits for it, so the write happens while the operation is pending."""
+
+    def __init__(self, store):
+        super().__init__(store)
+        self.thread_error = None
+
+    def __setitem__(self, k, v):
+        import threading
+
+        def work():
+            try:
+                self.store[k] = v
+            except BaseException as e:  # noqa: BLE001
+                self.thread_error = e
+
+        t = threading.Thread(target=work)
+        t.start()
+        t.join()
+
+
+def test_a_same_key_write_from_another_thread_during_an_operation_is_refused():
+    from cosmodol.routing import ConcurrentWriteError
+
+    c = _container()
+    before = _state(c)
+    store = _ThreadedWrite(_leaf(c))
+    with pytest.raises(RedirectError):  # the write the operation asked for never arrived
+        replace_item(store, "app_doc1", {"payload": "X"})
+    assert isinstance(store.thread_error, ConcurrentWriteError)
+    assert _state(c) == before  # the thread's write did not run as an ordinary upsert
+
+
+def test_writes_of_other_keys_and_after_the_operation_are_not_refused():
+    import threading
+
+    from cosmodol.routing import redirect
+
+    c = _container()
+    leaf = _leaf(c)
+    started, release = threading.Event(), threading.Event()
+
+    def slow(k, v):
+        started.set()
+        release.wait(5)
+
+    def op():
+        with redirect(leaf, "__setitem__", "app_doc1", slow):
+            leaf["app_doc1"] = {"payload": "X"}
+
+    t = threading.Thread(target=op)
+    t.start()
+    assert started.wait(5)
+    errors = []
+
+    def other():
+        try:
+            leaf["elsewhere"] = {"payload": "O"}
+        except BaseException as e:  # noqa: BLE001
+            errors.append(e)
+
+    t2 = threading.Thread(target=other)
+    t2.start()
+    t2.join()
+    release.set()
+    t.join()
+    assert errors == []
+    leaf["app_doc1"] = {"payload": "after"}  # the window has closed
+    assert _state(c)["app_doc1"] == "after" and _state(c)["elsewhere"] == "O"
+
+
+def test_a_failed_batch_names_the_keys_the_layers_already_acted_on():
+    c = _container()
+    store = _WriteThrough(_leaf(c))
+    with pytest.raises(Exception) as exc:
+        batch_items(
+            store,
+            [
+                ("upsert", "app_new", {"payload": "N"}),
+                ("replace", "app_doc1", {"payload": "E"}, {"if_match_etag": "stale"}),
+            ],
+        )
+    notes = getattr(exc.value, "__notes__", [])
+    assert any("app_new" in n and "app_doc1" in n for n in notes)
+    assert "app_doc1" in store.cache  # the layer did act, as the note says
+
+
+def test_pending_redirects_are_compared_by_identity_not_by_store_contents():
+    from cosmodol.routing import _pending, redirect
+
+    a, b = _leaf(_container()), _leaf(_container())  # equal contents
+    assert a == b
+    with redirect(a, "__setitem__", "k", lambda k, v: None):
+        with redirect(b, "__setitem__", "k", lambda k, v: None):
+            pass
+        assert [p.leaf for p in _pending] == [a]  # only the inner one went
+    assert _pending == []
+
+
+def test_a_failed_batch_on_a_subclassed_leaf_also_gets_the_note():
+    class Cached(CosmosItems):
+        def __setitem__(self, k, v):
+            super().__setitem__(k, v)
+
+    c = _container()
+    leaf = Cached(c, partition_key_value=PK, partition_key_path="/_pk")
+    with pytest.raises(Exception) as exc:
+        batch_items(
+            leaf,
+            [("replace", "app_doc1", {"payload": "E"}, {"if_match_etag": "stale"})],
+        )
+    assert any("app_doc1" in n for n in getattr(exc.value, "__notes__", []))