Fix race conditions in concurrent server publishing (#528)

## Summary

Fixes race conditions in concurrent server publishing that could result
in data integrity violations:
- Multiple versions marked as `isLatest=true` for the same server
- Duplicate (serverId, version) combinations
- Same server name assigned different serverIDs

## Solution

Uses two complementary approaches for correctness:

### 1. PostgreSQL Advisory Locks
- Serializes concurrent publishes per server name using
`pg_advisory_xact_lock`
- Uses FNV-1a hash to convert server name to lock ID
- Transaction-scoped locks auto-release on commit/rollback
- **Required after #465**: The serverId introduction means concurrent
publishes must atomically coordinate whether to reuse an existing
serverId (for additional versions) or generate a new one (for first
version). Without locks, race conditions would assign different
serverIDs to versions of the same server.

### 2. Unique Constraints (Migration 007)
Two new indexes provide data integrity guarantees:
- **idx_unique_server_version**: Ensures (serverId, version) is unique
across all versions
- **idx_unique_latest_version**: Partial index ensuring only one version
per serverId has `isLatest=true`

These constraints serve as defense-in-depth safety nets. Even if lock
logic had a bug, the database would reject integrity violations.

Also drops redundant `idx_servers_version_id` since `version_id` is
already the primary key.

## Why WithPublishLock validation is needed after #465

The `ctx.Err()` check at the start of `WithPublishLock` is critical
because:
1. After #465 introduced serverId, the publish flow must atomically
check existing versions and assign/reuse serverId
2. If the context is cancelled during lock acquisition, we could
deadlock
3. Checking ctx validity before acquiring locks prevents these deadlocks

## Testing

Added `TestPublishConcurrentVersionsNoRace` with 100 concurrent
goroutines publishing different versions of the same server. All succeed
with:
- Exactly one version marked as latest
- All versions sharing the same serverId
- No duplicate (serverId, version) combinations

NB: this only covers the in memory db, not the postgres db. It'd be nice
to write tests for this one as well, but our current setup is a bit
annoying for this. Separately looking into ways to easily run all our
tests with real postgres and ideally scrap memory.go completely.

internal/database/database.go

@@ -44,10 +44,34 @@ type Database interface {
 	CreateServer(ctx context.Context, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error)
 	// UpdateServer updates an existing server record
 	UpdateServer(ctx context.Context, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error)
+	// WithPublishLock executes a function with an exclusive lock for publishing a server
+	// This prevents race conditions when multiple versions are published concurrently
+	WithPublishLock(ctx context.Context, serverName string, fn func(ctx context.Context) error) error
 	// Close closes the database connection
 	Close() error
 }
 
+// WithPublishLockT is a generic helper that wraps WithPublishLock for functions returning a value
+// This exists because Go does not support generic methods on interfaces - only the Database interface
+// method WithPublishLock (without generics) can exist, so we provide this generic wrapper function.
+// This is a common pattern in Go for working around this language limitation.
+func WithPublishLockT[T any](ctx context.Context, db Database, serverName string, fn func(ctx context.Context) (T, error)) (T, error) {
+	var result T
+	var fnErr error
+
+	err := db.WithPublishLock(ctx, serverName, func(lockCtx context.Context) error {
+		result, fnErr = fn(lockCtx)
+		return fnErr
+	})
+
+	if err != nil {
+		var zero T
+		return zero, err
+	}
+
+	return result, nil
+}
+
 // ConnectionType represents the type of database connection
 type ConnectionType string
 

internal/database/memory.go

@@ -12,15 +12,18 @@ import (
 
 // MemoryDB is an in-memory implementation of the Database interface
 type MemoryDB struct {
-	entries map[string]*apiv0.ServerJSON // maps registry metadata version_id to ServerJSON
-	mu      sync.RWMutex
+	entries     map[string]*apiv0.ServerJSON // maps registry metadata version_id to ServerJSON
+	mu          sync.RWMutex
+	publishLocks map[string]*sync.Mutex       // per-server-name locks for publish operations
+	locksMu      sync.Mutex                   // protects publishLocks map
 }
 
 func NewMemoryDB() *MemoryDB {
 	// Convert input ServerJSON entries to have proper metadata
 	serverRecords := make(map[string]*apiv0.ServerJSON)
 	return &MemoryDB{
-		entries: serverRecords,
+		entries:      serverRecords,
+		publishLocks: make(map[string]*sync.Mutex),
 	}
 }
 
@@ -233,6 +236,27 @@ func (db *MemoryDB) UpdateServer(ctx context.Context, id string, server *apiv0.S
 	return server, nil
 }
 
+func (db *MemoryDB) WithPublishLock(ctx context.Context, serverName string, fn func(ctx context.Context) error) error {
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
+	// Get or create a lock for this specific server name
+	db.locksMu.Lock()
+	lock, exists := db.publishLocks[serverName]
+	if !exists {
+		lock = &sync.Mutex{}
+		db.publishLocks[serverName] = lock
+	}
+	db.locksMu.Unlock()
+
+	// Acquire the server-specific lock
+	lock.Lock()
+	defer lock.Unlock()
+
+	return fn(ctx)
+}
+
 // For an in-memory database, this is a no-op
 func (db *MemoryDB) Close() error {
 	return nil

internal/database/migrations/007_add_publish_constraints.sql

@@ -0,0 +1,21 @@
+-- Add unique constraints to prevent race conditions in concurrent publishing
+-- These serve as safety nets alongside advisory locks
+
+-- Drop redundant index (version_id is already the primary key)
+DROP INDEX IF EXISTS idx_servers_version_id;
+
+-- Prevent duplicate versions of the same server
+-- Ensures (serverId, version) is unique across all server versions
+CREATE UNIQUE INDEX idx_unique_server_version
+ON servers (
+    (value->'_meta'->'io.modelcontextprotocol.registry/official'->>'serverId'),
+    (value->>'version')
+);
+
+-- Prevent multiple versions marked as latest for the same server
+-- Ensures only one version per serverId can have isLatest=true
+CREATE UNIQUE INDEX idx_unique_latest_version
+ON servers (
+    (value->'_meta'->'io.modelcontextprotocol.registry/official'->>'serverId')
+)
+WHERE (value->'_meta'->'io.modelcontextprotocol.registry/official'->>'isLatest')::boolean = true;

internal/database/postgres.go

@@ -402,6 +402,61 @@ func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, server *apiv0
 	return server, nil
 }
 
+// WithPublishLock executes a function with an exclusive advisory lock for publishing a server
+// This prevents race conditions when multiple versions are published concurrently
+func (db *PostgreSQL) WithPublishLock(ctx context.Context, serverName string, fn func(ctx context.Context) error) error {
+	if ctx.Err() != nil {
+		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)
+	}
+
+	// Execute the function
+	if err := fn(ctx); err != nil {
+		return 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)
+	}
+
+	return nil
+}
+
+// hashServerName creates a consistent hash of the server name for advisory locking
+// We use FNV-1a hash and mask to 63 bits to fit in PostgreSQL's bigint range
+func hashServerName(name string) int64 {
+	// FNV-1a 64-bit hash
+	const (
+		offset64 = 14695981039346656037
+		prime64  = 1099511628211
+	)
+	hash := uint64(offset64)
+	for i := 0; i < len(name); i++ {
+		hash ^= uint64(name[i])
+		hash *= prime64
+	}
+	// Use only 63 bits to ensure positive int64
+	//nolint:gosec // Intentional conversion with masking to 63 bits
+	return int64(hash & 0x7FFFFFFFFFFFFFFF)
+}
+
 // Close closes the database connection
 func (db *PostgreSQL) Close() error {
 	db.pool.Close()

internal/service/registry_service.go

@@ -131,76 +131,85 @@ func (s *registryServiceImpl) Publish(req apiv0.ServerJSON) (*apiv0.ServerJSON,
 		return nil, err
 	}
 
-	publishTime := time.Now()
-	serverJSON := req
+	// Acquire advisory lock for this server name to prevent race conditions
+	result, err := database.WithPublishLockT(ctx, s.db, req.Name, func(lockCtx context.Context) (*apiv0.ServerJSON, error) {
+		publishTime := time.Now()
+		serverJSON := req
 
-	// Check for duplicate remote URLs
-	if err := s.validateNoDuplicateRemoteURLs(ctx, serverJSON); err != nil {
-		return nil, err
-	}
-
-	filter := &database.ServerFilter{Name: &serverJSON.Name}
-	existingServerVersions, _, err := s.db.List(ctx, filter, "", maxServerVersionsPerServer)
-	if err != nil && !errors.Is(err, database.ErrNotFound) {
-		return nil, err
-	}
-
-	// Check we haven't exceeded the maximum versions allowed for a server
-	if len(existingServerVersions) >= maxServerVersionsPerServer {
-		return nil, database.ErrMaxServersReached
-	}
-
-	// Check this isn't a duplicate version
-	for _, server := range existingServerVersions {
-		existingVersion := server.Version
-		if existingVersion == serverJSON.Version {
-			return nil, database.ErrInvalidVersion
+		// Check for duplicate remote URLs
+		if err := s.validateNoDuplicateRemoteURLs(lockCtx, serverJSON); err != nil {
+			return nil, err
 		}
-	}
 
-	// Determine if this version should be marked as latest
-	existingLatest := s.getCurrentLatestVersion(existingServerVersions)
-	isNewLatest := true
-	if existingLatest != nil {
-		var existingPublishedAt time.Time
-		if existingLatest.Meta != nil && existingLatest.Meta.Official != nil {
-			existingPublishedAt = existingLatest.Meta.Official.PublishedAt
+		filter := &database.ServerFilter{Name: &serverJSON.Name}
+		existingServerVersions, _, err := s.db.List(lockCtx, filter, "", maxServerVersionsPerServer)
+		if err != nil && !errors.Is(err, database.ErrNotFound) {
+			return nil, err
 		}
-		isNewLatest = CompareVersions(
-			serverJSON.Version,
-			existingLatest.Version,
-			publishTime,
-			existingPublishedAt,
-		) > 0
-	}
 
-	// Create complete server with metadata
-	server := s.createServerWithMetadata(serverJSON, existingServerVersions, publishTime, isNewLatest)
+		// Check we haven't exceeded the maximum versions allowed for a server
+		if len(existingServerVersions) >= maxServerVersionsPerServer {
+			return nil, database.ErrMaxServersReached
+		}
+
+		// Check this isn't a duplicate version
+		for _, server := range existingServerVersions {
+			existingVersion := server.Version
+			if existingVersion == serverJSON.Version {
+				return nil, database.ErrInvalidVersion
+			}
+		}
+
+		// Determine if this version should be marked as latest
+		existingLatest := s.getCurrentLatestVersion(existingServerVersions)
+		isNewLatest := true
+		if existingLatest != nil {
+			var existingPublishedAt time.Time
+			if existingLatest.Meta != nil && existingLatest.Meta.Official != nil {
+				existingPublishedAt = existingLatest.Meta.Official.PublishedAt
+			}
+			isNewLatest = CompareVersions(
+				serverJSON.Version,
+				existingLatest.Version,
+				publishTime,
+				existingPublishedAt,
+			) > 0
+		}
+
+		// Mark previous latest as no longer latest BEFORE creating new version
+		// This prevents violating the unique constraint on isLatest
+		if isNewLatest && existingLatest != nil {
+			var existingLatestVersionID string
+			if existingLatest.Meta != nil && existingLatest.Meta.Official != nil {
+				existingLatestVersionID = existingLatest.Meta.Official.VersionID
+			}
+			if existingLatestVersionID != "" {
+				// Update the existing server to set isLatest = false
+				existingLatest.Meta.Official.IsLatest = false
+				existingLatest.Meta.Official.UpdatedAt = time.Now()
+				if _, err := s.db.UpdateServer(lockCtx, existingLatestVersionID, existingLatest); err != nil {
+					return nil, err
+				}
+			}
+		}
+
+		// Create complete server with metadata
+		server := s.createServerWithMetadata(serverJSON, existingServerVersions, publishTime, isNewLatest)
+
+		// Create server in database
+		serverRecord, err := s.db.CreateServer(lockCtx, &server)
+		if err != nil {
+			return nil, err
+		}
+
+		return serverRecord, nil
+	})
 
-	// Create server in database
-	serverRecord, err := s.db.CreateServer(ctx, &server)
 	if err != nil {
 		return nil, err
 	}
 
-	// Mark previous latest as no longer latest
-	if isNewLatest && existingLatest != nil {
-		var existingLatestVersionID string
-		if existingLatest.Meta != nil && existingLatest.Meta.Official != nil {
-			existingLatestVersionID = existingLatest.Meta.Official.VersionID
-		}
-		if existingLatestVersionID != "" {
-			// Update the existing server to set isLatest = false
-			existingLatest.Meta.Official.IsLatest = false
-			existingLatest.Meta.Official.UpdatedAt = time.Now()
-			if _, err := s.db.UpdateServer(ctx, existingLatestVersionID, existingLatest); err != nil {
-				return nil, err
-			}
-		}
-	}
-
-	// Return the server record directly
-	return serverRecord, nil
+	return result, nil
 }
 
 // createServerWithMetadata creates a server with proper metadata including server_id and version_id

internal/service/registry_service_test.go

@@ -3,6 +3,8 @@ package service
 
 import (
 	"context"
+	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/modelcontextprotocol/registry/internal/config"
@@ -377,3 +379,52 @@ func TestGetAllVersionsByServerID(t *testing.T) {
 		})
 	}
 }
