fix: Handle concurrent MemDB writes

internal/datastore/memdb/memdb.go

@@ -374,7 +374,7 @@ func (mdb *memdbDatastore) ReadWriteTx(
 		// Create a snapshot and add it to the revisions slice
 		schemaHash := mdb.getCurrentSchemaHashNoLock()
 		snap := mdb.db.Snapshot()
-		mdb.revisions = append(mdb.revisions, snapshot{newRevision, schemaHash, snap})
+		mdb.insertRevisionSnapshot(newRevision, schemaHash, snap)
 		return newRevision, nil
 	}
 

internal/datastore/memdb/memdb_test.go

@@ -113,6 +113,75 @@ func TestConcurrentWriteRelsSucceed(t *testing.T) {
 	require.NoError(g.Wait())
 }
 
+// TestConcurrentWriteRevisionInversionLosesRead reproduces concurrent writes to MemDB
+func TestConcurrentWriteRevisionInversionLosesRead(t *testing.T) {
+	require := require.New(t)
+
+	ds, err := NewMemdbDatastore(0, 1*time.Hour, 1*time.Hour)
+	require.NoError(err)
+	mds := ds.(*memdbDatastore)
+
+	ctx := t.Context()
+
+	aEntered := make(chan struct{})
+	releaseA := make(chan struct{})
+	aDone := make(chan struct{})
+
+	var revA datastore.Revision
+	var errA error
+	go func() {
+		defer close(aDone)
+		revA, errA = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+			close(aEntered)
+			<-releaseA
+			return rwt.WriteRelationships(ctx, []tuple.RelationshipUpdate{
+				tuple.Touch(tuple.MustParse("document:doc-a#viewer@user:tom")),
+			})
+		}, options.WithDisableRetries(true))
+	}()
+
+	<-aEntered
+
+	revB, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+		return rwt.WriteRelationships(ctx, []tuple.RelationshipUpdate{
+			tuple.Touch(tuple.MustParse("document:doc-b#viewer@user:tom")),
+		})
+	}, options.WithDisableRetries(true))
+	require.NoError(err)
+
+	close(releaseA)
+	<-aDone
+	require.NoError(errA)
+
+	require.True(revB.GreaterThan(revA), "expected B's revision (%v) to be greater than A's (%v)", revB, revA)
+
+	indexOf := func(r datastore.Revision) int {
+		mds.RLock()
+		defer mds.RUnlock()
+		for i, snap := range mds.revisions {
+			if snap.revision.Equal(r) {
+				return i
+			}
+		}
+		return -1
+	}
+	t.Logf("storage order: A(rev=%v) at index %d, B(rev=%v) at index %d", revA, indexOf(revA), revB, indexOf(revB))
+
+	reader := ds.SnapshotReader(revA)
+	it, err := reader.QueryRelationships(ctx, datastore.RelationshipsFilter{OptionalResourceType: "document"})
+	require.NoError(err)
+	rels, err := datastore.IteratorToSlice(it)
+	require.NoError(err)
+
+	var sawA bool
+	for _, rel := range rels {
+		if rel.Resource.ObjectID == "doc-a" {
+			sawA = true
+		}
+	}
+	require.True(sawA, "relationship written by transaction A is invisible when reading at A's own returned revision %v", revA)
+}
+
 func TestAnythingAfterCloseDoesNotPanic(t *testing.T) {
 	require := require.New(t)
 

internal/datastore/memdb/revisions.go

@@ -2,8 +2,12 @@ package memdb
 
 import (
 	"context"
+	"slices"
+	"sort"
 	"time"
 
+	"github.com/hashicorp/go-memdb"
+
 	"github.com/authzed/spicedb/internal/datastore/revisions"
 	"github.com/authzed/spicedb/pkg/datastore"
 )
@@ -38,6 +42,13 @@ func (mdb *memdbDatastore) newRevisionID() revisions.TimestampRevision {
 	return created
 }
 
+func (mdb *memdbDatastore) insertRevisionSnapshot(newRevision revisions.TimestampRevision, schemaHash string, snap *memdb.MemDB) {
+	insertAt := sort.Search(len(mdb.revisions), func(i int) bool {
+		return mdb.revisions[i].revision.GreaterThan(newRevision)
+	})
+	mdb.revisions = slices.Insert(mdb.revisions, insertAt, snapshot{newRevision, schemaHash, snap})
+}
+
 func (mdb *memdbDatastore) HeadRevision(_ context.Context) (datastore.RevisionWithSchemaHash, error) {
 	mdb.RLock()
 	defer mdb.RUnlock()