fix: Handle concurrent MemDB writes

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Caveats: compiled caveats (and their CEL environments) are now cached per schema version — hung off the stored schema (`ReadOnlyStoredSchema`) and rebuilt only when the schema changes — rather than rebuilt on every check, reducing check cost for schemas with many caveats (https://github.com/authzed/spicedb/pull/3166)
 
 ### Fixed
+- MemDB: overlapping write transactions could violate every snapshot-consistency invariant of the datastore — a committed write could be invisible at its own returned revision (breaking read-your-writes, e.g. an at-exact-snapshot `Check` right after `WriteRelationships`), a later commit could leak into reads at an earlier revision, a write visible at one revision could be missing at a later one (including head), and two concurrent transactions could be assigned the same revision. Revisions are now assigned at write-transaction acquisition, where writer serialization makes the two orders identical. (https://github.com/authzed/spicedb/pull/3239)
 - Fixed a nil pointer dereference panic in `CheckBulkPermissions` that could occur under concurrent load when a tracing-enabled check shared a singleflight dispatch with a non-tracing bulk check. Debug-enabled checks are no longer singleflighted together with non-debug checks. (https://github.com/authzed/spicedb/pull/3174)
 - Fixed a nil pointer dereference panic in the Postgres FDW (https://github.com/authzed/spicedb/pull/3235)
 - CockroachDB: deletes performed by CockroachDB's row-level TTL job for expired relationships are no longer emitted as `DELETE` events by the Watch API. On CockroachDB ≥ 24.1, SpiceDB sets the `ttl_disable_changefeed_replication` storage parameter on the relationship tables at startup (if it lacks `ALTER TABLE` privileges, it logs a warning with the statement to run manually); on older versions a startup warning is logged and TTL deletes continue to be emitted. Note that the parameter affects any changefeed over these tables — external changefeeds that want TTL deletes can opt back in with `ignore_disable_changefeed_replication`. Delete-only transactions also no longer write an internal transaction-metadata marker row, reducing write amplification. (https://github.com/authzed/spicedb/pull/3210)

internal/datastore/memdb/memdb.go

@@ -90,10 +90,11 @@
 	revisions.CommonDecoder
 
 	// NOTE: call checkNotClosed before using
-	db             *memdb.MemDB // GUARDED_BY(RWMutex)
-	revisions      []snapshot   // GUARDED_BY(RWMutex)
-	activeWriteTxn *memdb.Txn   // GUARDED_BY(RWMutex)
-	writeTxReady   *sync.Cond   // broadcast when activeWriteTxn becomes nil
+	db *memdb.MemDB // GUARDED_BY(RWMutex)
+	// revisions MUST be sorted in strictly increasing revision order.
+	revisions      []snapshot // GUARDED_BY(RWMutex)
+	activeWriteTxn *memdb.Txn // GUARDED_BY(RWMutex)
+	writeTxReady   *sync.Cond // broadcast when activeWriteTxn becomes nil
 
 	negativeGCWindow        int64
 	quantizationPeriod      int64
@@ -116,6 +117,9 @@
 	return mdb.uniqueID, nil
 }
 
+// SnapshotReader returns a reader for the snapshot visible at the given
+// revision: the first entry in mdb.revisions at or after it, located by
+// binary search.
 func (mdb *memdbDatastore) SnapshotReader(dr datastore.Revision) datastore.Reader {
 	mdb.RLock()
 	defer mdb.RUnlock()
@@ -188,6 +192,8 @@
 
 	for i := 0; i < txNumAttempts; i++ {
 		var tx *memdb.Txn
+		var newRevision revisions.TimestampRevision
+		rwt := &memdbReadWriteTx{}
 		createTxOnce := sync.Once{}
 		txSrc := func() (*memdb.Txn, error) {
 			var err error
@@ -208,13 +214,16 @@
 				tx = mdb.db.Txn(true)
 				tx.TrackChanges()
 				mdb.activeWriteTxn = tx
+
+				// Assign the transaction's revision now, while holding the exclusive write transaction.
+				// Writers are serialized from this point until the revision is appended at commit.
+				newRevision = mdb.newRevisionIDNoLock()
+				rwt.newRevision = newRevision
 			})
 
 			return tx, err
 		}
-
-		newRevision := mdb.newRevisionID()
-		rwt := &memdbReadWriteTx{memdbReader{&sync.Mutex{}, txSrc, nil, time.Now()}, newRevision}
+		rwt.memdbReader = memdbReader{&sync.Mutex{}, txSrc, nil, time.Now()}
 		if config.SchemaHashPrecondition != "" {
 			if err := assertSchemaHash(ctx, rwt, config.SchemaHashPrecondition); err != nil {
 				mdb.Lock()
@@ -247,123 +256,130 @@
 		}
 
 		mdb.Lock()
-		defer mdb.Unlock()
+		defer mdb.Unlock() // TODO is this defer correct? it runs at the end of the function, not at the end of the for loop's iteration
+
+		// The user function never used the transaction: nothing was written,
+		// so no new revision is created and the head revision is returned.
+		if tx == nil {
+			if err := mdb.checkNotClosed(); err != nil {
+				return datastore.NoRevision, err
+			}
+			return mdb.headRevisionNoLock(), nil
+		}
 
 		tracked := common.NewChanges(revisions.TimestampIDKeyFunc, datastore.WatchRelationships|datastore.WatchSchema, 0)
-		if tx != nil {
-			if config.Metadata != nil && len(config.Metadata.GetFields()) > 0 {
-				if err := tracked.AddRevisionMetadata(ctx, newRevision, config.Metadata.AsMap()); err != nil {
-					return datastore.NoRevision, err
-				}
+		if config.Metadata != nil && len(config.Metadata.GetFields()) > 0 {
+			if err := tracked.AddRevisionMetadata(ctx, newRevision, config.Metadata.AsMap()); err != nil {
+				return datastore.NoRevision, err
 			}
+		}
 
-			for _, change := range tx.Changes() {
-				switch change.Table {
-				case tableRelationship:
-					switch {
-					case change.After != nil:
-						rt, err := change.After.(*relationship).Relationship()
-						if err != nil {
-							return datastore.NoRevision, err
-						}
-
-						if err := tracked.AddRelationshipChange(ctx, newRevision, rt, tuple.UpdateOperationTouch); err != nil {
-							return datastore.NoRevision, err
-						}
-					case change.After == nil && change.Before != nil:
-						rt, err := change.Before.(*relationship).Relationship()
-						if err != nil {
-							return datastore.NoRevision, err
-						}
-
-						if err := tracked.AddRelationshipChange(ctx, newRevision, rt, tuple.UpdateOperationDelete); err != nil {
-							return datastore.NoRevision, err
-						}
-					default:
-						return datastore.NoRevision, spiceerrors.MustBugf("unexpected relationship change")
+		for _, change := range tx.Changes() {
+			switch change.Table {
+			case tableRelationship:
+				switch {
+				case change.After != nil:
+					rt, err := change.After.(*relationship).Relationship()
+					if err != nil {
+						return datastore.NoRevision, err
 					}
-				case tableNamespace:
-					switch {
-					case change.After != nil:
-						loaded := &corev1.NamespaceDefinition{}
-						if err := loaded.UnmarshalVT(change.After.(*namespace).configBytes); err != nil {
-							return datastore.NoRevision, err
-						}
-
-						err := tracked.AddChangedDefinition(ctx, newRevision, loaded)
-						if err != nil {
-							return datastore.NoRevision, err
-						}
-					case change.After == nil && change.Before != nil:
-						err := tracked.AddDeletedNamespace(ctx, newRevision, change.Before.(*namespace).name)
-						if err != nil {
-							return datastore.NoRevision, err
-						}
-					default:
-						return datastore.NoRevision, spiceerrors.MustBugf("unexpected namespace change")
+
+					if err := tracked.AddRelationshipChange(ctx, newRevision, rt, tuple.UpdateOperationTouch); err != nil {
+						return datastore.NoRevision, err
 					}
-				case tableCaveats:
-					switch {
-					case change.After != nil:
-						loaded := &corev1.CaveatDefinition{}
-						if err := loaded.UnmarshalVT(change.After.(*caveat).definition); err != nil {
-							return datastore.NoRevision, err
-						}
-
-						err := tracked.AddChangedDefinition(ctx, newRevision, loaded)
-						if err != nil {
-							return datastore.NoRevision, err
-						}
-					case change.After == nil && change.Before != nil:
-						err := tracked.AddDeletedCaveat(ctx, newRevision, change.Before.(*caveat).name)
-						if err != nil {
-							return datastore.NoRevision, err
-						}
-					default:
-						return datastore.NoRevision, spiceerrors.MustBugf("unexpected namespace change")
+				case change.After == nil && change.Before != nil:
+					rt, err := change.Before.(*relationship).Relationship()
+					if err != nil {
+						return datastore.NoRevision, err
 					}
-				}
-			}
 
-			changes := tracked.AsRevisionChanges(revisions.TimestampIDKeyLessThanFunc)
-			wroteChangelog := false
-			for rc, err := range changes {
-				if err != nil {
-					return datastore.NoRevision, err
+					if err := tracked.AddRelationshipChange(ctx, newRevision, rt, tuple.UpdateOperationDelete); err != nil {
+						return datastore.NoRevision, err
+					}
+				default:
+					return datastore.NoRevision, spiceerrors.MustBugf("unexpected relationship change")
 				}
+			case tableNamespace:
+				switch {
+				case change.After != nil:
+					loaded := &corev1.NamespaceDefinition{}
+					if err := loaded.UnmarshalVT(change.After.(*namespace).configBytes); err != nil {
+						return datastore.NoRevision, err
+					}
 
-				if wroteChangelog {
-					return datastore.NoRevision, spiceerrors.MustBugf("unexpected MemDB transaction with multiple revision changes")
+					err := tracked.AddChangedDefinition(ctx, newRevision, loaded)
+					if err != nil {
+						return datastore.NoRevision, err
+					}
+				case change.After == nil && change.Before != nil:
+					err := tracked.AddDeletedNamespace(ctx, newRevision, change.Before.(*namespace).name)
+					if err != nil {
+						return datastore.NoRevision, err
+					}
+				default:
+					return datastore.NoRevision, spiceerrors.MustBugf("unexpected namespace change")
 				}
+			case tableCaveats:
+				switch {
+				case change.After != nil:
+					loaded := &corev1.CaveatDefinition{}
+					if err := loaded.UnmarshalVT(change.After.(*caveat).definition); err != nil {
+						return datastore.NoRevision, err
+					}
 
-				change := &changelog{
-					revisionNanos: newRevision.TimestampNanoSec(),
-					changes:       rc,
-				}
-				if err := tx.Insert(tableChangelog, change); err != nil {
-					return datastore.NoRevision, fmt.Errorf("error writing changelog: %w", err)
+					err := tracked.AddChangedDefinition(ctx, newRevision, loaded)
+					if err != nil {
+						return datastore.NoRevision, err
+					}
+				case change.After == nil && change.Before != nil:
+					err := tracked.AddDeletedCaveat(ctx, newRevision, change.Before.(*caveat).name)
+					if err != nil {
+						return datastore.NoRevision, err
+					}
+				default:
+					return datastore.NoRevision, spiceerrors.MustBugf("unexpected namespace change")
 				}
+			}
+		}
 
-				wroteChangelog = true
+		changes := tracked.AsRevisionChanges(revisions.TimestampIDKeyLessThanFunc)
+		wroteChangelog := false
+		for rc, err := range changes {
+			if err != nil {
+				return datastore.NoRevision, err
 			}
 
-			// Always emit a changelog entry for the committed revision, even
-			// when the transaction produced no observable changes (e.g., a
-			// TOUCH that matched the existing relationship). The changes
-			// payload is intentionally empty — the watch goroutine constructs the
-			// checkpoint event itself based on each consumer's options.
-			if !wroteChangelog {
-				change := &changelog{
-					revisionNanos: newRevision.TimestampNanoSec(),
-					changes:       datastore.RevisionChanges{},
-				}
-				if err := tx.Insert(tableChangelog, change); err != nil {
-					return datastore.NoRevision, fmt.Errorf("error writing changelog: %w", err)
-				}
+			if wroteChangelog {
+				return datastore.NoRevision, spiceerrors.MustBugf("unexpected MemDB transaction with multiple revision changes")
 			}
 
-			tx.Commit()
+			change := &changelog{
+				revisionNanos: newRevision.TimestampNanoSec(),
+				changes:       rc,
+			}
+			if err := tx.Insert(tableChangelog, change); err != nil {
+				return datastore.NoRevision, fmt.Errorf("error writing changelog: %w", err)
+			}
+
+			wroteChangelog = true
+		}
+
+		// Always emit a changelog entry for the committed revision, even
+		// when the transaction produced no observable changes (e.g., a
+		// TOUCH that matched the existing relationship). The changes
+		// payload is intentionally empty — the watch goroutine constructs the
+		// checkpoint event itself based on each consumer's options.
+		if !wroteChangelog {
+			change := &changelog{
+				revisionNanos: newRevision.TimestampNanoSec(),
+				changes:       datastore.RevisionChanges{},
+			}
+			if err := tx.Insert(tableChangelog, change); err != nil {
+				return datastore.NoRevision, fmt.Errorf("error writing changelog: %w", err)
+			}
 		}
+
+		tx.Commit()
 		mdb.activeWriteTxn = nil
 		mdb.writeTxReady.Signal()
 

internal/datastore/memdb/revisions.go

@@ -14,25 +14,24 @@
 	return revisions.NewForTime(time.Now().UTC())
 }
 
-func (mdb *memdbDatastore) newRevisionID() revisions.TimestampRevision {
-	mdb.Lock()
-	defer mdb.Unlock()
-
+// newRevisionIDNoLock returns a revision strictly greater than the current head revision.
+// The caller must hold the RWMutex.
+func (mdb *memdbDatastore) newRevisionIDNoLock() revisions.TimestampRevision {
 	existing := mdb.revisions[len(mdb.revisions)-1].revision
 	created := nowRevision()
 
 	// NOTE: The time.Now().UTC() only appears to have *microsecond* level
 	// precision on macOS Monterey in Go 1.19.1. This means that HeadRevision
 	// and the result of a ReadWriteTx could return the *same* transaction ID
 	// if both are executed in sequence without any other forms of delay on
-	// macOS. We therefore check if the created transaction ID matches that
-	// previously created and, if not, add to it.
+	// macOS. We therefore check if the created transaction ID is at or before
+	// the head revision and, if so, advance past the head.
 	//
 	// See: https://github.com/golang/go/issues/22037 which appeared to fix
 	// this in Go 1.9.2, but there appears to have been a reversion with either
 	// the new version of macOS or Go.
-	if created.Equal(existing) {
-		return revisions.NewForTimestamp(created.TimestampNanoSec() + 1)
+	if !created.GreaterThan(existing) {
+		return revisions.NewForTimestamp(existing.TimestampNanoSec() + 1)
 	}
 
 	return created
@@ -120,7 +119,7 @@
 	// HEAD revision is behind it.
 	if dr.GreaterThan(now) {
 		// If the revision is in the "future", then check to ensure that it is <= of HEAD to handle
-		// the microsecond granularity on macos (see comment above in newRevisionID)
+		// the microsecond granularity on macos (see comment above in newRevisionIDNoLock)
 		headRevision := mdb.headRevisionNoLock()
 		if dr.LessThan(headRevision) || dr.Equal(headRevision) {
 			return nil

pkg/datastore/test/datastore.go

@@ -172,6 +172,10 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories)
 		t.Run("TestConcurrentWriteSerialization", runner(tester, ConcurrentWriteSerializationTest))
 		t.Run("TestConcurrentWriteDeadlock", runner(tester, ConcurrentWriteDeadlockTest))
 	}
+	// Not gated on the ConcurrentWrite category: the overlapping transactions
+	// never contend for the write lock, so even datastores with a global write
+	// lock (e.g. memdb) must pass it. See the test's doc comment.
+	t.Run("TestConcurrentWriteRevisionVisibility", runner(tester, ConcurrentWriteRevisionVisibilityTest))
 
 	t.Run("TestOrdering", runner(tester, OrderingTest))
 	t.Run("TestLimit", runner(tester, LimitTest))