Fix atomic latest version update to prevent missing isLatest flags

- Pass transaction via context from WithPublishLock to CreateServer
- Use pg_try_advisory_xact_lock with retry to avoid connection pool exhaustion
- Update test to check database state (source of truth) vs returned snapshots
- Increase publish timeout to 60s to handle high concurrency scenarios
- Test now passes with concurrency=100

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

Author: domdomegg

internal/database/postgres.go

@@ -15,6 +15,10 @@ import (
 	apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0"
 )
 
+type contextKey string
+
+const publishTxKey contextKey = "publishTx"
+
 // PostgreSQL is an implementation of the Database interface using PostgreSQL
 type PostgreSQL struct {
 	pool *pgxpool.Pool
@@ -328,8 +332,9 @@ func (db *PostgreSQL) GetAllVersionsByServerID(ctx context.Context, serverID str
 	return results, nil
 }
 
-// CreateServer atomically publishes a new server version, optionally unmarking a previous latest version
-// Must be called within WithPublishLock to ensure proper serialization
+// CreateServer publishes a new server version, optionally unmarking a previous latest version
+// MUST be called within WithPublishLock to ensure proper serialization and prevent race conditions
+// Executes UPDATE and INSERT atomically using the transaction from the context
 func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON, oldLatestVersionID *string) (*apiv0.ServerJSON, error) {
 	if ctx.Err() != nil {
 		return nil, ctx.Err()
@@ -347,14 +352,11 @@ func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON
 		return nil, fmt.Errorf("server must have both ServerID and VersionID in registry metadata")
 	}
 
-	// Begin a transaction for atomicity of UPDATE + INSERT
-	tx, err := db.pool.Begin(ctx)
-	if err != nil {
-		return nil, fmt.Errorf("failed to begin transaction: %w", err)
+	// Get transaction from context (must be called within WithPublishLock)
+	tx, ok := ctx.Value(publishTxKey).(pgx.Tx)
+	if !ok {
+		return nil, fmt.Errorf("CreateServer must be called within WithPublishLock")
 	}
-	defer func() {
-		_ = tx.Rollback(ctx)
-	}()
 
 	// If there's a previous latest version, unmark it
 	if oldLatestVersionID != nil && *oldLatestVersionID != "" {
@@ -389,11 +391,6 @@ func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON
 		return nil, fmt.Errorf("failed to insert server: %w", err)
 	}
 
-	// Commit the transaction
-	if err := tx.Commit(ctx); err != nil {
-		return nil, fmt.Errorf("failed to commit transaction: %w", err)
-	}
-
 	return server, nil
 }
 
@@ -440,34 +437,57 @@ func (db *PostgreSQL) WithPublishLock(ctx context.Context, serverName string, fn
 		return ctx.Err()
 	}
 
-	// Begin a transaction
-	tx, err := db.pool.Begin(ctx)
-	if err != nil {
-		return fmt.Errorf("failed to begin transaction: %w", err)
-	}
-	defer func() {
-		_ = tx.Rollback(ctx)
-	}()
-
-	// Acquire advisory lock based on server name hash
-	// Using pg_advisory_xact_lock which auto-releases on transaction end
 	lockID := hashServerName(serverName)
-	_, err = tx.Exec(ctx, "SELECT pg_advisory_xact_lock($1)", lockID)
-	if err != nil {
-		return fmt.Errorf("failed to acquire publish lock: %w", err)
-	}
+	backoff := 10 * time.Millisecond
 
-	// Execute the function
-	if err := fn(ctx); err != nil {
-		return err
-	}
+	for {
+		// Begin a transaction
+		tx, err := db.pool.Begin(ctx)
+		if err != nil {
+			return fmt.Errorf("failed to begin transaction: %w", err)
+		}
 
-	// Commit the transaction (which also releases the lock)
-	if err := tx.Commit(ctx); err != nil {
-		return fmt.Errorf("failed to commit transaction: %w", err)
-	}
+		// Try to acquire advisory lock (non-blocking)
+		// Using pg_try_advisory_xact_lock which returns immediately
+		var acquired bool
+		err = tx.QueryRow(ctx, "SELECT pg_try_advisory_xact_lock($1)", lockID).Scan(&acquired)
+		if err != nil {
+			_ = tx.Rollback(ctx)
+			return fmt.Errorf("failed to acquire publish lock: %w", err)
+		}
 
-	return nil
+		if !acquired {
+			// Lock not available, rollback and retry
+			_ = tx.Rollback(ctx)
+
+			// Check if context is still valid
+			if ctx.Err() != nil {
+				return fmt.Errorf("failed to acquire publish lock: timeout: %w", ctx.Err())
+			}
+
+			// Wait before retrying
+			time.Sleep(backoff)
+			backoff = min(backoff*2, 500*time.Millisecond)
+			continue
+		}
+
+		// Lock acquired! Store transaction in context
+		ctxWithTx := context.WithValue(ctx, publishTxKey, tx)
+
+		// Execute the function with the transaction-aware context
+		fnErr := fn(ctxWithTx)
+		if fnErr != nil {
+			_ = tx.Rollback(ctx)
+			return fnErr
+		}
+
+		// Commit the transaction (which also releases the lock)
+		if err := tx.Commit(ctx); err != nil {
+			return fmt.Errorf("failed to commit transaction: %w", err)
+		}
+
+		return nil
+	}
 }
 
 // hashServerName creates a consistent hash of the server name for advisory locking

internal/service/registry_service.go

@@ -123,7 +123,8 @@ func (s *registryServiceImpl) GetAllVersionsByServerID(serverID string) ([]apiv0
 // Publish publishes a server with flattened _meta extensions
 func (s *registryServiceImpl) Publish(req apiv0.ServerJSON) (*apiv0.ServerJSON, error) {
 	// Create a timeout context for the database operation
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	// Longer timeout to handle high concurrency scenarios where operations are serialized by advisory locks
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
 	defer cancel()
 
 	// Validate the request

internal/service/registry_service_test.go

@@ -384,7 +384,7 @@ func TestPublishConcurrentVersionsNoRace(t *testing.T) {
 	testDB := database.NewTestDB(t)
 	service := NewRegistryService(testDB, &config.Config{EnableRegistryValidation: false})
 
-	const concurrency = 1 // @Maintainers: Fix this and increase to higher number, previously 100
+	const concurrency = 100
 	results := make([]*apiv0.ServerJSON, concurrency)
 	errors := make([]error, concurrency)
 
@@ -419,12 +419,18 @@ func TestPublishConcurrentVersionsNoRace(t *testing.T) {
 		}
 	}
 
-	latestCount := 0
-	for _, result := range results {
-		if result != nil && result.Meta != nil && result.Meta.Official != nil &&
-			result.Meta.Official.IsLatest {
-			latestCount++
+	// Check what's actually in the database (the source of truth)
+	// Note: returned results are snapshots from when each was published, so they may show
+	// multiple versions as latest if they were latest at different points in time
+	filter := &database.ServerFilter{Name: func() *string { s := "com.example/test-concurrent"; return &s }()}
+	dbVersions, _, err := testDB.List(context.Background(), filter, "", 1000)
+	assert.NoError(t, err)
+	dbLatestCount := 0
+	for _, v := range dbVersions {
+		if v.Meta != nil && v.Meta.Official != nil && v.Meta.Official.IsLatest {
+			dbLatestCount++
 		}
 	}
-	assert.Equal(t, 1, latestCount, "should have exactly one latest version")
+
+	assert.Equal(t, 1, dbLatestCount, "should have exactly one latest version in database")
 }