Merge pull request #343 from ydb-platform/next-node

Refactoring of table client

Author: asmyasnikov

connection.go

@@ -116,7 +116,7 @@ type connection struct {
 	pool *conn.Pool
 
 	mtx      sync.Mutex
-	balancer balancer.Connection
+	balancer *balancer.Balancer
 
 	children    map[uint64]Connection
 	childrenMtx xsync.Mutex

internal/balancer/balancer.go

@@ -3,7 +3,6 @@ package balancer
 import (
 	"context"
 	"fmt"
-	"sync"
 
 	"google.golang.org/grpc"
 
@@ -16,24 +15,25 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/repeater"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
 var ErrNoEndpoints = xerrors.Wrap(fmt.Errorf("no endpoints"))
 
-type balancer struct {
+type Balancer struct {
 	driverConfig      config.Config
 	balancerConfig    balancerConfig.Config
 	pool              *conn.Pool
 	discovery         discovery.Client
 	discoveryRepeater repeater.Repeater
 	localDCDetector   func(ctx context.Context, endpoints []endpoint.Endpoint) (string, error)
 
-	m                sync.RWMutex
+	mu               xsync.RWMutex
 	connectionsState *connectionsState
 }
 
-func (b *balancer) clusterDiscovery(ctx context.Context) (err error) {
+func (b *Balancer) clusterDiscovery(ctx context.Context) (err error) {
 	var (
 		onDone = trace.DriverOnBalancerUpdate(
 			b.driverConfig.Trace(),
@@ -73,7 +73,7 @@ func (b *balancer) clusterDiscovery(ctx context.Context) (err error) {
 	return nil
 }
 
-func (b *balancer) applyDiscoveredEndpoints(ctx context.Context, endpoints []endpoint.Endpoint, localDC string) {
+func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, endpoints []endpoint.Endpoint, localDC string) {
 	connections := endpointsToConnections(b.pool, endpoints)
 	for _, c := range connections {
 		b.pool.Allow(ctx, c)
@@ -83,17 +83,16 @@ func (b *balancer) applyDiscoveredEndpoints(ctx context.Context, endpoints []end
 	info := balancerConfig.Info{SelfLocation: localDC}
 	state := newConnectionsState(connections, b.balancerConfig.IsPreferConn, info, b.balancerConfig.AllowFalback)
 
-	b.m.Lock()
-	defer b.m.Unlock()
-
-	b.connectionsState = state
+	b.mu.WithLock(func() {
+		b.connectionsState = state
+	})
 }
 
-func (b *balancer) Discovery() discovery.Client {
+func (b *Balancer) Discovery() discovery.Client {
 	return b.discovery
 }
 
-func (b *balancer) Close(ctx context.Context) (err error) {
+func (b *Balancer) Close(ctx context.Context) (err error) {
 	onDone := trace.DriverOnBalancerClose(
 		b.driverConfig.Trace(),
 		&ctx,
@@ -113,7 +112,7 @@ func (b *balancer) Close(ctx context.Context) (err error) {
 	}
 
 	if len(issues) > 0 {
-		return xerrors.WithStackTrace(xerrors.NewWithIssues("balancer close failed", issues...))
+		return xerrors.WithStackTrace(xerrors.NewWithIssues("Balancer close failed", issues...))
 	}
 
 	return nil
@@ -124,7 +123,7 @@ func New(
 	c config.Config,
 	pool *conn.Pool,
 	opts ...discoveryConfig.Option,
-) (_ Connection, err error) {
+) (b *Balancer, err error) {
 	onDone := trace.DriverOnBalancerInit(
 		c.Trace(),
 		&ctx,
@@ -133,7 +132,7 @@ func New(
 		onDone(err)
 	}()
 
-	b := &balancer{
+	b = &Balancer{
 		driverConfig:    c,
 		pool:            pool,
 		localDCDetector: detectLocalDC,
@@ -187,19 +186,19 @@ func New(
 	return b, nil
 }
 
-func (b *balancer) Endpoint() string {
+func (b *Balancer) Endpoint() string {
 	return b.driverConfig.Endpoint()
 }
 
-func (b *balancer) Name() string {
+func (b *Balancer) Name() string {
 	return b.driverConfig.Database()
 }
 
-func (b *balancer) Secure() bool {
+func (b *Balancer) Secure() bool {
 	return b.driverConfig.Secure()
 }
 
-func (b *balancer) Invoke(
+func (b *Balancer) Invoke(
 	ctx context.Context,
 	method string,
 	args interface{},
@@ -211,7 +210,7 @@ func (b *balancer) Invoke(
 	})
 }
 
-func (b *balancer) NewStream(
+func (b *Balancer) NewStream(
 	ctx context.Context,
 	desc *grpc.StreamDesc,
 	method string,
@@ -228,7 +227,7 @@ func (b *balancer) NewStream(
 	return nil, err
 }
 
-func (b *balancer) wrapCall(ctx context.Context, f func(ctx context.Context, cc conn.Conn) error) (err error) {
+func (b *Balancer) wrapCall(ctx context.Context, f func(ctx context.Context, cc conn.Conn) error) (err error) {
 	cc, err := b.getConn(ctx)
 	if err != nil {
 		return xerrors.WithStackTrace(err)
@@ -257,14 +256,14 @@ func (b *balancer) wrapCall(ctx context.Context, f func(ctx context.Context, cc
 	return nil
 }
 
-func (b *balancer) connections() *connectionsState {
-	b.m.RLock()
-	defer b.m.RUnlock()
+func (b *Balancer) connections() *connectionsState {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
 
 	return b.connectionsState
 }
 
-func (b *balancer) getConn(ctx context.Context) (c conn.Conn, err error) {
+func (b *Balancer) getConn(ctx context.Context) (c conn.Conn, err error) {
 	onDone := trace.DriverOnBalancerChooseEndpoint(
 		b.driverConfig.Trace(),
 		&ctx,
@@ -295,7 +294,7 @@ func (b *balancer) getConn(ctx context.Context) (c conn.Conn, err error) {
 	c, failedCount = state.GetConnection(ctx)
 	if c == nil {
 		return nil, xerrors.WithStackTrace(
-			fmt.Errorf("%w: cannot get connection from balancer after %d attempts", ErrNoEndpoints, failedCount),
+			fmt.Errorf("%w: cannot get connection from Balancer after %d attempts", ErrNoEndpoints, failedCount),
 		)
 	}
 	return c, nil

internal/balancer/ctx.go

@@ -14,6 +14,16 @@ func WithEndpoint(ctx context.Context, endpoint Endpoint) context.Context {
 	return context.WithValue(ctx, ctxEndpointKey{}, endpoint)
 }
 
+type nodeID uint32
+
+func (n nodeID) NodeID() uint32 {
+	return uint32(n)
+}
+
+func WithNodeID(ctx context.Context, id uint32) context.Context {
+	return WithEndpoint(ctx, nodeID(id))
+}
+
 func ContextEndpoint(ctx context.Context) (e Endpoint, ok bool) {
 	if e, ok = ctx.Value(ctxEndpointKey{}).(Endpoint); ok {
 		return e, true

internal/balancer/interfaces.go

@@ -141,7 +141,7 @@ func TestLocalDCDiscovery(t *testing.T) {
 	cfg := config.New(
 		config.WithBalancer(balancers.PreferLocalDC(balancers.Default())),
 	)
-	r := &balancer{
+	r := &Balancer{
 		driverConfig:   cfg,
 		balancerConfig: *cfg.Balancer(),
 		pool:           conn.NewPool(ctx, cfg),