registry/internal/database/memory.go
adam jones 60c801c0c2
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.
2025-09-22 12:16:05 +01:00

359 lines
8.7 KiB
Go

package database
import (
"context"
"fmt"
"sort"
"strings"
"sync"
apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0"
)
// 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
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,
publishLocks: make(map[string]*sync.Mutex),
}
}
func (db *MemoryDB) List(
ctx context.Context,
filter *ServerFilter,
cursor string,
limit int,
) ([]*apiv0.ServerJSON, string, error) {
if ctx.Err() != nil {
return nil, "", ctx.Err()
}
if limit <= 0 {
limit = 10 // Default limit
}
db.mu.RLock()
defer db.mu.RUnlock()
// Convert all entries to a slice for pagination
var allEntries []*apiv0.ServerJSON
for _, entry := range db.entries {
allEntries = append(allEntries, entry)
}
// Apply filtering and sorting
filteredEntries := db.filterAndSort(allEntries, filter)
// Find starting point for cursor-based pagination
startIdx := 0
if cursor != "" {
for i, entry := range filteredEntries {
if db.getRegistryID(entry) == cursor {
startIdx = i + 1 // Start after the cursor
break
}
}
}
// Apply pagination
endIdx := min(startIdx+limit, len(filteredEntries))
var result []*apiv0.ServerJSON
if startIdx < len(filteredEntries) {
result = filteredEntries[startIdx:endIdx]
} else {
result = []*apiv0.ServerJSON{}
}
// Determine next cursor
nextCursor := ""
if endIdx < len(filteredEntries) && len(result) > 0 {
nextCursor = db.getRegistryID(result[len(result)-1])
}
return result, nextCursor, nil
}
func (db *MemoryDB) GetByVersionID(ctx context.Context, versionID string) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
db.mu.RLock()
defer db.mu.RUnlock()
// Find entry by registry metadata version_id
if entry, exists := db.entries[versionID]; exists {
// Return a copy of the ServerRecord
entryCopy := *entry
return &entryCopy, nil
}
return nil, ErrNotFound
}
// GetByServerID retrieves the latest version of a server by server ID
func (db *MemoryDB) GetByServerID(ctx context.Context, serverID string) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
db.mu.RLock()
defer db.mu.RUnlock()
// Find the latest version with the given server ID
var latestEntry *apiv0.ServerJSON
for _, entry := range db.entries {
if entry.Meta != nil && entry.Meta.Official != nil &&
entry.Meta.Official.ServerID == serverID &&
entry.Meta.Official.IsLatest {
latestEntry = entry
break
}
}
if latestEntry == nil {
return nil, ErrNotFound
}
// Return a copy
entryCopy := *latestEntry
return &entryCopy, nil
}
// GetByServerIDAndVersion retrieves a specific version of a server by server ID and version
func (db *MemoryDB) GetByServerIDAndVersion(ctx context.Context, serverID string, version string) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
db.mu.RLock()
defer db.mu.RUnlock()
// Find the entry with matching server ID and version
for _, entry := range db.entries {
if entry.Meta != nil && entry.Meta.Official != nil &&
entry.Meta.Official.ServerID == serverID &&
entry.Version == version {
// Return a copy
entryCopy := *entry
return &entryCopy, nil
}
}
return nil, ErrNotFound
}
// GetAllVersionsByServerID retrieves all versions of a server by server ID
func (db *MemoryDB) GetAllVersionsByServerID(ctx context.Context, serverID string) ([]*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
db.mu.RLock()
defer db.mu.RUnlock()
var results []*apiv0.ServerJSON
for _, entry := range db.entries {
if entry.Meta != nil && entry.Meta.Official != nil &&
entry.Meta.Official.ServerID == serverID {
// Add a copy
entryCopy := *entry
results = append(results, &entryCopy)
}
}
if len(results) == 0 {
return nil, ErrNotFound
}
// Sort by published date, latest first
sort.Slice(results, func(i, j int) bool {
iTime := results[i].Meta.Official.PublishedAt
jTime := results[j].Meta.Official.PublishedAt
return iTime.After(jTime)
})
return results, nil
}
func (db *MemoryDB) CreateServer(ctx context.Context, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Get the VersionID from the registry metadata
if server.Meta == nil || server.Meta.Official == nil {
return nil, fmt.Errorf("server must have registry metadata with ServerID and VersionID")
}
versionID := server.Meta.Official.VersionID
if versionID == "" {
return nil, fmt.Errorf("server must have VersionID in registry metadata")
}
db.mu.Lock()
defer db.mu.Unlock()
// Store the record using registry metadata VersionID
db.entries[versionID] = server
return server, nil
}
func (db *MemoryDB) UpdateServer(ctx context.Context, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Validate that meta structure exists and VersionID matches path (consistent with PostgreSQL implementation)
if server.Meta == nil || server.Meta.Official == nil || server.Meta.Official.VersionID != id {
return nil, fmt.Errorf("%w: io.modelcontextprotocol.registry/official.version_id must match path id (%s)", ErrInvalidInput, id)
}
db.mu.Lock()
defer db.mu.Unlock()
_, exists := db.entries[id]
if !exists {
return nil, ErrNotFound
}
// Update the server
db.entries[id] = server
// Return the updated record
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
}
// filterAndSort applies filtering and sorting to the entries
func (db *MemoryDB) filterAndSort(allEntries []*apiv0.ServerJSON, filter *ServerFilter) []*apiv0.ServerJSON {
// Apply filtering
var filteredEntries []*apiv0.ServerJSON
for _, entry := range allEntries {
if db.matchesFilter(entry, filter) {
filteredEntries = append(filteredEntries, entry)
}
}
// Sort by registry metadata ID for consistent pagination
sort.Slice(filteredEntries, func(i, j int) bool {
iID := db.getRegistryID(filteredEntries[i])
jID := db.getRegistryID(filteredEntries[j])
return iID < jID
})
return filteredEntries
}
// matchesFilter checks if an entry matches the provided filter
//
//nolint:cyclop // Filter matching logic is inherently complex but clear
func (db *MemoryDB) matchesFilter(entry *apiv0.ServerJSON, filter *ServerFilter) bool {
if filter == nil {
return true
}
// Check name filter
if filter.Name != nil && entry.Name != *filter.Name {
return false
}
// Check remote URL filter
if filter.RemoteURL != nil {
found := false
for _, remote := range entry.Remotes {
if remote.URL == *filter.RemoteURL {
found = true
break
}
}
if !found {
return false
}
}
// Check updatedSince filter
if filter.UpdatedSince != nil {
if entry.Meta == nil || entry.Meta.Official == nil {
return false
}
if entry.Meta.Official.UpdatedAt.Before(*filter.UpdatedSince) ||
entry.Meta.Official.UpdatedAt.Equal(*filter.UpdatedSince) {
return false
}
}
// Check name search filter (substring match)
if filter.SubstringName != nil {
// Case-insensitive substring search
searchLower := strings.ToLower(*filter.SubstringName)
nameLower := strings.ToLower(entry.Name)
if !strings.Contains(nameLower, searchLower) {
return false
}
}
// Check exact version filter
if filter.Version != nil {
if entry.Version != *filter.Version {
return false
}
}
// Check isLatest filter
if filter.IsLatest != nil {
if entry.Meta == nil || entry.Meta.Official == nil {
return false
}
if entry.Meta.Official.IsLatest != *filter.IsLatest {
return false
}
}
return true
}
// getRegistryID safely extracts the registry version ID from an entry
func (db *MemoryDB) getRegistryID(entry *apiv0.ServerJSON) string {
if entry.Meta != nil && entry.Meta.Official != nil {
return entry.Meta.Official.VersionID
}
return ""
}