fix: use dual schema writer in MemDB to write both legacy and unified schemas

MemDB supports both legacy and unified schema storage, but was only
using LegacySchemaWriterAdapter which after our changes no longer writes
the legacy schema hash. This caused TestUnifiedSchemaHashWatch to fail
with a timeout waiting for hash updates.

Changes:
1. Updated SchemaWriter() to return a dualSchemaWriter that writes to
   both legacy and unified storage (SchemaModeReadNewWriteBoth)

2. Refactored schema writing methods to avoid deadlocks:
   - Added writeStoredSchemaNoLock() that uses the transaction's
     db handle without acquiring additional locks
   - Added writeSchemaHashNoLock() for transaction-scoped hash writes
   - Added writeLegacySchemaHashFromDefinitionsNoLock() for legacy hash

3. Updated WriteStoredSchema() and WriteLegacySchemaHashFromDefinitions()
   to use the transaction's db handle via txSource()

This ensures MemDB writes to both storages and properly updates the
schema hash for watchers, fixing the unified schema hash watch test.

Author: josephschorr

internal/datastore/memdb/memdb.go

@@ -461,45 +461,27 @@ func (mdb *memdbDatastore) readStoredSchemaInternal() (*corev1.StoredSchema, err
 	return mdb.storedSchema.CloneVT(), nil
 }
 
-// writeStoredSchemaInternal is an internal method for transactions to write the stored schema.
-// This should NOT be called directly - use transactions instead.
-func (mdb *memdbDatastore) writeStoredSchemaInternal(schema *corev1.StoredSchema) error {
-	if schema == nil {
-		return errors.New("stored schema cannot be nil")
-	}
-
-	if schema.Version == 0 {
-		return errors.New("stored schema version cannot be 0")
-	}
-
-	mdb.Lock()
-	defer mdb.Unlock()
-
-	if err := mdb.checkNotClosed(); err != nil {
-		return err
-	}
-
+// writeStoredSchemaNoLock writes the stored schema using the provided transaction.
+// This is called from within an existing transaction, so it doesn't acquire locks or commit.
+func (mdb *memdbDatastore) writeStoredSchemaNoLock(tx *memdb.Txn, schema *corev1.StoredSchema) error {
 	// Store a copy to prevent external mutations
 	mdb.storedSchema = schema.CloneVT()
 
 	// Write the schema hash to the schema revision table for fast lookups
-	if err := mdb.writeSchemaHashInternal(schema); err != nil {
+	if err := mdb.writeSchemaHashNoLock(tx, schema); err != nil {
 		return fmt.Errorf("failed to write schema hash: %w", err)
 	}
 
 	return nil
 }
 
-// writeSchemaHashInternal writes the schema hash to the in-memory schema revision table
-func (mdb *memdbDatastore) writeSchemaHashInternal(schema *corev1.StoredSchema) error {
+// writeSchemaHashNoLock writes the schema hash to the in-memory schema revision table using the provided transaction.
+func (mdb *memdbDatastore) writeSchemaHashNoLock(tx *memdb.Txn, schema *corev1.StoredSchema) error {
 	v1 := schema.GetV1()
 	if v1 == nil {
 		return fmt.Errorf("unsupported schema version: %d", schema.Version)
 	}
 
-	tx := mdb.db.Txn(true)
-	defer tx.Abort()
-
 	// Delete existing hash (if any)
 	if existing, err := tx.First(tableSchemaRevision, indexID, "current"); err == nil && existing != nil {
 		if err := tx.Delete(tableSchemaRevision, existing); err != nil {
@@ -517,12 +499,14 @@ func (mdb *memdbDatastore) writeSchemaHashInternal(schema *corev1.StoredSchema)
 		return fmt.Errorf("failed to insert hash: %w", err)
 	}
 
-	tx.Commit()
+	// Note: Don't commit here - the caller will commit the transaction
 	return nil
 }
 
 // writeLegacySchemaHashFromDefinitionsInternal writes the schema hash computed from the given definitions
-func (mdb *memdbDatastore) writeLegacySchemaHashFromDefinitionsInternal(ctx context.Context, namespaces []datastore.RevisionedNamespace, caveats []datastore.RevisionedCaveat) error {
+// writeLegacySchemaHashFromDefinitionsNoLock writes the schema hash using the provided transaction.
+// This is called from within an existing transaction, so it doesn't acquire locks.
+func (mdb *memdbDatastore) writeLegacySchemaHashFromDefinitionsNoLock(ctx context.Context, tx *memdb.Txn, namespaces []datastore.RevisionedNamespace, caveats []datastore.RevisionedCaveat) error {
 	// Build schema definitions list
 	definitions := make([]compiler.SchemaDefinition, 0, len(namespaces)+len(caveats))
 	for _, ns := range namespaces {
@@ -547,16 +531,6 @@ func (mdb *memdbDatastore) writeLegacySchemaHashFromDefinitionsInternal(ctx cont
 	hash := sha256.Sum256([]byte(schemaText))
 	schemaHash := hex.EncodeToString(hash[:])
 
-	mdb.Lock()
-	defer mdb.Unlock()
-
-	if err := mdb.checkNotClosed(); err != nil {
-		return err
-	}
-
-	tx := mdb.db.Txn(true)
-	defer tx.Abort()
-
 	// Delete existing hash (if any)
 	if existing, err := tx.First(tableSchemaRevision, indexID, "current"); err == nil && existing != nil {
 		if err := tx.Delete(tableSchemaRevision, existing); err != nil {
@@ -574,7 +548,7 @@ func (mdb *memdbDatastore) writeLegacySchemaHashFromDefinitionsInternal(ctx cont
 		return fmt.Errorf("failed to insert hash: %w", err)
 	}
 
-	tx.Commit()
+	// Note: Don't commit here - the caller will commit the transaction
 	return nil
 }
 

internal/datastore/memdb/readwrite.go

@@ -348,17 +348,30 @@ func (rwt *memdbReadWriteTx) LegacyDeleteNamespaces(ctx context.Context, nsNames
 }
 
 func (rwt *memdbReadWriteTx) SchemaWriter() (datastore.SchemaWriter, error) {
-	return schemautil.NewLegacySchemaWriterAdapter(rwt, rwt), nil
+	// MemDB supports both legacy and unified schema storage.
+	// Use write-to-both mode to ensure both are updated.
+	return schemautil.NewSchemaWriter(rwt, rwt, options.SchemaModeReadNewWriteBoth), nil
 }
 
 // WriteStoredSchema implements datastore.SingleStoreSchemaWriter
 func (rwt *memdbReadWriteTx) WriteStoredSchema(ctx context.Context, schema *core.StoredSchema) error {
-	return rwt.datastore.writeStoredSchemaInternal(schema)
+	// Called from within a transaction - use the transaction's db handle directly
+	tx, err := rwt.txSource()
+	if err != nil {
+		return err
+	}
+	return rwt.datastore.writeStoredSchemaNoLock(tx, schema)
 }
 
 // WriteLegacySchemaHashFromDefinitions implements datastore.LegacySchemaHashWriter
 func (rwt *memdbReadWriteTx) WriteLegacySchemaHashFromDefinitions(ctx context.Context, namespaces []datastore.RevisionedNamespace, caveats []datastore.RevisionedCaveat) error {
-	return rwt.datastore.writeLegacySchemaHashFromDefinitionsInternal(ctx, namespaces, caveats)
+	// Called from within a transaction - use the transaction's db handle directly
+	// without trying to acquire additional locks
+	tx, err := rwt.txSource()
+	if err != nil {
+		return err
+	}
+	return rwt.datastore.writeLegacySchemaHashFromDefinitionsNoLock(ctx, tx, namespaces, caveats)
 }
 
 func (rwt *memdbReadWriteTx) BulkLoad(ctx context.Context, iter datastore.BulkWriteRelationshipSource) (uint64, error) {