Handle Client.Close on cloned clients

client/client.go

@@ -386,6 +386,11 @@ type (
 		WorkflowService() workflowservice.WorkflowServiceClient
 
 		// Close client and clean up underlying resources.
+		//
+		// If this client was created via NewClientFromExisting or this client has
+		// been used in that call, Close() on may not necessarily close the
+		// underlying connection. Only the final close of all existing clients will
+		// close the underlying connection.
 		Close()
 	}
 
@@ -474,6 +479,11 @@ func NewClient(options Options) (Client, error) {
 // this package and cannot be wrapped. Currently, this always attempts an eager
 // connection even if the existing client was created with NewLazyClient and has
 // not made any calls yet.
+//
+// Close() on the resulting client may not necessarily close the underlying
+// connection if there are any other clients using the connection. All clients
+// associated with the existing client must call Close() and only the last one
+// actually performs the connection close.
 func NewClientFromExisting(existingClient Client, options Options) (Client, error) {
 	return internal.NewClientFromExisting(existingClient, options)
 }

internal/client.go

@@ -28,6 +28,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"sync/atomic"
 	"time"
 
 	commonpb "go.temporal.io/api/common/v1"
@@ -689,12 +690,18 @@ func newClient(options ClientOptions, existing *WorkflowClient) (Client, error)
 		if client.capabilities, err = existing.loadCapabilities(); err != nil {
 			return nil, err
 		}
-	} else if !options.ConnectionOptions.disableEagerConnection {
-		if _, err := client.loadCapabilities(); err != nil {
-			client.Close()
-			return nil, err
+		client.unclosedClients = existing.unclosedClients
+	} else {
+		if !options.ConnectionOptions.disableEagerConnection {
+			if _, err := client.loadCapabilities(); err != nil {
+				client.Close()
+				return nil, err
+			}
 		}
+		var unclosedClients int32
+		client.unclosedClients = &unclosedClients
 	}
+	atomic.AddInt32(client.unclosedClients, 1)
 
 	return client, nil
 }

internal/internal_workflow_client.go

@@ -29,8 +29,10 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"reflect"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/pborman/uuid"
@@ -83,6 +85,10 @@ type (
 		excludeInternalFromRetry *uberatomic.Bool
 		capabilities             *workflowservice.GetSystemInfoResponse_Capabilities
 		capabilitiesLock         sync.RWMutex
+
+		// The pointer value is shared across multiple clients. If non-nil, only
+		// access/mutate atomically.
+		unclosedClients *int32
 	}
 
 	// namespaceClient is the client for managing namespaces.
@@ -942,11 +948,25 @@ func (wc *WorkflowClient) ensureInitialized() error {
 
 // Close client and clean up underlying resources.
 func (wc *WorkflowClient) Close() {
-	if wc.conn == nil {
-		return
+	// If there's a set of unclosed clients, we have to decrement it and then
+	// set it to a new pointer of max to prevent decrementing on repeated Close
+	// calls to this client. If the count has not reached zero, this close call is
+	// ignored.
+	if wc.unclosedClients != nil {
+		remainingUnclosedClients := atomic.AddInt32(wc.unclosedClients, -1)
+		// Set the unclosed clients to max value so we never try this again
+		var maxUnclosedClients int32 = math.MaxInt32
+		wc.unclosedClients = &maxUnclosedClients
+		// If there are any remaining, do not close
+		if remainingUnclosedClients > 0 {
+			return
+		}
 	}
-	if err := wc.conn.Close(); err != nil {
-		wc.logger.Warn("unable to close connection", tagError, err)
+
+	if wc.conn != nil {
+		if err := wc.conn.Close(); err != nil {
+			wc.logger.Warn("unable to close connection", tagError, err)
+		}
 	}
 }
 

internal/internal_workflow_client_test.go

@@ -34,10 +34,12 @@ import (
 	workflowpb "go.temporal.io/api/workflow/v1"
 	uberatomic "go.uber.org/atomic"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 
 	ilog "go.temporal.io/sdk/internal/log"
 
 	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	commonpb "go.temporal.io/api/common/v1"
 	enumspb "go.temporal.io/api/enums/v1"
@@ -1402,3 +1404,45 @@ func serializeEvents(events []*historypb.HistoryEvent) *commonpb.DataBlob {
 		Data:         blob.Data,
 	}
 }
+
+func TestClientCloseCount(t *testing.T) {
+	// Create primary client
+	server, err := startTestGRPCServer()
+	require.NoError(t, err)
+	defer server.Stop()
+	client, err := DialClient(ClientOptions{HostPort: server.addr})
+	require.NoError(t, err)
+	workflowClient := client.(*WorkflowClient)
+
+	// Confirm there is 1 unclosed client
+	require.EqualValues(t, 1, *workflowClient.unclosedClients)
+
+	// Create two more and confirm counts
+	client2, err := NewClientFromExisting(client, ClientOptions{})
+	require.NoError(t, err)
+	require.EqualValues(t, 2, *workflowClient.unclosedClients)
+	require.Same(t, workflowClient.unclosedClients, client2.(*WorkflowClient).unclosedClients)
+	client3, err := NewClientFromExisting(client, ClientOptions{})
+	require.NoError(t, err)
+	require.EqualValues(t, 3, *workflowClient.unclosedClients)
+	require.Same(t, workflowClient.unclosedClients, client3.(*WorkflowClient).unclosedClients)
+
+	// Close the third one 3 times and confirm counts and that connection not
+	// closed
+	client3.Close()
+	client3.Close()
+	client3.Close()
+	require.EqualValues(t, 2, *workflowClient.unclosedClients)
+	require.NotSame(t, workflowClient.unclosedClients, client3.(*WorkflowClient).unclosedClients)
+	require.Less(t, workflowClient.conn.GetState(), connectivity.Shutdown)
+
+	// Close the primary one and confirm not closed
+	client.Close()
+	require.EqualValues(t, 1, *client2.(*WorkflowClient).unclosedClients)
+	require.NotSame(t, workflowClient.unclosedClients, client2.(*WorkflowClient).unclosedClients)
+	require.Less(t, workflowClient.conn.GetState(), connectivity.Shutdown)
+
+	// Now close the last one (the second) and confirm it actually gets closed
+	client2.Close()
+	require.Equal(t, connectivity.Shutdown, workflowClient.conn.GetState())
+}