Support for server-side load balancing of sessions (#1813)

Co-authored-by: Aleksey Myasnikov <asmyasnikov@ydb.tech>

Author: kprokopenko

CHANGELOG.md

@@ -1,3 +1,6 @@
+* Fixed the support of server-side session balancing in `database/sql` driver
+* Added `ydb.WithDisableSessionBalancer()` driver option for disable server-side session balancing on table and query clients
+
 ## v3.111.3
 * Fixed session closing in `ydb.WithExecuteDataQueryOverQueryClient(true)` scenario
 

internal/config/config.go

@@ -16,6 +16,8 @@ type Common struct {
 	traceRetry           trace.Retry
 	retryBudget          budget.Budget
 
+	disableSessionBalancer bool
+
 	panicCallback func(e interface{})
 }
 
@@ -60,6 +62,10 @@ func (c *Common) RetryBudget() budget.Budget {
 	return c.retryBudget
 }
 
+func (c *Common) DisableSessionBalancer() bool {
+	return c.disableSessionBalancer
+}
+
 // SetOperationTimeout define the maximum amount of time a YDB server will process
 // an operation. After timeout exceeds YDB will try to cancel operation and
 // regardless of the cancellation appropriate error will be returned to
@@ -97,3 +103,7 @@ func SetTraceRetry(c *Common, t *trace.Retry, opts ...trace.RetryComposeOption)
 func SetRetryBudget(c *Common, b budget.Budget) {
 	c.retryBudget = b
 }
+
+func (c *Common) SetDisableSessionBalancer() {
+	c.disableSessionBalancer = true
+}

internal/query/client.go

@@ -10,6 +10,7 @@ import (
 	"google.golang.org/grpc"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/closer"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/operation"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/pool"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/config"
@@ -595,6 +596,10 @@ func New(ctx context.Context, cc grpc.ClientConnInterface, cfg *config.Config) *
 				}
 				defer cancelCreate()
 
+				if !cfg.DisableSessionBalancer() {
+					createCtx = meta.WithAllowFeatures(createCtx, meta.HintSessionBalancer)
+				}
+
 				s, err := createSession(createCtx, client,
 					WithConn(cc),
 					WithDeleteTimeout(cfg.SessionDeleteTimeout()),

internal/query/config/options.go

@@ -85,3 +85,9 @@ func WithLazyTx(lazyTx bool) Option {
 		c.lazyTx = lazyTx
 	}
 }
+
+func WithDisableSessionBalancer() Option {
+	return func(c *Config) {
+		c.SetDisableSessionBalancer()
+	}
+}

internal/query/session_core.go

@@ -13,10 +13,12 @@ import (
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/closer"
 	"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/meta"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/pool"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
@@ -90,6 +92,19 @@ func (core *sessionCore) Status() string {
 	}
 }
 
+func (core *sessionCore) checkCloseHint(md metadata.MD) {
+	for header, values := range md {
+		if header != meta.HeaderServerHints {
+			continue
+		}
+		for _, hint := range values {
+			if hint == meta.HintSessionClose {
+				core.SetStatus(StatusClosing)
+			}
+		}
+	}
+}
+
 type Option func(*sessionCore)
 
 func WithConn(cc grpc.ClientConnInterface) Option {
@@ -150,7 +165,7 @@ func Open(
 	if core.cc != nil {
 		core.Client = Ydb_Query_V1.NewQueryServiceClient(
 			conn.WithContextModifier(core.cc, func(ctx context.Context) context.Context {
-				return balancerContext.WithNodeID(ctx, core.NodeID())
+				return meta.WithTrailerCallback(balancerContext.WithNodeID(ctx, core.NodeID()), core.checkCloseHint)
 			}),
 		)
 	}

internal/table/client.go

@@ -11,6 +11,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/proto"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/pool"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/table/config"
@@ -56,6 +57,10 @@ func New(ctx context.Context, cc grpc.ClientConnInterface, config *config.Config
 			}),
 			pool.WithClock[*Session, Session](config.Clock()),
 			pool.WithCreateItemFunc[*Session, Session](func(ctx context.Context) (*Session, error) {
+				if !config.DisableSessionBalancer() {
+					ctx = meta.WithAllowFeatures(ctx, meta.HintSessionBalancer)
+				}
+
 				return newSession(ctx, cc, config)
 			}),
 			pool.WithTrace[*Session, Session](&pool.Trace{

internal/table/config/config.go

@@ -186,6 +186,12 @@ func UseQuerySession(b bool) Option {
 	}
 }
 
+func WithDisableSessionBalancer() Option {
+	return func(c *Config) {
+		c.SetDisableSessionBalancer()
+	}
+}
+
 // WithClock replaces default clock
 func WithClock(clock clockwork.Clock) Option {
 	return func(c *Config) {

internal/xsql/connector.go

@@ -14,6 +14,7 @@ import (
 	"google.golang.org/grpc"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/bind"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
 	internalQuery "github.com/ydb-platform/ydb-go-sdk/v3/internal/query"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
@@ -117,6 +118,10 @@ func (c *Connector) Connect(ctx context.Context) (_ driver.Conn, finalErr error)
 		stack.FunctionID("database/sql.(*Connector).Connect", stack.Package("database/sql")),
 	)
 
+	if !c.disableServerBalancer {
+		ctx = meta.WithAllowFeatures(ctx, meta.HintSessionBalancer)
+	}
+
 	switch c.processor {
 	case QUERY:
 		s, err := internalQuery.CreateSession(ctx, Ydb_Query_V1.NewQueryServiceClient(c.balancer), c.queryConfig)

options.go

@@ -252,6 +252,20 @@ func WithDatabase(database string) Option {
 	}
 }
 
+// WithDisableSessionBalancer returns an Option that disables session balancing.
+func WithDisableSessionBalancer() Option {
+	return func(ctx context.Context, d *Driver) error {
+		d.queryOptions = append(d.queryOptions, queryConfig.WithDisableSessionBalancer())
+		d.tableOptions = append(d.tableOptions, tableConfig.WithDisableSessionBalancer())
+
+		// implicit disable session balancer for SQL driver;
+		// rule of thumb: if user disables session balancer anywhere, we disable it everywhere
+		d.databaseSQLOptions = append(d.databaseSQLOptions, WithDisableServerBalancer())
+
+		return nil
+	}
+}
+
 // WithSecure defines secure option
 //
 // Warning: use ydb.Open with required Driver string parameter instead

tests/integration/helpers_test.go

@@ -127,39 +127,50 @@ func (scope *scopeT) driverNamed(name string, opts ...ydb.Option) *ydb.Driver {
 		connectionString := scope.ConnectionString()
 		scope.Logf("Connect with connection string, driver name %q: %v", name, connectionString)
 
-		token := scope.AuthToken()
-		if token == "" {
-			scope.Logf("With empty auth token")
-			opts = append(opts, ydb.WithAnonymousCredentials())
-		} else {
-			scope.Logf("With auth token")
-			opts = append(opts, ydb.WithAccessTokenCredentials(token))
-		}
-		cert := scope.CertFile()
-		if cert == "" {
-			scope.Logf("Without tls")
-			opts = append(opts, ydb.WithTLSSInsecureSkipVerify())
-		} else {
-			scope.Logf("With tls")
-			opts = append(opts, ydb.WithCertificatesFromFile(cert))
-		}
+		driver := scope.NonCachingDriver(opts...)
 
-		connectionContext, cancel := context.WithTimeout(scope.Ctx, time.Second*10)
-		defer cancel()
-
-		driver, err := ydb.Open(connectionContext, connectionString, opts...)
 		clean := func() {
 			if driver != nil {
 				scope.Require.NoError(driver.Close(scope.Ctx))
 			}
 		}
 
-		return fixenv.NewGenericResultWithCleanup(driver, clean), err
+		return fixenv.NewGenericResultWithCleanup(driver, clean), nil
 	}
 
 	return fixenv.CacheResult(scope.Env, f, fixenv.CacheOptions{CacheKey: name})
 }
 
+func (scope *scopeT) NonCachingDriver(opts ...ydb.Option) *ydb.Driver {
+	connectionString := scope.ConnectionString()
+	scope.Logf("Connect with connection string: %v", connectionString)
+
+	token := scope.AuthToken()
+	if token == "" {
+		scope.Logf("With empty auth token")
+		opts = append(opts, ydb.WithAnonymousCredentials())
+	} else {
+		scope.Logf("With auth token")
+		opts = append(opts, ydb.WithAccessTokenCredentials(token))
+	}
+	cert := scope.CertFile()
+	if cert == "" {
+		scope.Logf("Without tls")
+		opts = append(opts, ydb.WithTLSSInsecureSkipVerify())
+	} else {
+		scope.Logf("With tls")
+		opts = append(opts, ydb.WithCertificatesFromFile(cert))
+	}
+
+	connectionContext, cancel := context.WithTimeout(scope.Ctx, time.Second*10)
+	defer cancel()
+
+	driver, err := ydb.Open(connectionContext, connectionString, opts...)
+	scope.Require.NoError(err)
+
+	return driver
+}
+
 func (scope *scopeT) SQLDriver(opts ...ydb.ConnectorOption) *sql.DB {
 	f := func() (*fixenv.GenericResult[*sql.DB], error) {
 		driver := scope.Driver()