Fix TxnVarMap new-key concurrency bugs (#51)

Three related fixes for when a transaction adds a new key to a
TxnVarMap and the underlying TxnVar doesn't exist yet:

- Commit-time lock fallback: TxnLogUpdateVarMapEntry.lock now falls
  back to the map's structural commitLock when getTxnVar returns None,
  ensuring new-key commits serialise with structural changes.

- addOrUpdate TOCTOU fix: moved the value.get + match inside the
  internalStructureLock to eliminate the race between checking for
  key existence and performing the add/update.

- delete TOCTOU fix: same pattern — moved the check-then-act inside
  the lock.

Author: gvonness-apolitical

src/main/scala/bengal/stm/model/TxnVarMap.scala

@@ -82,38 +82,34 @@ case class TxnVarMap[F[_]: STM: Async, K, V](
   ): F[TxnVarRuntimeId] =
     Async[F].delay(getRuntimeExistentialId(key).addParent(runtimeId))
 
-  // Only called when key is known to not exist
-  private def add(newKey: K, newValue: V): F[Unit] =
-    for {
-      newTxnVar <- TxnVar.of(newValue)
-      _ <- withLock(internalStructureLock)(
-             value.update(_ += (newKey -> newTxnVar))
-           )
-    } yield ()
-
   private[stm] def addOrUpdate(key: K, newValue: V): F[Unit] =
-    for {
-      txnVarMap <- value.get
-      _ <- txnVarMap.get(key) match {
-             case Some(tVar) =>
-               withLock(internalStructureLock)(
+    withLock(internalStructureLock) {
+      for {
+        txnVarMap <- value.get
+        _ <- txnVarMap.get(key) match {
+               case Some(tVar) =>
                  tVar.set(newValue)
-               )
-             case None =>
-               add(key, newValue)
-           }
-    } yield ()
+               case None =>
+                 for {
+                   newTxnVar <- TxnVar.of(newValue)
+                   _         <- value.update(_ += (key -> newTxnVar))
+                 } yield ()
+             }
+      } yield ()
+    }
 
   private[stm] def delete(key: K): F[Unit] =
-    for {
-      txnVarMap <- value.get
-      _ <- txnVarMap.get(key) match {
-             case Some(_) =>
-               withLock(internalStructureLock)(value.update(_ -= key))
-             case None =>
-               Async[F].unit
-           }
-    } yield ()
+    withLock(internalStructureLock) {
+      for {
+        txnVarMap <- value.get
+        _ <- txnVarMap.get(key) match {
+               case Some(_) =>
+                 value.update(_ -= key)
+               case None =>
+                 Async[F].unit
+             }
+      } yield ()
+    }
 }
 
 object TxnVarMap {

src/main/scala/bengal/stm/runtime/TxnLogContext.scala

@@ -276,7 +276,7 @@ private[stm] trait TxnLogContext[F[_]] {
     override private[stm] lazy val lock: F[Option[Semaphore[F]]] =
       for {
         oTxnVar <- txnVarMap.getTxnVar(key)
-      } yield oTxnVar.map(_.commitLock)
+      } yield Some(oTxnVar.map(_.commitLock).getOrElse(txnVarMap.commitLock))
 
     override private[stm] lazy val idFootprint: F[IdFootprint] =
       txnVarMap.getRuntimeId(key).map { rid =>

src/test/scala/bengal/stm/runtime/TxnLogEntrySpec.scala

@@ -485,5 +485,13 @@ class TxnLogEntrySpec extends AsyncFreeSpec with AsyncIOSpec with Matchers {
         footprint.updatedIds shouldBe Set(rid)
       }
     }
+
+    "lock returns Some for new key (falls back to map commitLock)" in withRuntime { implicit stm =>
+      for {
+        tvarMap <- TxnVarMap.of(Map("a" -> 1))
+        entry = stm.TxnLogUpdateVarMapEntry[String, Int]("newkey", None, Some(5), tvarMap)
+        lock <- entry.lock
+      } yield lock shouldBe Some(tvarMap.commitLock)
+    }
   }
 }

src/test/scala/model/TxnVarMapSpec.scala

@@ -18,6 +18,7 @@ package ai.entrolution
 package model
 
 import cats.effect.IO
+import cats.effect.implicits._
 import cats.effect.testing.scalatest.AsyncIOSpec
 import org.scalatest.EitherValues
 import org.scalatest.freespec.AsyncFreeSpec
@@ -235,4 +236,38 @@ class TxnVarMapSpec extends AsyncFreeSpec with AsyncIOSpec with Matchers with Ei
         .asserting(_ shouldBe None)
     }
   }
+
+  "concurrent new-key operations" - {
+    "concurrent set of same new key" in {
+      STM
+        .runtime[IO]
+        .flatMap { implicit stm =>
+          for {
+            tVarMap <- TxnVarMap.of(Map.empty[String, Int])
+            _ <- (
+                   tVarMap.set("x", 1).commit,
+                   tVarMap.set("x", 2).commit
+                 ).parTupled
+            result <- tVarMap.get("x").commit
+          } yield result
+        }
+        .asserting(_ shouldBe defined)
+    }
+
+    "concurrent delete of different keys" in {
+      STM
+        .runtime[IO]
+        .flatMap { implicit stm =>
+          for {
+            tVarMap <- TxnVarMap.of(Map("a" -> 1, "b" -> 2))
+            _ <- (
+                   tVarMap.remove("a").commit,
+                   tVarMap.remove("b").commit
+                 ).parTupled
+            result <- tVarMap.get.commit
+          } yield result
+        }
+        .asserting(_ shouldBe empty)
+    }
+  }
 }