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.
This commit is contained in:
adam jones 2025-09-22 12:16:05 +01:00 • committed by GitHub
commit 60c801c0c2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 3189 additions and 60 deletions

View file

@ -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

View file

@ -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

View file

@ -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;

View file

@ -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()