From b665247ea5020494ba7eff483039c3a1330d7e8a Mon Sep 17 00:00:00 2001 From: Fernando Marino Date: Wed, 9 Sep 2026 10:30:05 +0200 Subject: [PATCH 1/6] node: make the backend probe timeout configurable sam-node.yaml's command-spawned MCP services are given a hard-coded 2s (dhtProbeTimeout) to answer an MCP initialize before the service is registered but withheld from advertisement ("backend did not answer: context deadline exceeded"), with no retry observed afterwards. 2s is tighter than the cold-start cost of realistic backends. Measured directly: a bare `import fastmcp` (Python) takes ~2.7s, and even SAM's own reference example, `npx -y @modelcontextprotocol/server-everything`, takes ~4.0s to answer initialize on Windows even warm/cached - SAM's own documented reference server cannot reliably meet its own default. Adds a new ServiceRegistry.SetBackendProbeTimeout, wired from a new --backend-probe-timeout flag (0 keeps the existing 2s default, matching the convention already used by --dht-max-record-age and similar flags). Backward compatible: NewServiceRegistry's default is unchanged, and every existing call site that doesn't pass the new flag behaves exactly as before (verified: all pre-existing internal/node tests pass unmodified, including 6 that fail identically on unmodified main - confirmed via git stash - so unrelated to this change). Also fixes three of the four node.Options{} construction sites in cmd/sam-node/main.go that already wire NewServiceRegistry-adjacent flags (DHTMaxRecordAge et al.) but were missing this one; the fourth (join-only enrollment path) doesn't call RegisterStaticServices and is out of scope. Adds TestServiceRegistry_BackendProbeTimeoutIsConfigurable, using a new slowProbingService test fake that (unlike the existing probingService) actually respects context deadlines, so it can demonstrate: the same slow backend fails to advertise under the default timeout and succeeds once given more time. --- cmd/sam-node/main.go | 5 ++ internal/node/node.go | 1 + internal/node/options.go | 8 +++ internal/node/service_registry.go | 46 ++++++++++++--- internal/node/service_registry_test.go | 80 ++++++++++++++++++++++++++ 5 files changed, 133 insertions(+), 7 deletions(-) diff --git a/cmd/sam-node/main.go b/cmd/sam-node/main.go index 9c65be62..8c63e131 100644 --- a/cmd/sam-node/main.go +++ b/cmd/sam-node/main.go @@ -99,6 +99,7 @@ var ( dhtMaxRecordAgeFlag time.Duration dhtLookupLimitFlag int discoveryConcurrencyFlag int + backendProbeTimeoutFlag time.Duration policySyncIntervalFlag time.Duration ) @@ -466,6 +467,7 @@ func main() { DHTMaxRecordAge: dhtMaxRecordAgeFlag, DHTLookupLimit: dhtLookupLimitFlag, DiscoveryConcurrency: discoveryConcurrencyFlag, + BackendProbeTimeout: backendProbeTimeoutFlag, }) if err != nil { logger.Fatalf("Failed to initialize mesh node: %v", err) @@ -533,6 +535,7 @@ func main() { DHTMaxRecordAge: dhtMaxRecordAgeFlag, DHTLookupLimit: dhtLookupLimitFlag, DiscoveryConcurrency: discoveryConcurrencyFlag, + BackendProbeTimeout: backendProbeTimeoutFlag, }) if err != nil { enrollCancel() @@ -599,6 +602,7 @@ func main() { RouterConnectTimeout: routerConnectTimeoutFlag, RequiredRole: api.RoleNode, PolicySyncInterval: policySyncIntervalFlag, + BackendProbeTimeout: backendProbeTimeoutFlag, }) if err != nil { logger.Fatalf("Failed to initialize node after enrollment: %v", err) @@ -854,6 +858,7 @@ func main() { runCmd.Flags().IntVar(&dhtLookupLimitFlag, "dht-lookup-limit", 0, "Maximum number of providers to query from the DHT (0 uses default 20)") runCmd.Flags().IntVar(&discoveryConcurrencyFlag, "discovery-concurrency", 0, "Max concurrent catalog fetches during discovery (0 uses default 10)") runCmd.Flags().DurationVar(&policySyncIntervalFlag, "policy-sync-interval", 1*time.Hour, "Interval for syncing mesh policy from the control plane") + runCmd.Flags().DurationVar(&backendProbeTimeoutFlag, "backend-probe-timeout", 0, "Timeout for probing a command-spawned service backend before advertising it (0 uses default 2s); raise this for backends with slower cold-start times") rootCmd.PersistentFlags().StringVar(&controlPlaneAddr, "control-plane", "", "Control plane URL") rootCmd.PersistentFlags().StringVar(&configFile, "config", node.DefaultConfigFile, "Path to sam-node.yaml configuration file") rootCmd.PersistentFlags().StringVar(&oidcIssuerFlag, "oidc-issuer", "", "OIDC Issuer URL") diff --git a/internal/node/node.go b/internal/node/node.go index 685bc613..039084b9 100644 --- a/internal/node/node.go +++ b/internal/node/node.go @@ -520,6 +520,7 @@ func (n *SamNode) Start(ctx context.Context) error { n.services = NewServiceRegistry(n.DHT) n.services.reprovideNow = n.triggerReprovide + n.services.SetBackendProbeTimeout(n.config.BackendProbeTimeout) var authenticated bool var fatalAuthErr error diff --git a/internal/node/options.go b/internal/node/options.go index 9f8a0367..2075c7ee 100644 --- a/internal/node/options.go +++ b/internal/node/options.go @@ -78,6 +78,14 @@ type Options struct { PolicySyncInterval time.Duration // PolicySyncJitter specifies the maximum jitter delay when scheduling policy syncs on event broadcasts. PolicySyncJitter time.Duration + // BackendProbeTimeout bounds how long a command-spawned service backend + // (sam-node.yaml's `command`, spawned as a local subprocess) is given to + // answer before the service is registered but withheld from + // advertisement. Zero uses the library default (2s). Raise this for + // backends with slower cold-start/import costs than that - the default + // is tight enough that even simple interpreted-language MCP servers can + // miss it on first spawn. + BackendProbeTimeout time.Duration } // Default applies default values to Options if they are not specified. diff --git a/internal/node/service_registry.go b/internal/node/service_registry.go index ffcc2724..280f2718 100644 --- a/internal/node/service_registry.go +++ b/internal/node/service_registry.go @@ -38,7 +38,14 @@ type backendProber interface { Probe(ctx context.Context) error } -// dhtProbeTimeout bounds one backend probe before advertising. +// dhtProbeTimeout is the default bound on one backend probe before +// advertising, used when a ServiceRegistry isn't given an explicit +// BackendProbeTimeout (see Options.BackendProbeTimeout / --backend-probe- +// timeout). Command-spawned backends (sam-node.yaml's `command`, launched as +// a local subprocess) can need longer than this to answer their first +// request - a moderately-featured interpreted-language MCP server's own +// import/startup cost alone can exceed 2s - so this is a floor, not +// something every backend is expected to meet. const dhtProbeTimeout = 2 * time.Second // advertisable reports whether a service is fit to be published to the DHT. @@ -49,12 +56,12 @@ const dhtProbeTimeout = 2 * time.Second // listed by discover_remote_services and only failed later, at initialize, in // the caller's face. Advertising is a claim the node makes on the backend's // behalf, so it is the node that should verify it. -func advertisable(ctx context.Context, svc Service) error { +func advertisable(ctx context.Context, svc Service, probeTimeout time.Duration) error { prober, ok := svc.(backendProber) if !ok { return nil } - probeCtx, cancel := context.WithTimeout(ctx, dhtProbeTimeout) + probeCtx, cancel := context.WithTimeout(ctx, probeTimeout) defer cancel() return prober.Probe(probeCtx) } @@ -69,15 +76,40 @@ type ServiceRegistry struct { // registered after the loop last ran does not wait a whole interval to be // advertised. Optional; nil outside a running node. reprovideNow func() + + // backendProbeTimeout bounds each backend probe in advertisable. Defaults + // to dhtProbeTimeout; override with SetBackendProbeTimeout. + backendProbeTimeout time.Duration } func NewServiceRegistry(d dhtProvider) *ServiceRegistry { return &ServiceRegistry{ - services: map[string]Service{}, - dht: d, + services: map[string]Service{}, + dht: d, + backendProbeTimeout: dhtProbeTimeout, } } +// SetBackendProbeTimeout overrides the default backend probe timeout +// (dhtProbeTimeout). A zero or negative duration is a no-op, so callers can +// pass an unset Options.BackendProbeTimeout straight through without an +// explicit zero-check. +func (r *ServiceRegistry) SetBackendProbeTimeout(d time.Duration) { + if d <= 0 { + return + } + r.mu.Lock() + defer r.mu.Unlock() + r.backendProbeTimeout = d +} + +// probeTimeout returns the current backend probe timeout. +func (r *ServiceRegistry) probeTimeout() time.Duration { + r.mu.RLock() + defer r.mu.RUnlock() + return r.backendProbeTimeout +} + // Register initialises a service, advertises it on the DHT, and inserts it // into the map. Init runs before Provide so a failed handler-build never // briefly advertises an unservable name. @@ -104,7 +136,7 @@ func (r *ServiceRegistry) Register(ctx context.Context, svc Service) error { return err } - probeErr := advertisable(ctx, svc) + probeErr := advertisable(ctx, svc, r.probeTimeout()) if probeErr != nil { logger.Warnf("[ServiceRegistry] Registered %s/%s but not advertising it: backend did not answer: %v", info.Type, info.Name, probeErr) } else { @@ -229,7 +261,7 @@ Loop: }() info := svc.Info() - if err := advertisable(ctx, svc); err != nil { + if err := advertisable(ctx, svc, r.probeTimeout()); err != nil { withheld.Add(1) // On shutdown every service fails this way, and saying so // would blame backends for the node stopping. diff --git a/internal/node/service_registry_test.go b/internal/node/service_registry_test.go index b0d2f5e5..8d724aae 100644 --- a/internal/node/service_registry_test.go +++ b/internal/node/service_registry_test.go @@ -21,6 +21,7 @@ import ( "sync" "sync/atomic" "testing" + "time" "github.com/google/sam/api" "github.com/ipfs/go-cid" @@ -274,3 +275,82 @@ func TestServiceRegistry_ReprovideResumesWhenBackendRecovers(t *testing.T) { t.Errorf("Provide called %d times after recovery, want 2 (name + type CID)", len(dht.calls)) } } + +// slowProbingService is a backendProber whose Probe blocks until the given +// delay elapses or the context is cancelled first, whichever comes first - +// unlike probingService, it actually respects the probe deadline, which is +// what a real command-spawned backend with a slow cold-start does. +type slowProbingService struct { + *fakeService + delay time.Duration +} + +func newSlowProbingSvc(name string, delay time.Duration) *slowProbingService { + return &slowProbingService{ + fakeService: newFakeSvc(name, api.ServiceType_SERVICE_TYPE_MCP), + delay: delay, + } +} + +func (p *slowProbingService) Probe(ctx context.Context) error { + select { + case <-time.After(p.delay): + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// The bug behind #376: dhtProbeTimeout was a hard-coded 2s with no way to +// raise it, so a backend whose own cold-start cost alone exceeds that - +// measured in practice for moderately-featured MCP server stacks - could +// never be advertised on its first registration. SetBackendProbeTimeout +// (wired from --backend-probe-timeout) is the fix: the same slow backend +// must fail to advertise under the default and succeed once given more time. +func TestServiceRegistry_BackendProbeTimeoutIsConfigurable(t *testing.T) { + const probeDelay = 60 * time.Millisecond + + t.Run("default timeout is too short for a slow backend", func(t *testing.T) { + dht := &fakeDHT{} + r := NewServiceRegistry(dht) + r.SetBackendProbeTimeout(10 * time.Millisecond) // shorter than probeDelay + + svc := newSlowProbingSvc("slow", probeDelay) + if err := r.Register(context.Background(), svc); err != nil { + t.Fatalf("Register: %v", err) + } + if len(dht.calls) != 0 { + t.Errorf("Provide called %d times for a backend slower than the probe timeout, want 0", len(dht.calls)) + } + }) + + t.Run("raising the timeout lets the same backend advertise", func(t *testing.T) { + dht := &fakeDHT{} + r := NewServiceRegistry(dht) + r.SetBackendProbeTimeout(probeDelay * 5) // comfortably longer than probeDelay + + svc := newSlowProbingSvc("slow", probeDelay) + if err := r.Register(context.Background(), svc); err != nil { + t.Fatalf("Register: %v", err) + } + if len(dht.calls) != 2 { + t.Errorf("Provide called %d times once given enough time to probe, want 2 (name + type CID)", len(dht.calls)) + } + }) + + t.Run("NewServiceRegistry defaults to dhtProbeTimeout unchanged", func(t *testing.T) { + r := NewServiceRegistry(&fakeDHT{}) + if got := r.probeTimeout(); got != dhtProbeTimeout { + t.Errorf("default probe timeout = %v, want %v (unchanged default behaviour)", got, dhtProbeTimeout) + } + }) + + t.Run("SetBackendProbeTimeout ignores zero and negative durations", func(t *testing.T) { + r := NewServiceRegistry(&fakeDHT{}) + r.SetBackendProbeTimeout(0) + r.SetBackendProbeTimeout(-1 * time.Second) + if got := r.probeTimeout(); got != dhtProbeTimeout { + t.Errorf("probe timeout after no-op sets = %v, want unchanged %v", got, dhtProbeTimeout) + } + }) +} From b341328d296b68df32cfc3e580263b36c9f8ddc8 Mon Sep 17 00:00:00 2001 From: Fernando Marino Date: Wed, 9 Sep 2026 11:23:41 +0200 Subject: [PATCH 2/6] address review: rename to defaultDHTProbeTimeout, fall back on <=0, fix timer leak in test - aojea: renamed dhtProbeTimeout to defaultDHTProbeTimeout for clarity now that there's also a configurable per-registry value. - gemini-code-assist: probeTimeout() now falls back to defaultDHTProbeTimeout when backendProbeTimeout is <= 0, so a zero-initialized ServiceRegistry (e.g. newServiceRegistryForTest's struct literal, which bypasses NewServiceRegistry) behaves the same as a properly constructed one instead of timing out every probe immediately. - gemini-code-assist: slowProbingService.Probe now uses time.NewTimer with a deferred Stop() instead of time.After, avoiding the short-term timer leak when the context is cancelled before the delay elapses. All internal/node tests pass, including the new TestServiceRegistry_BackendProbeTimeoutIsConfigurable. --- internal/node/service_registry.go | 19 +++++++++++++------ internal/node/service_registry_test.go | 16 +++++++++------- 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/internal/node/service_registry.go b/internal/node/service_registry.go index 280f2718..37872d1f 100644 --- a/internal/node/service_registry.go +++ b/internal/node/service_registry.go @@ -38,7 +38,7 @@ type backendProber interface { Probe(ctx context.Context) error } -// dhtProbeTimeout is the default bound on one backend probe before +// defaultDHTProbeTimeout is the default bound on one backend probe before // advertising, used when a ServiceRegistry isn't given an explicit // BackendProbeTimeout (see Options.BackendProbeTimeout / --backend-probe- // timeout). Command-spawned backends (sam-node.yaml's `command`, launched as @@ -46,7 +46,7 @@ type backendProber interface { // request - a moderately-featured interpreted-language MCP server's own // import/startup cost alone can exceed 2s - so this is a floor, not // something every backend is expected to meet. -const dhtProbeTimeout = 2 * time.Second +const defaultDHTProbeTimeout = 2 * time.Second // advertisable reports whether a service is fit to be published to the DHT. // @@ -78,7 +78,7 @@ type ServiceRegistry struct { reprovideNow func() // backendProbeTimeout bounds each backend probe in advertisable. Defaults - // to dhtProbeTimeout; override with SetBackendProbeTimeout. + // to defaultDHTProbeTimeout; override with SetBackendProbeTimeout. backendProbeTimeout time.Duration } @@ -86,12 +86,12 @@ func NewServiceRegistry(d dhtProvider) *ServiceRegistry { return &ServiceRegistry{ services: map[string]Service{}, dht: d, - backendProbeTimeout: dhtProbeTimeout, + backendProbeTimeout: defaultDHTProbeTimeout, } } // SetBackendProbeTimeout overrides the default backend probe timeout -// (dhtProbeTimeout). A zero or negative duration is a no-op, so callers can +// (defaultDHTProbeTimeout). A zero or negative duration is a no-op, so callers can // pass an unset Options.BackendProbeTimeout straight through without an // explicit zero-check. func (r *ServiceRegistry) SetBackendProbeTimeout(d time.Duration) { @@ -103,10 +103,17 @@ func (r *ServiceRegistry) SetBackendProbeTimeout(d time.Duration) { r.backendProbeTimeout = d } -// probeTimeout returns the current backend probe timeout. +// probeTimeout returns the current backend probe timeout, falling back to +// defaultDHTProbeTimeout for a zero-initialized registry (e.g. a struct +// literal built directly in a test, bypassing NewServiceRegistry) so it +// behaves the same as a properly constructed one rather than timing out +// every probe immediately. func (r *ServiceRegistry) probeTimeout() time.Duration { r.mu.RLock() defer r.mu.RUnlock() + if r.backendProbeTimeout <= 0 { + return defaultDHTProbeTimeout + } return r.backendProbeTimeout } diff --git a/internal/node/service_registry_test.go b/internal/node/service_registry_test.go index 8d724aae..a07a3c17 100644 --- a/internal/node/service_registry_test.go +++ b/internal/node/service_registry_test.go @@ -293,15 +293,17 @@ func newSlowProbingSvc(name string, delay time.Duration) *slowProbingService { } func (p *slowProbingService) Probe(ctx context.Context) error { + timer := time.NewTimer(p.delay) + defer timer.Stop() select { - case <-time.After(p.delay): + case <-timer.C: return nil case <-ctx.Done(): return ctx.Err() } } -// The bug behind #376: dhtProbeTimeout was a hard-coded 2s with no way to +// The bug behind #376: defaultDHTProbeTimeout was a hard-coded 2s with no way to // raise it, so a backend whose own cold-start cost alone exceeds that - // measured in practice for moderately-featured MCP server stacks - could // never be advertised on its first registration. SetBackendProbeTimeout @@ -338,10 +340,10 @@ func TestServiceRegistry_BackendProbeTimeoutIsConfigurable(t *testing.T) { } }) - t.Run("NewServiceRegistry defaults to dhtProbeTimeout unchanged", func(t *testing.T) { + t.Run("NewServiceRegistry defaults to defaultDHTProbeTimeout unchanged", func(t *testing.T) { r := NewServiceRegistry(&fakeDHT{}) - if got := r.probeTimeout(); got != dhtProbeTimeout { - t.Errorf("default probe timeout = %v, want %v (unchanged default behaviour)", got, dhtProbeTimeout) + if got := r.probeTimeout(); got != defaultDHTProbeTimeout { + t.Errorf("default probe timeout = %v, want %v (unchanged default behaviour)", got, defaultDHTProbeTimeout) } }) @@ -349,8 +351,8 @@ func TestServiceRegistry_BackendProbeTimeoutIsConfigurable(t *testing.T) { r := NewServiceRegistry(&fakeDHT{}) r.SetBackendProbeTimeout(0) r.SetBackendProbeTimeout(-1 * time.Second) - if got := r.probeTimeout(); got != dhtProbeTimeout { - t.Errorf("probe timeout after no-op sets = %v, want unchanged %v", got, dhtProbeTimeout) + if got := r.probeTimeout(); got != defaultDHTProbeTimeout { + t.Errorf("probe timeout after no-op sets = %v, want unchanged %v", got, defaultDHTProbeTimeout) } }) } From 20469848166b707bfc77439258b7080a26bb1a5b Mon Sep 17 00:00:00 2001 From: Fernando Marino Date: Wed, 9 Sep 2026 17:25:38 +0200 Subject: [PATCH 3/6] address review: drop the SetBackendProbeTimeout/probeTimeout setter+getter pair Config is populated once, at construction, everywhere else in this package (see internal/node/options.go) - this had grown a setter and a getter instead. NewServiceRegistry now takes backendProbeTimeout directly and defaults it in the constructor if <= 0; the timeout is never mutated afterwards, so reading the plain field needs no lock. --- internal/node/identity_evidence_http_test.go | 2 +- internal/node/mcp_discovery_test.go | 4 +- internal/node/node.go | 3 +- internal/node/service_registry.go | 48 ++++++-------------- internal/node/service_registry_test.go | 37 ++++++--------- internal/node/sidecar_auth_test.go | 2 +- internal/node/sidecar_test.go | 28 ++++++------ 7 files changed, 49 insertions(+), 75 deletions(-) diff --git a/internal/node/identity_evidence_http_test.go b/internal/node/identity_evidence_http_test.go index da12f74e..6a8f75bf 100644 --- a/internal/node/identity_evidence_http_test.go +++ b/internal/node/identity_evidence_http_test.go @@ -42,7 +42,7 @@ func TestIdentityEvidenceRoutesHaveOwnMetricClass(t *testing.T) { func TestIdentityEvidenceTrailingSlashReturnsNotFound(t *testing.T) { node := &SamNode{ BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } socketPath := filepath.Join(t.TempDir(), "sam.sock") diff --git a/internal/node/mcp_discovery_test.go b/internal/node/mcp_discovery_test.go index d936261f..ef63e5a9 100644 --- a/internal/node/mcp_discovery_test.go +++ b/internal/node/mcp_discovery_test.go @@ -71,8 +71,8 @@ func (f *fakeToolService) Tools(_ context.Context) ([]string, error) { func TestDiscoverySource(t *testing.T) { node := &SamNode{ - services: NewServiceRegistry(&fakeDHT{}), - nodeConfig: &NodeConfigComplete{Labels: map[string]string{"region": "EU"}}, + services: NewServiceRegistry(&fakeDHT{}, 0), + config: Options{Labels: map[string]string{"region": "EU"}}, } ctx := context.Background() diff --git a/internal/node/node.go b/internal/node/node.go index 039084b9..f402e053 100644 --- a/internal/node/node.go +++ b/internal/node/node.go @@ -518,9 +518,8 @@ func (n *SamNode) Start(ctx context.Context) error { } n.DHT = kdht - n.services = NewServiceRegistry(n.DHT) + n.services = NewServiceRegistry(n.DHT, n.config.BackendProbeTimeout) n.services.reprovideNow = n.triggerReprovide - n.services.SetBackendProbeTimeout(n.config.BackendProbeTimeout) var authenticated bool var fatalAuthErr error diff --git a/internal/node/service_registry.go b/internal/node/service_registry.go index 37872d1f..398583ce 100644 --- a/internal/node/service_registry.go +++ b/internal/node/service_registry.go @@ -77,44 +77,26 @@ type ServiceRegistry struct { // advertised. Optional; nil outside a running node. reprovideNow func() - // backendProbeTimeout bounds each backend probe in advertisable. Defaults - // to defaultDHTProbeTimeout; override with SetBackendProbeTimeout. + // backendProbeTimeout bounds each backend probe in advertisable. Set once + // at construction (see NewServiceRegistry); never mutated afterwards, so + // reading it needs no lock. backendProbeTimeout time.Duration } -func NewServiceRegistry(d dhtProvider) *ServiceRegistry { +// NewServiceRegistry constructs a registry bounding backend probes by +// backendProbeTimeout. A zero or negative value falls back to +// defaultDHTProbeTimeout, so callers can pass an unset +// Options.BackendProbeTimeout straight through without an explicit +// zero-check. +func NewServiceRegistry(d dhtProvider, backendProbeTimeout time.Duration) *ServiceRegistry { + if backendProbeTimeout <= 0 { + backendProbeTimeout = defaultDHTProbeTimeout + } return &ServiceRegistry{ services: map[string]Service{}, dht: d, - backendProbeTimeout: defaultDHTProbeTimeout, - } -} - -// SetBackendProbeTimeout overrides the default backend probe timeout -// (defaultDHTProbeTimeout). A zero or negative duration is a no-op, so callers can -// pass an unset Options.BackendProbeTimeout straight through without an -// explicit zero-check. -func (r *ServiceRegistry) SetBackendProbeTimeout(d time.Duration) { - if d <= 0 { - return - } - r.mu.Lock() - defer r.mu.Unlock() - r.backendProbeTimeout = d -} - -// probeTimeout returns the current backend probe timeout, falling back to -// defaultDHTProbeTimeout for a zero-initialized registry (e.g. a struct -// literal built directly in a test, bypassing NewServiceRegistry) so it -// behaves the same as a properly constructed one rather than timing out -// every probe immediately. -func (r *ServiceRegistry) probeTimeout() time.Duration { - r.mu.RLock() - defer r.mu.RUnlock() - if r.backendProbeTimeout <= 0 { - return defaultDHTProbeTimeout + backendProbeTimeout: backendProbeTimeout, } - return r.backendProbeTimeout } // Register initialises a service, advertises it on the DHT, and inserts it @@ -143,7 +125,7 @@ func (r *ServiceRegistry) Register(ctx context.Context, svc Service) error { return err } - probeErr := advertisable(ctx, svc, r.probeTimeout()) + probeErr := advertisable(ctx, svc, r.backendProbeTimeout) if probeErr != nil { logger.Warnf("[ServiceRegistry] Registered %s/%s but not advertising it: backend did not answer: %v", info.Type, info.Name, probeErr) } else { @@ -268,7 +250,7 @@ Loop: }() info := svc.Info() - if err := advertisable(ctx, svc, r.probeTimeout()); err != nil { + if err := advertisable(ctx, svc, r.backendProbeTimeout); err != nil { withheld.Add(1) // On shutdown every service fails this way, and saying so // would blame backends for the node stopping. diff --git a/internal/node/service_registry_test.go b/internal/node/service_registry_test.go index a07a3c17..f813e06b 100644 --- a/internal/node/service_registry_test.go +++ b/internal/node/service_registry_test.go @@ -71,8 +71,9 @@ func newFakeSvc(name string, st api.ServiceType) *fakeService { // newServiceRegistryForTest builds a registry against the fake DHT for tests. func newServiceRegistryForTest(d dhtProvider) *ServiceRegistry { return &ServiceRegistry{ - services: map[string]Service{}, - dht: d, + services: map[string]Service{}, + dht: d, + backendProbeTimeout: defaultDHTProbeTimeout, } } @@ -306,16 +307,16 @@ func (p *slowProbingService) Probe(ctx context.Context) error { // The bug behind #376: defaultDHTProbeTimeout was a hard-coded 2s with no way to // raise it, so a backend whose own cold-start cost alone exceeds that - // measured in practice for moderately-featured MCP server stacks - could -// never be advertised on its first registration. SetBackendProbeTimeout -// (wired from --backend-probe-timeout) is the fix: the same slow backend -// must fail to advertise under the default and succeed once given more time. +// never be advertised on its first registration. NewServiceRegistry's +// backendProbeTimeout parameter (wired from --backend-probe-timeout) is the +// fix: the same slow backend must fail to advertise under the default and +// succeed once constructed with more time. func TestServiceRegistry_BackendProbeTimeoutIsConfigurable(t *testing.T) { const probeDelay = 60 * time.Millisecond t.Run("default timeout is too short for a slow backend", func(t *testing.T) { dht := &fakeDHT{} - r := NewServiceRegistry(dht) - r.SetBackendProbeTimeout(10 * time.Millisecond) // shorter than probeDelay + r := NewServiceRegistry(dht, 10*time.Millisecond) // shorter than probeDelay svc := newSlowProbingSvc("slow", probeDelay) if err := r.Register(context.Background(), svc); err != nil { @@ -328,8 +329,7 @@ func TestServiceRegistry_BackendProbeTimeoutIsConfigurable(t *testing.T) { t.Run("raising the timeout lets the same backend advertise", func(t *testing.T) { dht := &fakeDHT{} - r := NewServiceRegistry(dht) - r.SetBackendProbeTimeout(probeDelay * 5) // comfortably longer than probeDelay + r := NewServiceRegistry(dht, probeDelay*5) // comfortably longer than probeDelay svc := newSlowProbingSvc("slow", probeDelay) if err := r.Register(context.Background(), svc); err != nil { @@ -340,19 +340,12 @@ func TestServiceRegistry_BackendProbeTimeoutIsConfigurable(t *testing.T) { } }) - t.Run("NewServiceRegistry defaults to defaultDHTProbeTimeout unchanged", func(t *testing.T) { - r := NewServiceRegistry(&fakeDHT{}) - if got := r.probeTimeout(); got != defaultDHTProbeTimeout { - t.Errorf("default probe timeout = %v, want %v (unchanged default behaviour)", got, defaultDHTProbeTimeout) - } - }) - - t.Run("SetBackendProbeTimeout ignores zero and negative durations", func(t *testing.T) { - r := NewServiceRegistry(&fakeDHT{}) - r.SetBackendProbeTimeout(0) - r.SetBackendProbeTimeout(-1 * time.Second) - if got := r.probeTimeout(); got != defaultDHTProbeTimeout { - t.Errorf("probe timeout after no-op sets = %v, want unchanged %v", got, defaultDHTProbeTimeout) + t.Run("zero or negative backendProbeTimeout falls back to defaultDHTProbeTimeout", func(t *testing.T) { + for _, d := range []time.Duration{0, -1 * time.Second} { + r := NewServiceRegistry(&fakeDHT{}, d) + if got := r.backendProbeTimeout; got != defaultDHTProbeTimeout { + t.Errorf("NewServiceRegistry(dht, %v).backendProbeTimeout = %v, want %v", d, got, defaultDHTProbeTimeout) + } } }) } diff --git a/internal/node/sidecar_auth_test.go b/internal/node/sidecar_auth_test.go index fc8b311b..0fc04e19 100644 --- a/internal/node/sidecar_auth_test.go +++ b/internal/node/sidecar_auth_test.go @@ -66,7 +66,7 @@ func TestConstantTimeEqual(t *testing.T) { func TestMetricsGatedOnTCPButNotOnTheSocket(t *testing.T) { node := &SamNode{ BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } socketPath := filepath.Join(t.TempDir(), "sam.sock") diff --git a/internal/node/sidecar_test.go b/internal/node/sidecar_test.go index e8d8986e..bbbe2990 100644 --- a/internal/node/sidecar_test.go +++ b/internal/node/sidecar_test.go @@ -70,7 +70,7 @@ func waitForSocket(t *testing.T, path string) *http.Client { func TestSidecarSocketAuthorizesWithoutToken(t *testing.T) { node := &SamNode{ BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } socketPath := filepath.Join(t.TempDir(), "sam.sock") @@ -120,7 +120,7 @@ func TestSidecarSocketAuthorizesWithoutToken(t *testing.T) { func TestSidecarSocketOnly(t *testing.T) { node := &SamNode{ BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } socketPath := filepath.Join(t.TempDir(), "sam.sock") @@ -138,7 +138,7 @@ func TestSidecarSocketOnly(t *testing.T) { } func TestStartSidecarServerRequiresAListener(t *testing.T) { - node := &SamNode{services: NewServiceRegistry(&fakeDHT{})} + node := &SamNode{services: NewServiceRegistry(&fakeDHT{}, 0)} if _, err := StartSidecarServer(node, "", "", "token", "", "", ""); err == nil { t.Fatal("expected an error when neither a TCP address nor a socket is configured") } @@ -156,7 +156,7 @@ func TestSidecarSocketFailureKeepsTCPServing(t *testing.T) { node := &SamNode{ BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } srv, err := StartSidecarServer(node, "127.0.0.1:0", socketPath, "test-token", "", "", "") if err != nil { @@ -170,7 +170,7 @@ func TestSidecarSocketFailureKeepsTCPServing(t *testing.T) { t.Errorf("BoundSocketPath = %q, want empty after a failed socket", node.BoundSocketPath) } - socketOnly := &SamNode{services: NewServiceRegistry(&fakeDHT{})} + socketOnly := &SamNode{services: NewServiceRegistry(&fakeDHT{}, 0)} if _, err := StartSidecarServer(socketOnly, "", socketPath, "", "", "", ""); err == nil { t.Error("expected an error when the socket is the only configured listener") } @@ -363,7 +363,7 @@ func TestWithAuth(t *testing.T) { func TestSidecarServerAuthEnforcement(t *testing.T) { node := &SamNode{ BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } // We use a dummy token token := "test-token" @@ -453,7 +453,7 @@ func TestSidecarServerAuthEnforcement(t *testing.T) { func TestSidecarAuthorizationFallbackScope(t *testing.T) { node := &SamNode{ BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } token := "test-token" @@ -546,7 +546,7 @@ func TestRegisterService(t *testing.T) { time.Sleep(100 * time.Millisecond) node := &SamNode{BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(d), + services: NewServiceRegistry(d, 0), DHT: d, } @@ -575,7 +575,7 @@ func TestRegisterService(t *testing.T) { func TestUnregisterService(t *testing.T) { node := &SamNode{BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } node.services.insertService(&testService{info: &api.ServiceInfo{Name: "test-service"}}) @@ -601,7 +601,7 @@ func TestHandleDiscoverService(t *testing.T) { defer func() { _ = d.Close() }() node := &SamNode{BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(d), + services: NewServiceRegistry(d, 0), DHT: d, Host: h, BoundHTTPAddr: "127.0.0.1:8080", @@ -660,7 +660,7 @@ func TestHandleDiscoverService(t *testing.T) { func TestListLocalServices(t *testing.T) { node := &SamNode{BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } service1 := &api.ServiceInfo{Type: api.ServiceType_SERVICE_TYPE_MCP, Name: "service1"} @@ -678,7 +678,7 @@ func TestListLocalServices(t *testing.T) { func TestListLocalServices_TypeFilter(t *testing.T) { node := &SamNode{BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } mcpA := &api.ServiceInfo{Type: api.ServiceType_SERVICE_TYPE_MCP, Name: "mcp-a"} mcpB := &api.ServiceInfo{Type: api.ServiceType_SERVICE_TYPE_MCP, Name: "mcp-b"} @@ -777,7 +777,7 @@ func TestServiceKeyToCID_Equivalence(t *testing.T) { func TestRegisterService_Validation(t *testing.T) { node := &SamNode{BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), } tests := []struct { @@ -858,7 +858,7 @@ func TestDiscoverService_Pagination(t *testing.T) { node := &SamNode{ BiscuitTimeout: 500 * time.Millisecond, - services: NewServiceRegistry(d), + services: NewServiceRegistry(d, 0), DHT: d, Host: h, BoundHTTPAddr: "127.0.0.1:8080", From 4e181ea81cd5106acbc9487ea1fb0b5dac827fe7 Mon Sep 17 00:00:00 2001 From: Fernando Marino Date: Thu, 10 Sep 2026 08:17:20 +0200 Subject: [PATCH 4/6] address review: default BackendProbeTimeout in Options.Default(), matching every other option --- internal/node/options.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/internal/node/options.go b/internal/node/options.go index 2075c7ee..730ecd1e 100644 --- a/internal/node/options.go +++ b/internal/node/options.go @@ -142,6 +142,9 @@ func (o *Options) Default() { if o.NodeConfig == nil { o.NodeConfig = &NodeConfigComplete{} } + if o.BackendProbeTimeout <= 0 { + o.BackendProbeTimeout = defaultDHTProbeTimeout + } } // Validate verifies that the required options are provided and valid. From 35debe29b73606781bbd15fbd10db8834cb665f9 Mon Sep 17 00:00:00 2001 From: Fernando Marino Date: Thu, 10 Sep 2026 09:38:12 +0200 Subject: [PATCH 5/6] fix rebase conflict resolution: TestDiscoverySource's labels moved to nodeConfig upstream, not Options --- internal/node/mcp_discovery_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/node/mcp_discovery_test.go b/internal/node/mcp_discovery_test.go index ef63e5a9..127304b7 100644 --- a/internal/node/mcp_discovery_test.go +++ b/internal/node/mcp_discovery_test.go @@ -71,8 +71,8 @@ func (f *fakeToolService) Tools(_ context.Context) ([]string, error) { func TestDiscoverySource(t *testing.T) { node := &SamNode{ - services: NewServiceRegistry(&fakeDHT{}, 0), - config: Options{Labels: map[string]string{"region": "EU"}}, + services: NewServiceRegistry(&fakeDHT{}, 0), + nodeConfig: &NodeConfigComplete{Labels: map[string]string{"region": "EU"}}, } ctx := context.Background() From 1d047d12ef09844178db5a570dc6a7c11243d5b5 Mon Sep 17 00:00:00 2001 From: Fernando Marino Date: Thu, 10 Sep 2026 09:41:29 +0200 Subject: [PATCH 6/6] address review: make the probe-timeout test deterministic, no real-time sleeps --- internal/node/service_registry_test.go | 49 +++++++++++++++----------- 1 file changed, 29 insertions(+), 20 deletions(-) diff --git a/internal/node/service_registry_test.go b/internal/node/service_registry_test.go index f813e06b..dfbd2a3f 100644 --- a/internal/node/service_registry_test.go +++ b/internal/node/service_registry_test.go @@ -277,31 +277,33 @@ func TestServiceRegistry_ReprovideResumesWhenBackendRecovers(t *testing.T) { } } -// slowProbingService is a backendProber whose Probe blocks until the given -// delay elapses or the context is cancelled first, whichever comes first - -// unlike probingService, it actually respects the probe deadline, which is -// what a real command-spawned backend with a slow cold-start does. +// slowProbingService is a backendProber whose Probe is deadline-inspecting +// and context-controlled rather than timer-based: it records the deadline +// it was given, and for the timeout case blocks on ctx.Done() instead of +// sleeping a real duration. This keeps the tests below deterministic and +// free of real-time dependencies - no CI flakiness from scheduling jitter, +// no slow test suite from real sleeps - while still exercising the same +// behavior a real command-spawned backend with a slow cold-start would hit. type slowProbingService struct { *fakeService - delay time.Duration + shouldTimeout bool + lastDeadline time.Time + hasDeadline bool } -func newSlowProbingSvc(name string, delay time.Duration) *slowProbingService { +func newSlowProbingSvc(name string) *slowProbingService { return &slowProbingService{ fakeService: newFakeSvc(name, api.ServiceType_SERVICE_TYPE_MCP), - delay: delay, } } func (p *slowProbingService) Probe(ctx context.Context) error { - timer := time.NewTimer(p.delay) - defer timer.Stop() - select { - case <-timer.C: - return nil - case <-ctx.Done(): + p.lastDeadline, p.hasDeadline = ctx.Deadline() + if p.shouldTimeout { + <-ctx.Done() return ctx.Err() } + return nil } // The bug behind #376: defaultDHTProbeTimeout was a hard-coded 2s with no way to @@ -312,13 +314,12 @@ func (p *slowProbingService) Probe(ctx context.Context) error { // fix: the same slow backend must fail to advertise under the default and // succeed once constructed with more time. func TestServiceRegistry_BackendProbeTimeoutIsConfigurable(t *testing.T) { - const probeDelay = 60 * time.Millisecond - t.Run("default timeout is too short for a slow backend", func(t *testing.T) { dht := &fakeDHT{} - r := NewServiceRegistry(dht, 10*time.Millisecond) // shorter than probeDelay + r := NewServiceRegistry(dht, 10*time.Millisecond) - svc := newSlowProbingSvc("slow", probeDelay) + svc := newSlowProbingSvc("slow") + svc.shouldTimeout = true if err := r.Register(context.Background(), svc); err != nil { t.Fatalf("Register: %v", err) } @@ -327,17 +328,25 @@ func TestServiceRegistry_BackendProbeTimeoutIsConfigurable(t *testing.T) { } }) - t.Run("raising the timeout lets the same backend advertise", func(t *testing.T) { + t.Run("raising the timeout applies the configured duration to the probe context", func(t *testing.T) { dht := &fakeDHT{} - r := NewServiceRegistry(dht, probeDelay*5) // comfortably longer than probeDelay + timeout := 500 * time.Millisecond + r := NewServiceRegistry(dht, timeout) - svc := newSlowProbingSvc("slow", probeDelay) + svc := newSlowProbingSvc("slow") if err := r.Register(context.Background(), svc); err != nil { t.Fatalf("Register: %v", err) } if len(dht.calls) != 2 { t.Errorf("Provide called %d times once given enough time to probe, want 2 (name + type CID)", len(dht.calls)) } + if !svc.hasDeadline { + t.Fatal("expected probe context to have a deadline") + } + remaining := time.Until(svc.lastDeadline) + if remaining > timeout || remaining < timeout-100*time.Millisecond { + t.Errorf("expected probe deadline to be close to %v, got remaining %v", timeout, remaining) + } }) t.Run("zero or negative backendProbeTimeout falls back to defaultDHTProbeTimeout", func(t *testing.T) {