diff --git a/internal/dispatch/combined/combined.go b/internal/dispatch/combined/combined.go index e4f10e6e80..71acc36317 100644 --- a/internal/dispatch/combined/combined.go +++ b/internal/dispatch/combined/combined.go @@ -292,6 +292,7 @@ func NewDispatcher(options ...Option) (dispatch.Dispatcher, error) { secondaryClients[name] = remote.SecondaryDispatch{ Name: name, Client: v1.NewDispatchServiceClient(secondaryConn), + Conn: secondaryConn, MaximumPrimaryHedgingDelay: maximumHedgingDelay, } } diff --git a/internal/dispatch/remote/cluster.go b/internal/dispatch/remote/cluster.go index eefddba17a..6ad32daa7e 100644 --- a/internal/dispatch/remote/cluster.go +++ b/internal/dispatch/remote/cluster.go @@ -111,6 +111,10 @@ type SecondaryDispatch struct { // Client is the client to use for dispatching to the secondary. Client ClusterClient + // Conn is the underlying gRPC connection for the secondary dispatcher, used for closing. + // May be nil if the client is a mock or test double. + Conn *grpc.ClientConn + // MaximumHedgingDelay is the maximum delay that the primary will wait before dispatching // when this secondary is active. MaximumPrimaryHedgingDelay time.Duration @@ -914,10 +918,20 @@ func (cr *clusterDispatcher) DispatchQueryPlan(req *v1.DispatchQueryPlanRequest, } func (cr *clusterDispatcher) Close() error { + var firstErr error + for _, secondary := range cr.secondaryDispatch { + if secondary.Conn != nil { + if err := secondary.Conn.Close(); err != nil && firstErr == nil { + firstErr = err + } + } + } if cr.conn != nil { - return cr.conn.Close() + if err := cr.conn.Close(); err != nil && firstErr == nil { + firstErr = err + } } - return nil + return firstErr } // ReadyState returns whether the underlying dispatch connection is available