feat(bigtable): add ip protocol for bigtable connection (#13520)

sushanb · web-flow · commit 6f86983e4b30 · 2025-12-29T18:27:49.000-05:00
we will need this for populating direct_access/compatible as it has an
`ip_preference` metric field.
diff --git a/bigtable/internal/transport/connpool.go b/bigtable/internal/transport/connpool.go
@@ -22,6 +22,7 @@ import (
 	"log"
 	"math"
 	"math/rand"
+	"net"
 	"net/url"
 	"slices"
 	"sort"
@@ -47,6 +48,30 @@ var maxDrainingTimeout = 30 * time.Minute
 
 const requestParamsHeader = "x-goog-request-params"
 
+// ipProtocol represents the type of IP protocol used.
+type ipProtocol int32
+
+const (
+	// unknown represents an unknown or undetermined IP protocol.
+	unknown ipProtocol = iota - 1
+	// ipv6 represents the IPv4 protocol.
+	ipv4
+	// ipv6 represents the IPv6 protocol.
+	ipv6
+)
+
+// AddressType returns the string representation of the IPProtocol.
+func (ip ipProtocol) addressType() string {
+	switch ip {
+	case ipv4:
+		return "ipv4"
+	case ipv6:
+		return "ipv6"
+	default:
+		return "unknown"
+	}
+}
+
 // BigtableChannelPoolOption options for configurable
 type BigtableChannelPoolOption func(*BigtableChannelPool)
 
@@ -111,6 +136,13 @@ type BigtableConn struct {
 	*grpc.ClientConn
 	isALTSConn atomic.Bool
 	createdAt  atomic.Int64
+	// remoteAddrType stores the  type: -1 (unknown/nil), 0 (ipv4), 1 (ipv6)
+	remoteAddrType atomic.Int32
+}
+
+// ipProtocol returns the IP protocol as a string: "ipv4", "ipv6", or "unknown".
+func (bc *BigtableConn) ipProtocol() string {
+	return ipProtocol(bc.remoteAddrType.Load()).addressType()
 }
 
 // Prime sends a PingAndWarm request to warm up the connection.
@@ -137,6 +169,19 @@ func (bc *BigtableConn) Prime(ctx context.Context, fullInstanceName, appProfileI
 		return err
 	}
 
+	// ip protocol will be -1 if it addr is nil/default, 0 is ipv4 and 1 if ipv6.
+	if p.Addr != nil {
+		if tcpAddr, ok := p.Addr.(*net.TCPAddr); ok {
+			if tcpAddr.IP != nil {
+				if tcpAddr.IP.To4() != nil {
+					bc.remoteAddrType.Store(int32(ipv4))
+				} else {
+					bc.remoteAddrType.Store(int32(ipv6))
+				}
+			}
+		}
+	}
+
 	if p.AuthInfo != nil {
 		if _, ok := p.AuthInfo.(alts.AuthInfo); ok {
 			bc.isALTSConn.Store(true)
@@ -173,8 +218,8 @@ func NewBigtableConn(conn *grpc.ClientConn) *BigtableConn {
 		ClientConn: conn,
 	}
 	bc.createdAt.Store(time.Now().UnixMilli())
+	bc.remoteAddrType.Store(int32(unknown))
 	return bc
-
 }
 
 // createdAt returns the creation time of the connection in int64. milliseconds since epoch
diff --git a/bigtable/internal/transport/connpool_test.go b/bigtable/internal/transport/connpool_test.go
@@ -63,6 +63,33 @@ func entryIndex(s []*connEntry, e *connEntry) int {
 	return -1
 }
 
+func TestBigtableConnIpProtocol(t *testing.T) {
+	bc := NewBigtableConn(nil)
+	if got := bc.ipProtocol(); got != "unknown" {
+		t.Errorf("NewBigtableConn default ipProtocol() got %q, want %q", got, unknown)
+	}
+
+	tests := []struct {
+		name     string
+		addrType ipProtocol
+		want     string
+	}{
+		{name: "IPv4", addrType: ipv4, want: "ipv4"},
+		{name: "IPv6", addrType: ipv6, want: "ipv6"},
+		{name: "Unknown", addrType: unknown, want: "unknown"},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			conn := &BigtableConn{}
+			conn.remoteAddrType.Store(int32(tc.addrType))
+			if got := conn.ipProtocol(); got != tc.want {
+				t.Errorf("ipProtocol() with remoteAddrType %d got %q, want %q", tc.addrType, got, tc.want)
+			}
+		})
+	}
+}
+
 func TestSelectRoundRobin(t *testing.T) {
 	pool := &BigtableChannelPool{rrIndex: 0}
 
