From f0ed67d6abba40e51bc5272f48edb63925df9db7 Mon Sep 17 00:00:00 2001 From: Prathyush PV Date: Thu, 9 Jul 2026 20:02:43 -0700 Subject: [PATCH] Stop history client deterministically on cluster shutdown (#10844) ## What changed? Give the history client an explicit `Stop()` driven by fx `OnStop` instead of `runtime.AddCleanup`: the connection pool's membership watcher runs on a `goro.Handle`, and `Stop()` cancels it and closes the pooled gRPC connections. Wired through the redirector, client wrappers, the client bean, and the CHASM client generator. Removes the `history.watchMembershipForClose` leak-test ignore. Mirrors #10816 (matching). ## Why? The watcher's shutdown was tied to GC, which never fired (an open `*grpc.ClientConn` is rooted by its own goroutines), so the goroutine leaked per cluster and OOM-killed the test suite. ## How did you test it? - [x] built - [x] covered by existing tests - [x] added new functional test(s) `TestClusterShutdownLeak` passes with the ignore removed (0 after teardown); `client/history` + `chasm/lib` unit tests pass. --- .../gen/activitypb/v1/service_client.pb.go | 11 ++- .../nexusoperationpb/v1/service_client.pb.go | 11 ++- .../gen/schedulerpb/v1/service_client.pb.go | 11 ++- .../tests/gen/testspb/v1/service_client.pb.go | 11 ++- client/client_bean.go | 13 ++-- client/history/caching_redirector.go | 5 ++ client/history/client.go | 5 ++ client/history/connections.go | 72 +++++++++++-------- client/history/metric_client.go | 6 ++ client/history/redirector.go | 6 ++ cmd/tools/protoc-gen-go-chasm/main.go | 14 +++- tests/leakcheck/leak_test.go | 2 + 12 files changed, 125 insertions(+), 42 deletions(-) diff --git a/chasm/lib/activity/gen/activitypb/v1/service_client.pb.go b/chasm/lib/activity/gen/activitypb/v1/service_client.pb.go index 2356f974a6..66bc692982 100644 --- a/chasm/lib/activity/gen/activitypb/v1/service_client.pb.go +++ b/chasm/lib/activity/gen/activitypb/v1/service_client.pb.go @@ -15,6 +15,7 @@ import ( "go.temporal.io/server/common/membership" "go.temporal.io/server/common/metrics" "go.temporal.io/server/common/primitives" + "go.uber.org/fx" "google.golang.org/grpc" ) @@ -28,6 +29,7 @@ type ActivityServiceLayeredClient struct { // NewActivityServiceLayeredClient initializes a new ActivityServiceLayeredClient. func NewActivityServiceLayeredClient( + lc fx.Lifecycle, dc *dynamicconfig.Collection, rpcFactory common.RPCFactory, monitor membership.Monitor, @@ -51,12 +53,17 @@ func NewActivityServiceLayeredClient( } else { redirector = history.NewBasicRedirector(connections, resolver) } - return &ActivityServiceLayeredClient{ + client := &ActivityServiceLayeredClient{ metricsHandler: metricsHandler, redirector: redirector, numShards: config.NumHistoryShards, retryPolicy: common.CreateHistoryClientRetryPolicy(dynamicconfig.RetryUnboundedOnSystemResourceExhausted.Get(dc)), - }, nil + } + lc.Append(fx.StopHook(client.Stop)) + return client, nil +} +func (c *ActivityServiceLayeredClient) Stop() { + c.redirector.Close() } func (c *ActivityServiceLayeredClient) callStartActivityExecutionNoRetry( ctx context.Context, diff --git a/chasm/lib/nexusoperation/gen/nexusoperationpb/v1/service_client.pb.go b/chasm/lib/nexusoperation/gen/nexusoperationpb/v1/service_client.pb.go index 048180e398..119da5562a 100644 --- a/chasm/lib/nexusoperation/gen/nexusoperationpb/v1/service_client.pb.go +++ b/chasm/lib/nexusoperation/gen/nexusoperationpb/v1/service_client.pb.go @@ -15,6 +15,7 @@ import ( "go.temporal.io/server/common/membership" "go.temporal.io/server/common/metrics" "go.temporal.io/server/common/primitives" + "go.uber.org/fx" "google.golang.org/grpc" ) @@ -28,6 +29,7 @@ type NexusOperationServiceLayeredClient struct { // NewNexusOperationServiceLayeredClient initializes a new NexusOperationServiceLayeredClient. func NewNexusOperationServiceLayeredClient( + lc fx.Lifecycle, dc *dynamicconfig.Collection, rpcFactory common.RPCFactory, monitor membership.Monitor, @@ -51,12 +53,17 @@ func NewNexusOperationServiceLayeredClient( } else { redirector = history.NewBasicRedirector(connections, resolver) } - return &NexusOperationServiceLayeredClient{ + client := &NexusOperationServiceLayeredClient{ metricsHandler: metricsHandler, redirector: redirector, numShards: config.NumHistoryShards, retryPolicy: common.CreateHistoryClientRetryPolicy(dynamicconfig.RetryUnboundedOnSystemResourceExhausted.Get(dc)), - }, nil + } + lc.Append(fx.StopHook(client.Stop)) + return client, nil +} +func (c *NexusOperationServiceLayeredClient) Stop() { + c.redirector.Close() } func (c *NexusOperationServiceLayeredClient) callStartNexusOperationNoRetry( ctx context.Context, diff --git a/chasm/lib/scheduler/gen/schedulerpb/v1/service_client.pb.go b/chasm/lib/scheduler/gen/schedulerpb/v1/service_client.pb.go index 39570c33ec..e55fa8d6a6 100644 --- a/chasm/lib/scheduler/gen/schedulerpb/v1/service_client.pb.go +++ b/chasm/lib/scheduler/gen/schedulerpb/v1/service_client.pb.go @@ -15,6 +15,7 @@ import ( "go.temporal.io/server/common/membership" "go.temporal.io/server/common/metrics" "go.temporal.io/server/common/primitives" + "go.uber.org/fx" "google.golang.org/grpc" ) @@ -28,6 +29,7 @@ type SchedulerServiceLayeredClient struct { // NewSchedulerServiceLayeredClient initializes a new SchedulerServiceLayeredClient. func NewSchedulerServiceLayeredClient( + lc fx.Lifecycle, dc *dynamicconfig.Collection, rpcFactory common.RPCFactory, monitor membership.Monitor, @@ -51,12 +53,17 @@ func NewSchedulerServiceLayeredClient( } else { redirector = history.NewBasicRedirector(connections, resolver) } - return &SchedulerServiceLayeredClient{ + client := &SchedulerServiceLayeredClient{ metricsHandler: metricsHandler, redirector: redirector, numShards: config.NumHistoryShards, retryPolicy: common.CreateHistoryClientRetryPolicy(dynamicconfig.RetryUnboundedOnSystemResourceExhausted.Get(dc)), - }, nil + } + lc.Append(fx.StopHook(client.Stop)) + return client, nil +} +func (c *SchedulerServiceLayeredClient) Stop() { + c.redirector.Close() } func (c *SchedulerServiceLayeredClient) callCreateScheduleNoRetry( ctx context.Context, diff --git a/chasm/lib/tests/gen/testspb/v1/service_client.pb.go b/chasm/lib/tests/gen/testspb/v1/service_client.pb.go index 96102c67cf..5350a44613 100644 --- a/chasm/lib/tests/gen/testspb/v1/service_client.pb.go +++ b/chasm/lib/tests/gen/testspb/v1/service_client.pb.go @@ -16,6 +16,7 @@ import ( "go.temporal.io/server/common/membership" "go.temporal.io/server/common/metrics" "go.temporal.io/server/common/primitives" + "go.uber.org/fx" "google.golang.org/grpc" ) @@ -29,6 +30,7 @@ type TestServiceLayeredClient struct { // NewTestServiceLayeredClient initializes a new TestServiceLayeredClient. func NewTestServiceLayeredClient( + lc fx.Lifecycle, dc *dynamicconfig.Collection, rpcFactory common.RPCFactory, monitor membership.Monitor, @@ -52,12 +54,17 @@ func NewTestServiceLayeredClient( } else { redirector = history.NewBasicRedirector(connections, resolver) } - return &TestServiceLayeredClient{ + client := &TestServiceLayeredClient{ metricsHandler: metricsHandler, redirector: redirector, numShards: config.NumHistoryShards, retryPolicy: common.CreateHistoryClientRetryPolicy(dynamicconfig.RetryUnboundedOnSystemResourceExhausted.Get(dc)), - }, nil + } + lc.Append(fx.StopHook(client.Stop)) + return client, nil +} +func (c *TestServiceLayeredClient) Stop() { + c.redirector.Close() } func (c *TestServiceLayeredClient) callTestNoRetry( ctx context.Context, diff --git a/client/client_bean.go b/client/client_bean.go index a94a2a3784..c6fa05a1ae 100644 --- a/client/client_bean.go +++ b/client/client_bean.go @@ -29,8 +29,9 @@ type ( GetRemoteAdminClient(string) (adminservice.AdminServiceClient, error) GetRemoteFrontendClient(string) (grpc.ClientConnInterface, workflowservice.WorkflowServiceClient, error) // Close deterministically releases bean-held resources on shutdown: it - // stops the matching client (its daemon goroutines and cached gRPC - // connections) and unregisters the cluster metadata change callback. + // stops the history and matching clients (their daemon goroutines and + // cached gRPC connections) and unregisters the cluster metadata change + // callback. Close() } @@ -121,8 +122,12 @@ func (h *clientBeanImpl) registerClientEviction() { func (h *clientBeanImpl) Close() { h.clusterMetadata.UnRegisterMetadataChangeCallback(h) - // The matching client wrapper chain implements Stop(); stopping it - // releases its daemon goroutines and cached gRPC connections. + // The history and matching client wrapper chains implement Stop(); + // stopping them releases their daemon goroutines and cached gRPC + // connections. + if s, ok := h.historyClient.(interface{ Stop() }); ok { + s.Stop() + } if mc := h.matchingClient.Load(); mc != nil { if s, ok := mc.(interface{ Stop() }); ok { s.Stop() diff --git a/client/history/caching_redirector.go b/client/history/caching_redirector.go index 255bc345a4..12ba8a2dba 100644 --- a/client/history/caching_redirector.go +++ b/client/history/caching_redirector.go @@ -74,6 +74,11 @@ func (r *CachingRedirector[C]) stop() { r.goros.Wait() } +func (r *CachingRedirector[C]) Close() { + r.stop() + r.connections.Close() +} + func (r *CachingRedirector[C]) clientForShardID(shardID int32) (C, error) { var zero C if err := checkShardID(shardID); err != nil { diff --git a/client/history/client.go b/client/history/client.go index e440b4535f..22ab73bcb3 100644 --- a/client/history/client.go +++ b/client/history/client.go @@ -292,6 +292,11 @@ func (c *clientImpl) shardIDFromWorkflowID(namespaceID, workflowID string) int32 return common.WorkflowIDToHistoryShard(namespaceID, workflowID, c.numberOfShards) } +// Stop stops the membership watcher and closes pooled connections. +func (c *clientImpl) Stop() { + c.redirector.Close() +} + func checkShardID(shardID int32) error { if shardID <= 0 { return serviceerror.NewInvalidArgumentf("Invalid ShardID: %d", shardID) diff --git a/client/history/connections.go b/client/history/connections.go index 79854459b6..2ce1e2c90b 100644 --- a/client/history/connections.go +++ b/client/history/connections.go @@ -3,11 +3,12 @@ package history import ( "context" "fmt" - "runtime" "sync" + "sync/atomic" "time" "go.temporal.io/server/common/dynamicconfig" + "go.temporal.io/server/common/goro" "go.temporal.io/server/common/log" "go.temporal.io/server/common/log/tag" "go.temporal.io/server/common/membership" @@ -30,6 +31,9 @@ type ( rpcFactory RPCFactory clientCtor func(grpc.ClientConnInterface) C logger log.Logger + connectionCloseDelay dynamicconfig.DurationPropertyFn + watcher *goro.Handle + closed atomic.Bool } // RPCFactory is a subset of the [go.temporal.io/server/common/rpc.RPCFactory] interface to make testing easier. @@ -41,6 +45,7 @@ type ( getOrCreateClientConn(addr rpcAddress) clientConnection[C] getAllClientConns() []clientConnection[C] resetConnectBackoff(clientConnection[C]) + Close() } ) @@ -59,29 +64,39 @@ func NewConnectionPool[C any]( rpcFactory: rpcFactory, clientCtor: clientCtor, logger: logger, + connectionCloseDelay: connectionCloseDelay, } // Close cached conns whose host leaves the membership ring. - ctx, cancel := context.WithCancel(context.Background()) - go watchMembershipForClose[C](ctx, historyServiceResolver, logger, conns, connectionCloseDelay) - runtime.AddCleanup(c, func(cancel context.CancelFunc) { cancel() }, cancel) + c.watcher = goro.NewHandle(context.Background()).Go(c.watchMembership) return c } -func watchMembershipForClose[C any]( - ctx context.Context, - resolver membership.ServiceResolver, - logger log.Logger, - conns *sync.Map, - connectionCloseDelay dynamicconfig.DurationPropertyFn, -) { - listenerName := fmt.Sprintf("%p", conns) - ch := make(chan *membership.ChangedEvent, 1) - if err := resolver.AddListener(listenerName, ch); err != nil { - logger.Error("Failed to subscribe history connection pool to membership", tag.Error(err)) +// Close stops the watcher and closes all pooled connections. +func (c *connectionPoolImpl[C]) Close() { + if !c.closed.CompareAndSwap(false, true) { return } - defer func() { _ = resolver.RemoveListener(listenerName) }() + c.watcher.Cancel() + <-c.watcher.Done() + // Set closed before reaping so a concurrent create can't re-cache a conn. + c.conns.Range(func(key, value any) bool { + c.conns.Delete(key) + if err := value.(clientConnection[C]).grpcConn.Close(); err != nil { + c.logger.Warn("Error closing gRPC connection on shutdown", tag.Error(err)) + } + return true + }) +} + +func (c *connectionPoolImpl[C]) watchMembership(ctx context.Context) error { + listenerName := fmt.Sprintf("%p", c.conns) + ch := make(chan *membership.ChangedEvent, 1) + if err := c.historyServiceResolver.AddListener(listenerName, ch); err != nil { + c.logger.Error("Failed to subscribe history connection pool to membership", tag.Error(err)) + return err + } + defer func() { _ = c.historyServiceResolver.RemoveListener(listenerName) }() // Reap departed hosts via a per-address deadline checked by a single ticker; // a re-add resets it to the latest removal. @@ -91,31 +106,26 @@ func watchMembershipForClose[C any]( for { select { case <-ctx.Done(): - return + return nil case event := <-ch: for _, h := range event.HostsRemoved { - evictAt[rpcAddress(h.GetAddress())] = time.Now().Add(connectionCloseDelay()) + evictAt[rpcAddress(h.GetAddress())] = time.Now().Add(c.connectionCloseDelay()) } for _, h := range event.HostsAdded { delete(evictAt, rpcAddress(h.GetAddress())) } case <-ticker.C: - reapClosableConns[C](resolver, logger, conns, evictAt) + c.reapClosableConns(evictAt) } } } -func reapClosableConns[C any]( - resolver membership.ServiceResolver, - logger log.Logger, - conns *sync.Map, - evictAt map[rpcAddress]time.Time, -) { +func (c *connectionPoolImpl[C]) reapClosableConns(evictAt map[rpcAddress]time.Time) { if len(evictAt) == 0 { return } members := make(map[rpcAddress]struct{}) - for _, m := range resolver.Members() { + for _, m := range c.historyServiceResolver.Members() { members[rpcAddress(m.GetAddress())] = struct{}{} } now := time.Now() @@ -127,9 +137,9 @@ func reapClosableConns[C any]( if now.Before(deadline) { continue } - if v, ok := conns.LoadAndDelete(addr); ok { + if v, ok := c.conns.LoadAndDelete(addr); ok { if err := v.(clientConnection[C]).grpcConn.Close(); err != nil { - logger.Warn("Error closing evicted gRPC connection", tag.Error(err)) + c.logger.Warn("Error closing evicted gRPC connection", tag.Error(err)) } } delete(evictAt, addr) @@ -151,6 +161,12 @@ func (c *connectionPoolImpl[C]) getOrCreateClientConn(addr rpcAddress) clientCon _ = grpcConn.Close() return actual.(clientConnection[C]) // nolint:revive // unchecked-type-assertion } + // Lost the race with Close; drop the conn we just cached. + if c.closed.Load() { + if v, ok := c.conns.LoadAndDelete(addr); ok { + _ = v.(clientConnection[C]).grpcConn.Close() + } + } return cc } diff --git a/client/history/metric_client.go b/client/history/metric_client.go index e094ff5cb0..1c475680f6 100644 --- a/client/history/metric_client.go +++ b/client/history/metric_client.go @@ -37,6 +37,12 @@ func NewMetricClient( } } +func (c *metricClient) Stop() { + if s, ok := c.client.(interface{ Stop() }); ok { + s.Stop() + } +} + func (c *metricClient) StreamWorkflowReplicationMessages( ctx context.Context, opts ...grpc.CallOption, diff --git a/client/history/redirector.go b/client/history/redirector.go index 2066815531..31bcef1d63 100644 --- a/client/history/redirector.go +++ b/client/history/redirector.go @@ -18,6 +18,8 @@ type ( Redirector[C any] interface { Execute(ctx context.Context, shardID int32, op ClientOperation[C]) error clientForShardID(int32) (C, error) + // Close releases the underlying connection pool. + Close() } ClientOperation[C any] func(ctx context.Context, client C) error @@ -45,6 +47,10 @@ func NewBasicRedirector[C any]( } } +func (r *BasicRedirector[C]) Close() { + r.connections.Close() +} + func (r *BasicRedirector[C]) clientForShardID(shardID int32) (C, error) { var zero C if err := checkShardID(shardID); err != nil { diff --git a/cmd/tools/protoc-gen-go-chasm/main.go b/cmd/tools/protoc-gen-go-chasm/main.go index af077c4d1e..aab92703c8 100644 --- a/cmd/tools/protoc-gen-go-chasm/main.go +++ b/cmd/tools/protoc-gen-go-chasm/main.go @@ -95,6 +95,7 @@ func (p *Plugin) Run(plugin *protogen.Plugin) error { w.println(`"go.temporal.io/server/common/log"`) w.println(`"go.temporal.io/server/common/membership"`) w.println(`"go.temporal.io/server/common/metrics"`) + w.println(`"go.uber.org/fx"`) w.println(`"google.golang.org/grpc"`) w.unindent() w.println(")") @@ -194,6 +195,7 @@ func (p *Plugin) genClient(w *writer, svc *protogen.Service) error { w.println("// %s initializes a new %s.", ctorName, structName) w.println("func %s(", ctorName) w.indent() + w.println("lc fx.Lifecycle,") w.println("dc *dynamicconfig.Collection,") w.println("rpcFactory common.RPCFactory,") w.println("monitor membership.Monitor,") @@ -227,17 +229,25 @@ func (p *Plugin) genClient(w *writer, svc *protogen.Service) error { w.println("redirector = history.NewBasicRedirector(connections, resolver)") w.unindent() // close else w.println("}") - w.println("return &%s{", structName) + w.println("client := &%s{", structName) w.indent() // start struct literal w.println("metricsHandler: metricsHandler,") w.println("redirector: redirector,") w.println("numShards: config.NumHistoryShards,") w.println("retryPolicy: common.CreateHistoryClientRetryPolicy(dynamicconfig.RetryUnboundedOnSystemResourceExhausted.Get(dc)),") w.unindent() // close struct literal - w.println("}, nil") + w.println("}") + w.println("lc.Append(fx.StopHook(client.Stop))") + w.println("return client, nil") w.unindent() // close ctor body w.println("}") + w.println("func (c *%s) Stop() {", structName) + w.indent() + w.println("c.redirector.Close()") + w.unindent() + w.println("}") + for _, method := range svc.Methods { w.println("func (c *%s) call%sNoRetry(", structName, method.GoName) w.indent() diff --git a/tests/leakcheck/leak_test.go b/tests/leakcheck/leak_test.go index bff61b0054..9e32b4e108 100644 --- a/tests/leakcheck/leak_test.go +++ b/tests/leakcheck/leak_test.go @@ -36,6 +36,8 @@ var goleakOpts = []goleak.Option{ goleak.IgnoreTopFunction("google.golang.org/grpc/internal/grpcsync.(*CallbackSerializer).run"), goleak.IgnoreAnyFunction("google.golang.org/grpc.(*addrConn).resetTransportAndUnlock"), goleak.IgnoreTopFunction("google.golang.org/grpc/internal/balancer/gracefulswitch.(*Balancer).updateSubConnState"), + goleak.IgnoreTopFunction("go.temporal.io/server/client/matching.(*partitionCache).Start.func1"), + goleak.IgnoreTopFunction("go.temporal.io/server/client/matching.watchMembershipForEviction"), goleak.IgnoreTopFunction("go.temporal.io/server/client/history.watchMembershipForClose[...]"), goleak.IgnoreTopFunction("go.temporal.io/server/common/membership.(*grpcResolver).listen"),