Fix P1 #30 STM lost-wakeup race; harden atom/lru-cache

ober

25cbfac18cd4812630bcaaa195dcf220a691186a

diff --git a/atom.ss b/atom.ss
index 1046ce6..cd785f9 100644
--- a/atom.ss
+++ b/atom.ss
@@ -21,50 +21,40 @@
 (defrules with-mutex ()
   ((_ mx body ...) (call-with-mutex mx (lambda () body ...))))
 
-(def +jc-atom-tag+ 'jerboa-compat-atom)
+;; A proper struct (not a tagged vector) so atom? is an exact type predicate:
+;; an arbitrary 3-vector carrying the old tag can no longer pass as an atom.
+(defstruct jc-atom (value mutex))
 
 (def (atom initial-value)
-  (vector +jc-atom-tag+ initial-value (make-mutex 'jerboa-compat-atom)))
+  (make-jc-atom initial-value (make-mutex 'jerboa-compat-atom)))
 
-(def (atom? x)
-  (and (vector? x)
-       (= (vector-length x) 3)
-       (eq? (vector-ref x 0) +jc-atom-tag+)))
+(def atom? jc-atom?)
 
 (def (check-atom who a)
   (unless (atom? a)
     (error who "not an atom" a)))
 
-(def (atom-value a)
-  (vector-ref a 1))
-
-(def (atom-value-set! a value)
-  (vector-set! a 1 value))
-
-(def (atom-mutex a)
-  (vector-ref a 2))
-
 (def (atom-deref a)
   (check-atom 'atom-deref a)
-  (with-mutex (atom-mutex a)
-    (atom-value a)))
+  (with-mutex (jc-atom-mutex a)
+    (jc-atom-value a)))
 
 (def (atom-reset! a value)
   (check-atom 'atom-reset! a)
-  (with-mutex (atom-mutex a)
-    (atom-value-set! a value))
+  (with-mutex (jc-atom-mutex a)
+    (jc-atom-value-set! a value))
   value)
 
 (def (atom-swap! a proc)
   (check-atom 'atom-swap! a)
-  (with-mutex (atom-mutex a)
-    (let ((value (proc (atom-value a))))
-      (atom-value-set! a value)
+  (with-mutex (jc-atom-mutex a)
+    (let ((value (proc (jc-atom-value a))))
+      (jc-atom-value-set! a value)
       value)))
 
 (def (atom-update! a proc . args)
   (check-atom 'atom-update! a)
-  (with-mutex (atom-mutex a)
-    (let ((value (apply proc (atom-value a) args)))
-      (atom-value-set! a value)
+  (with-mutex (jc-atom-mutex a)
+    (let ((value (apply proc (jc-atom-value a) args)))
+      (jc-atom-value-set! a value)
       value)))
diff --git a/lru-cache.ss b/lru-cache.ss
index a9d994d..a669375 100644
--- a/lru-cache.ss
+++ b/lru-cache.ss
@@ -25,7 +25,9 @@
 (def (make-lru-cache cap)
   (unless (and (integer? cap) (> cap 0))
     (error 'make-lru-cache "capacity must be a positive integer" cap))
-  (make-lruc cap (make-hash-table) #f #f 0 0 0 (make-mutex 'lru-cache)))
+  ;; equal? keys: equal?-but-not-eq? keys (e.g. fresh strings) must hit the same
+  ;; slot, or lookups phantom-miss and the table grows unbounded past capacity.
+  (make-lruc cap (make-hash-table test: equal?) #f #f 0 0 0 (make-mutex 'lru-cache)))
 
 (def (lru-cache? x) (lruc? x))
 (def (lru-cache-size c) (lruc-size c))
@@ -98,7 +100,7 @@
 
 (def (lru-cache-clear! c)
   (with-mutex (lruc-mutex c)
-    (lruc-table-set! c (make-hash-table))
+    (lruc-table-set! c (make-hash-table test: equal?))
     (lruc-head-set! c #f)
     (lruc-tail-set! c #f)
     (lruc-size-set! c 0)))
diff --git a/stm.ss b/stm.ss
index 9dc192b..c077d6d 100644
--- a/stm.ss
+++ b/stm.ss
@@ -134,26 +134,54 @@
           (for-each (lambda (e) (mutex-unlock! (tvx-lock (car e))))
                     (reverse sorted-locks)))))))
 
+(def (stm-unregister-waiter! reads w)
+  (for-each (lambda (e)
+              (let* ((tv (car e)) (mx (tvx-lock tv)))
+                (mutex-lock! mx)
+                (tvx-waiters-set! tv (remq-eq w (tvx-waiters tv)))
+                (mutex-unlock! mx)))
+            reads))
+
+;; Register a waiter on every read TVar and, while still holding each TVar's
+;; lock, re-validate its recorded version. A version that moved since the read
+;; means a writer committed inside the window between the failed attempt and
+;; waiter registration; its notify-waiters! ran before we were on the list, so
+;; blocking now would lose that wakeup forever. In that case unregister and
+;; return #f so the caller retries immediately. Return #t only after actually
+;; blocking on a delivered notification.
 (def (stm-wait-on-reads! tx)
   (let ((reads (txr-reads tx)) (w (make-thread-waiter)))
-    (for-each (lambda (e)
-                (let* ((tv (car e)) (mx (tvx-lock tv)))
-                  (mutex-lock! mx)
-                  (tvx-waiters-set! tv (cons w (tvx-waiters tv)))
-                  (mutex-unlock! mx)))
-              reads)
-    (waiter-wait! w)
-    (for-each (lambda (e)
-                (let* ((tv (car e)) (mx (tvx-lock tv)))
-                  (mutex-lock! mx)
-                  (tvx-waiters-set! tv (remq-eq w (tvx-waiters tv)))
-                  (mutex-unlock! mx)))
-              reads)))
+    (and (not (null? reads))
+         (let (stale?
+                (let loop ((rs reads) (stale? #f))
+                  (if (null? rs) stale?
+                    (let* ((e (car rs)) (tv (car e)) (mx (tvx-lock tv)))
+                      (mutex-lock! mx)
+                      (tvx-waiters-set! tv (cons w (tvx-waiters tv)))
+                      (let (moved? (not (= (tvx-version tv) (cadr e))))
+                        (mutex-unlock! mx)
+                        (loop (cdr rs) (or stale? moved?)))))))
+           (if stale?
+             (begin (stm-unregister-waiter! reads w) #f)
+             (begin (waiter-wait! w) (stm-unregister-waiter! reads w) #t))))))
+
+;; Exponential backoff for contended retries, capped so a steady writer cannot
+;; keep a reader in a tight retry loop (livelock). *stm-max-retries* bounds the
+;; total number of non-committing attempts so the loop always terminates.
+(def *stm-max-retries* 4096)
+(def +stm-backoff-base+ 0.0001)
+(def +stm-backoff-cap+ 0.05)
+(def (stm-backoff-delay n)
+  (let (shift (if (> n 12) 12 n))
+    (let (d (* +stm-backoff-base+ (expt 2 shift)))
+      (if (> d +stm-backoff-cap+) +stm-backoff-cap+ d))))
 
 (def (%run-atomically thunk)
   (if (*current-tx*)
     (thunk)                                  ; nested: join the outer transaction
-    (let loop ()
+    (let loop ((attempts 0) (backoff 0))
+      (when (>= attempts *stm-max-retries*)
+        (error 'atomically "transaction retry limit exceeded" *stm-max-retries*))
       (let (tx (make-tx))
         (parameterize ((*current-tx* tx))
           (let (outcome
@@ -164,8 +192,15 @@
                      (if (stm-retry-exn? e) 'wait (raise e)))))
             (cond
               ((pair? outcome) (cdr outcome))
-              ((eq? outcome 'wait) (stm-wait-on-reads! tx) (loop))
-              (else (loop)))))))))
+              ((eq? outcome 'wait)
+               (if (stm-wait-on-reads! tx)
+                 (loop (+ attempts 1) 0)     ; woke on a real write: backoff reset
+                 (begin                        ; missed wakeup / no reads: back off
+                   (thread-sleep! (stm-backoff-delay backoff))
+                   (loop (+ attempts 1) (+ backoff 1)))))
+              (else                            ; commit validation failed: back off
+                (thread-sleep! (stm-backoff-delay backoff))
+                (loop (+ attempts 1) (+ backoff 1))))))))))
 
 (defrules atomically ()
   ((_ body ...) (%run-atomically (lambda () body ...))))
diff --git a/tests/stm-engine-test.ss b/tests/stm-engine-test.ss
index cf1f24f..bde9de8 100644
--- a/tests/stm-engine-test.ss
+++ b/tests/stm-engine-test.ss
@@ -2,7 +2,9 @@
 
 (import :std/sugar
         :jerboa-compat/stm
-        :jerboa-compat/engine)
+        :jerboa-compat/engine
+        :jerboa-compat/lru-cache
+        :jerboa-compat/atom)
 
 (def failures 0)
 
@@ -93,5 +95,92 @@
 (check! "timed-eval rejection leaves the thunk untouched"
         (not (vector-ref arbitrary-thunk-ran? 0)))
 
+;; Basic retry/wait still works: a reader blocks on (retry) until a concurrent
+;; commit changes the tvar it read, then observes the committed value.
+(def br-tv (make-tvar 0))
+(def br-reader (spawn (lambda ()
+                        (atomically
+                          (let (v (tvar-read br-tv))
+                            (when (= v 0) (retry))
+                            v)))))
+(thread-sleep! 0.05)              ; let the reader reach its retry/wait first
+(atomically (tvar-write! br-tv 7))
+(def br-timeout '#(basic-retry-timeout))
+(def br-r (thread-join! br-reader 5 br-timeout))
+(check! "retry blocks, then observes a concurrent commit"
+        (and (not (eq? br-r br-timeout)) (= br-r 7)))
+
+;; P1 #30 regression (deterministic): force the writer to commit AFTER the
+;; reader reads the tvar but BEFORE the reader registers its retry waiter — the
+;; exact lost-wakeup window. The reader reads 0, signals the writer, waits for
+;; the writer's commit, then (retry)s. With the bug the writer's notify-waiters!
+;; sees an empty waiter list and the reader blocks forever (join times out). The
+;; fix re-validates the read-set version after registering and retries at once.
+(def lw-rounds 50)
+(def lw-timeout '#(lost-wakeup-timeout))
+(def lw-failures 0)
+(let loop ((i 0))
+  (when (< i lw-rounds)
+    (let ((tv (make-tvar 0))
+          (reader-ready (vector #f))
+          (writer-done (vector #f)))
+      (let ((reader (spawn (lambda ()
+                             (atomically
+                               (let (v (tvar-read tv))
+                                 (when (= v 0)
+                                   (vector-set! reader-ready 0 #t)
+                                   (let spin ()
+                                     (unless (vector-ref writer-done 0)
+                                       (thread-yield!) (spin)))
+                                   (retry))
+                                 v)))))
+            (writer (spawn (lambda ()
+                             (let spin ()
+                               (unless (vector-ref reader-ready 0)
+                                 (thread-yield!) (spin)))
+                             (atomically (tvar-write! tv 42))
+                             (vector-set! writer-done 0 #t)))))
+        (let (r (thread-join! reader 3 lw-timeout))
+          (unless (and (not (eq? r lw-timeout)) (= r 42))
+            (set! lw-failures (+ lw-failures 1))))
+        (thread-join! writer 3)))
+    (loop (+ i 1))))
+(check! "a commit inside the retry window never loses the reader's wakeup"
+        (= lw-failures 0))
+
+;; equal?-but-not-eq? string keys must hit the same slot (no phantom miss, no
+;; duplicate insert that grows the table past capacity).
+(def lru (make-lru-cache 4))
+(lru-cache-put! lru (string-copy "alpha") 1)
+(check! "lru-cache hits an equal?-but-not-eq? string key"
+        (= (lru-cache-get lru (string-copy "alpha") 'miss) 1))
+(check! "lru-cache contains? an equal?-but-not-eq? string key"
+        (lru-cache-contains? lru (string-copy "alpha")))
+(lru-cache-put! lru (string-copy "alpha") 2)
+(check! "lru-cache does not duplicate an equal? string key"
+        (= (lru-cache-size lru) 1))
+(check! "lru-cache update through an equal? key takes effect"
+        (= (lru-cache-get lru "alpha" 'miss) 2))
+(def lru-cap (make-lru-cache 2))
+(lru-cache-put! lru-cap (string-copy "k") 1)
+(lru-cache-put! lru-cap (string-copy "k") 2)
+(lru-cache-put! lru-cap (string-copy "k") 3)
+(check! "lru-cache stays within capacity for equal? string keys"
+        (= (lru-cache-size lru-cap) 1))
+
+;; A plain 3-vector carrying the legacy tag must not pass atom? (structural
+;; forgery); a real atom must, and the ops must work through the struct.
+(def forged (vector 'jerboa-compat-atom 0 (make-mutex 'forged-atom)))
+(check! "atom? rejects a structurally forged tagged vector" (not (atom? forged)))
+(check! "atom-deref rejects a forged atom"
+        (try (begin (atom-deref forged) #f) (catch (e) #t)))
+(def a (atom 10))
+(check! "atom? accepts a real atom" (atom? a))
+(check! "atom-deref reads the initial value" (= (atom-deref a) 10))
+(check! "atom-swap! updates the value" (= (atom-swap! a (lambda (x) (+ x 5))) 15))
+(check! "atom-update! applies with args" (= (atom-update! a + 100) 115))
+(check! "atom-reset! sets the value" (= (atom-reset! a 1) 1))
+(check! "atom-deref reads the reset value" (= (atom-deref a) 1))
+
 (displayln "stm-engine-tests: ok")
 (exit (if (> failures 0) 1 0))