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:
parent
53f808a2c1
commit
60c801c0c2
4 changed files with 3189 additions and 60 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
21
internal/database/migrations/007_add_publish_constraints.sql
Normal file
21
internal/database/migrations/007_add_publish_constraints.sql
Normal 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;
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in a new issue