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
GitHub
parent
53f808a2c1
commit
60c801c0c2
2 changed files with 3189 additions and 60 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue