* Added balancers.WithEndpoint() context modifier for define per request the YDB endpoint by NodeID

Author: asmyasnikov

CHANGELOG.md

@@ -1,5 +1,6 @@
 * Removed check the node is available for query and table service sessions 
 * Refactored the `balancers.PreferLocations()` function - it is a clean/pure function
+* Added `balancers.WithNodeID()` context modifier for define per request the YDB endpoint by NodeID
 
 ## v3.74.2
 * Added description to scan errors with use query service client scanner

balancers/context.go

@@ -0,0 +1,11 @@
+package balancers
+
+import (
+	"context"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
+)
+
+func WithNodeID(ctx context.Context, nodeID uint32) context.Context {
+	return endpoint.WithNodeID(ctx, nodeID)
+}

internal/balancer/connections_state.go

@@ -5,6 +5,7 @@ import (
 
 	balancerConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xrand"
 )
 
@@ -73,8 +74,8 @@ func (s *connectionsState) GetConnection(ctx context.Context) (_ conn.Conn, fail
 }
 
 func (s *connectionsState) preferConnection(ctx context.Context) conn.Conn {
-	if e, hasPreferEndpoint := ContextEndpoint(ctx); hasPreferEndpoint {
-		c := s.connByNodeID[e.NodeID()]
+	if nodeID, hasPreferEndpoint := endpoint.ContextNodeID(ctx); hasPreferEndpoint {
+		c := s.connByNodeID[nodeID]
 		if c != nil && isOkConnection(c, true) {
 			return c
 		}

internal/balancer/connections_state_test.go

@@ -9,6 +9,7 @@ import (
 
 	balancerConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/mock"
 )
 
@@ -447,7 +448,7 @@ func TestConnection(t *testing.T) {
 			&mock.Conn{AddrField: "1", State: conn.Online, NodeIDField: 1},
 			&mock.Conn{AddrField: "2", State: conn.Online, NodeIDField: 2},
 		}, nil, balancerConfig.Info{}, false)
-		c, failed := s.GetConnection(WithEndpoint(context.Background(), &mock.Endpoint{AddrField: "2", NodeIDField: 2}))
+		c, failed := s.GetConnection(endpoint.WithNodeID(context.Background(), 2))
 		require.Equal(t, &mock.Conn{AddrField: "2", State: conn.Online, NodeIDField: 2}, c)
 		require.Equal(t, 0, failed)
 	})
@@ -456,7 +457,7 @@ func TestConnection(t *testing.T) {
 			&mock.Conn{AddrField: "1", State: conn.Online, NodeIDField: 1},
 			&mock.Conn{AddrField: "2", State: conn.Unknown, NodeIDField: 2},
 		}, nil, balancerConfig.Info{}, false)
-		c, failed := s.GetConnection(WithEndpoint(context.Background(), &mock.Endpoint{AddrField: "2", NodeIDField: 2}))
+		c, failed := s.GetConnection(endpoint.WithNodeID(context.Background(), 2))
 		require.Equal(t, &mock.Conn{AddrField: "1", State: conn.Online, NodeIDField: 1}, c)
 		require.Equal(t, 0, failed)
 	})

internal/balancer/ctx.go

@@ -0,0 +1,19 @@
+package endpoint
+
+import "context"
+
+type (
+	ctxEndpointKey struct{}
+)
+
+func WithNodeID(ctx context.Context, nodeID uint32) context.Context {
+	return context.WithValue(ctx, ctxEndpointKey{}, nodeID)
+}
+
+func ContextNodeID(ctx context.Context) (nodeID uint32, ok bool) {
+	if nodeID, ok = ctx.Value(ctxEndpointKey{}).(uint32); ok {
+		return nodeID, true
+	}
+
+	return 0, false
+}

internal/endpoint/context.go

@@ -6,27 +6,31 @@ import (
 	"time"
 )
 
-type Info interface {
-	NodeID() uint32
-	Address() string
-	Location() string
-	LastUpdated() time.Time
-	LoadFactor() float32
-
-	// Deprecated: LocalDC check "local" by compare endpoint location with discovery "selflocation" field.
-	// It work good only if connection url always point to local dc.
-	// Will be removed after Oct 2024.
-	// Read about versioning policy: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#deprecated
-	LocalDC() bool
-}
-
-type Endpoint interface {
-	Info
+type (
+	NodeID interface {
+		NodeID() uint32
+	}
+	Info interface {
+		NodeID
+		Address() string
+		Location() string
+		LastUpdated() time.Time
+		LoadFactor() float32
+
+		// Deprecated: LocalDC check "local" by compare endpoint location with discovery "selflocation" field.
+		// It work good only if connection url always point to local dc.
+		// Will be removed after Oct 2024.
+		// Read about versioning policy: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#deprecated
+		LocalDC() bool
+	}
+	Endpoint interface {
+		Info
 
-	String() string
-	Copy() Endpoint
-	Touch(opts ...Option)
-}
+		String() string
+		Copy() Endpoint
+		Touch(opts ...Option)
+	}
+)
 
 type endpoint struct { //nolint:maligned
 	mu       sync.RWMutex

internal/endpoint/endpoint.go

@@ -17,8 +17,8 @@ import (
 	"google.golang.org/grpc/metadata"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
-	balancerContext "github.com/ydb-platform/ydb-go-sdk/v3/internal/balancer"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
+	balancerContext "github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/feature"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/operation"
@@ -159,7 +159,7 @@ func newSession(ctx context.Context, cc grpc.ClientConnInterface, config *config
 	s.tableService = Ydb_Table_V1.NewTableServiceClient(
 		conn.WithBeforeFunc(
 			conn.WithContextModifier(cc, func(ctx context.Context) context.Context {
-				return meta.WithTrailerCallback(balancerContext.WithEndpoint(ctx, s), s.checkCloseHint)
+				return meta.WithTrailerCallback(balancerContext.WithNodeID(ctx, s.NodeID()), s.checkCloseHint)
 			}),
 			func() {
 				s.lastUsage.Store(time.Now().Unix())