+
+func TestPublishConcurrentVersionsNoRace(t *testing.T) {
+	memDB := database.NewMemoryDB()
+	service := NewRegistryService(memDB, &config.Config{EnableRegistryValidation: false})
+
+	const concurrency = 100
+	results := make([]*apiv0.ServerJSON, concurrency)
+	errors := make([]error, concurrency)
+
+	var wg sync.WaitGroup
+	for i := range concurrency {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			result, err := service.Publish(apiv0.ServerJSON{
+				Name:        "com.example/test-concurrent",
+				Description: fmt.Sprintf("Version %d", idx),
+				Version:     fmt.Sprintf("1.0.%d", idx),
+			})
+			results[idx] = result
+			errors[idx] = err
+		}(i)
+	}
+	wg.Wait()
+
+	for i, err := range errors {
+		assert.NoError(t, err, "publish %d failed", i)
+	}
+
+	var sharedServerID string
+	for i, result := range results {
+		if result != nil && result.Meta != nil && result.Meta.Official != nil {
+			if sharedServerID == "" {
+				sharedServerID = result.Meta.Official.ServerID
+			}
+			assert.Equal(t, sharedServerID, result.Meta.Official.ServerID,
+				"version %d has different serverID", i)
+		}
+	}
+
+	latestCount := 0
+	for _, result := range results {
+		if result != nil && result.Meta != nil && result.Meta.Official != nil &&
+			result.Meta.Official.IsLatest {
+			latestCount++
+		}
+	}
+	assert.Equal(t, 1, latestCount, "should have exactly one latest version")
+}