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/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..127304b7 100644 --- a/internal/node/mcp_discovery_test.go +++ b/internal/node/mcp_discovery_test.go @@ -71,7 +71,7 @@ func (f *fakeToolService) Tools(_ context.Context) ([]string, error) { func TestDiscoverySource(t *testing.T) { node := &SamNode{ - services: NewServiceRegistry(&fakeDHT{}), + services: NewServiceRegistry(&fakeDHT{}, 0), nodeConfig: &NodeConfigComplete{Labels: map[string]string{"region": "EU"}}, } ctx := context.Background() diff --git a/internal/node/node.go b/internal/node/node.go index 685bc613..f402e053 100644 --- a/internal/node/node.go +++ b/internal/node/node.go @@ -518,7 +518,7 @@ 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 var authenticated bool diff --git a/internal/node/options.go b/internal/node/options.go index 9f8a0367..730ecd1e 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. @@ -134,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. diff --git a/internal/node/service_registry.go b/internal/node/service_registry.go index ffcc2724..398583ce 100644 --- a/internal/node/service_registry.go +++ b/internal/node/service_registry.go @@ -38,8 +38,15 @@ type backendProber interface { Probe(ctx context.Context) error } -// dhtProbeTimeout bounds one backend probe before advertising. -const dhtProbeTimeout = 2 * time.Second +// 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 +// 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 defaultDHTProbeTimeout = 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,12 +76,26 @@ 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. 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, + services: map[string]Service{}, + dht: d, + backendProbeTimeout: backendProbeTimeout, } } @@ -104,7 +125,7 @@ func (r *ServiceRegistry) Register(ctx context.Context, svc Service) error { return err } - probeErr := advertisable(ctx, svc) + 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 { @@ -229,7 +250,7 @@ Loop: }() info := svc.Info() - if err := advertisable(ctx, svc); 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 b0d2f5e5..dfbd2a3f 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" @@ -70,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, } } @@ -274,3 +276,85 @@ 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 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 + shouldTimeout bool + lastDeadline time.Time + hasDeadline bool +} + +func newSlowProbingSvc(name string) *slowProbingService { + return &slowProbingService{ + fakeService: newFakeSvc(name, api.ServiceType_SERVICE_TYPE_MCP), + } +} + +func (p *slowProbingService) Probe(ctx context.Context) error { + 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 +// 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. 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) { + t.Run("default timeout is too short for a slow backend", func(t *testing.T) { + dht := &fakeDHT{} + r := NewServiceRegistry(dht, 10*time.Millisecond) + + svc := newSlowProbingSvc("slow") + svc.shouldTimeout = true + 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 applies the configured duration to the probe context", func(t *testing.T) { + dht := &fakeDHT{} + timeout := 500 * time.Millisecond + r := NewServiceRegistry(dht, timeout) + + 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) { + 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",