feat(internal): implement Bigtable specific channel pool optimizations

bigtable/bigtable.go

@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"hash/crc32"
 	"io"
+	"log"
 	"net/url"
 	"os"
 	"strconv"
@@ -80,6 +81,7 @@ type Client struct {
 	retryOption             gax.CallOption
 	executeQueryRetryOption gax.CallOption
 	enableDirectAccess      bool
+	logger                  *log.Logger
 }
 
 // ClientConfig has configurations for the client.
@@ -94,6 +96,10 @@ type ClientConfig struct {
 	//
 	// TODO: support user provided meter provider
 	MetricsProvider MetricsProvider
+
+	// Logger is the logger to use for this client. If it is nil, all logging
+	// will be directed to the standard logger.
+	Logger *log.Logger
 }
 
 // MetricsProvider is a wrapper for built in metrics meter provider
@@ -175,9 +181,29 @@ func NewClientWithConfig(ctx context.Context, project, instance string, config C
 	var connPoolErr error
 	enableBigtableConnPool := btopt.EnableBigtableConnectionPool()
 	if enableBigtableConnPool {
-		connPool, connPoolErr = btransport.NewBigtableChannelPool(defaultBigtableConnPoolSize, btopt.BigtableLoadBalancingStrategy(), func() (*grpc.ClientConn, error) {
-			return gtransport.Dial(ctx, o...)
-		})
+		// Return *BigtableConn.
+		dialFunc := func() (*btransport.BigtableConn, error) {
+			// First, establish the underlying gRPC connection.
+			grpcConn, err := gtransport.Dial(ctx, o...)
+			if err != nil {
+				return nil, err
+			}
+			// Then, wrap the grpcConn in a BigtableConn, providing instance and app profile.
+			return btransport.NewBigtableConn(grpcConn, instance, config.AppProfile), nil
+		}
+
+		// Initialize the BigtableChannelPool with the updated dialFunc.
+		// btransport is assumed to be the package where BigtableChannelPool resides.
+		connPool, connPoolErr = btransport.NewBigtableChannelPool(
+			ctx,
+			defaultBigtableConnPoolSize,
+			btopt.BigtableLoadBalancingStrategy(),
+			dialFunc,
+			config.Logger,
+			nil,
+			btransport.WithHealthCheckConfig(btransport.DefaultHealthCheckConfig()),
+			btransport.WithDynamicChannelPool(btransport.DynamicChannelPoolConfig{}),
+		)
 	} else {
 		// use to regular ConnPool
 		connPool, connPoolErr = gtransport.DialPool(ctx, o...)
@@ -198,6 +224,7 @@ func NewClientWithConfig(ctx context.Context, project, instance string, config C
 		retryOption:             retryOption,
 		executeQueryRetryOption: executeQueryRetryOption,
 		enableDirectAccess:      enableDirectAccess,
+		logger:                  config.Logger,
 	}, nil
 }
 

bigtable/internal/option/option.go

@@ -20,6 +20,7 @@ package option
 import (
 	"context"
 	"fmt"
+	"log"
 	"os"
 	"strconv"
 	"strings"
@@ -201,3 +202,27 @@ func EnableBigtableConnectionPool() bool {
 	}
 	return enableBigtableConnPool
 }
+
+// Logf logs the given message to the given logger, or the standard logger if
+// the given logger is nil.
+func logf(logger *log.Logger, format string, v ...interface{}) {
+	if logger == nil {
+		log.Printf(format, v...)
+	} else {
+		logger.Printf(format, v...)
+	}
+}
+
+var debug = os.Getenv("CBT_ENABLE_DEBUG") == "true"
+
+// Debugf logs the given message *only if* the global Debug flag is true.
+// It reuses Logf to handle the nil logger logic and prepends "DEBUG: "
+// to the message.
+func Debugf(logger *log.Logger, format string, v ...interface{}) {
+	// Only log if the Debug flag is set
+	if debug {
+		// Prepend "DEBUG: " to the format string
+		debugFormat := "DEBUG: " + format
+		logf(logger, debugFormat, v...)
+	}
+}