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.

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,

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,

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,

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,

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()

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 {

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)

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
 }
 

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,

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 {

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()

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"),