Merge pull request #709 from ydb-platform/bugfix-put-session

fix bug with no put session before second discovery

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed bug with returning session into pool before second re-discovery
+
 ## v3.44.0
 * Added `table/options.WithCallOptions` options for append custom grpc call options into `session.{BulkUpsert,Execute,StreamExecuteScanQuery}`
 * Supported fake transactions in `database/sql` driver over connector option `ydb.WithFakeTx(queryMode)` and connection string param `go_fake_tx`

connection.go

@@ -206,6 +206,7 @@ func (c *Driver) Table() table.Client {
 					c.tableOptions...,
 				)...,
 			),
+			c.balancer.Nodes(),
 		)
 		return c.table.Close
 	})

internal/balancer/balancer.go

@@ -43,6 +43,16 @@ type Balancer struct {
 	onApplyDiscoveredEndpoints []func(ctx context.Context, endpoints []endpoint.Info)
 }
 
+func (b *Balancer) Nodes() (nodes []uint32) {
+	b.mu.WithLock(func() {
+		nodes = make([]uint32, 0, len(b.connectionsState.connByNodeID))
+		for nodeID := range b.connectionsState.connByNodeID {
+			nodes = append(nodes, nodeID)
+		}
+	})
+	return nodes
+}
+
 func (b *Balancer) OnUpdate(onApplyDiscoveredEndpoints func(ctx context.Context, endpoints []endpoint.Info)) {
 	b.mu.WithLock(func() {
 		b.onApplyDiscoveredEndpoints = append(b.onApplyDiscoveredEndpoints, onApplyDiscoveredEndpoints)

internal/table/client.go

@@ -34,16 +34,17 @@ type balancerNotifier interface {
 	OnUpdate(onDiscovery func(ctx context.Context, endpoints []endpoint.Info))
 }
 
-func New(balancer balancerNotifier, config config.Config) *Client {
+func New(balancer balancerNotifier, config config.Config, knownNodeIDs []uint32) *Client {
 	return newClient(balancer, func(ctx context.Context, opts ...sessionBuilderOption) (s *session, err error) {
 		return newSession(ctx, balancer, config, opts...)
-	}, config)
+	}, config, knownNodeIDs)
 }
 
 func newClient(
 	balancer balancerNotifier,
 	builder sessionBuilder,
 	config config.Config,
+	knownNodeIDs []uint32,
 ) *Client {
 	var (
 		ctx    = context.Background()
@@ -55,7 +56,7 @@ func newClient(
 		cc:     balancer,
 		build:  builder,
 		index:  make(map[*session]sessionInfo),
-		nodes:  make(map[uint32]map[*session]struct{}),
+		nodes:  make(map[uint32]map[*session]struct{}, len(knownNodeIDs)*2),
 		idle:   list.New(),
 		waitQ:  list.New(),
 		limit:  config.SizeLimit(),
@@ -67,6 +68,9 @@ func newClient(
 		},
 		done: make(chan struct{}),
 	}
+	for _, nodeID := range knownNodeIDs {
+		c.nodes[nodeID] = make(map[*session]struct{})
+	}
 	if balancer != nil {
 		balancer.OnUpdate(c.updateNodes)
 	}
@@ -268,6 +272,8 @@ func (c *Client) appendSessionToNodes(s *session) {
 		nodeID := s.NodeID()
 		if _, has := c.nodes[nodeID]; has {
 			c.nodes[nodeID][s] = struct{}{}
+		} else {
+			fmt.Println("")
 		}
 	})
 }
@@ -286,8 +292,6 @@ func (c *Client) removeSessionFromNodes(s *session) {
 			delete(sessions, s)
 			if len(sessions) == 0 {
 				delete(c.nodes, nodeID)
-			} else {
-				c.nodes[nodeID] = sessions
 			}
 		}
 	})

internal/table/client_test.go

@@ -421,6 +421,7 @@ func TestSessionPoolRacyGet(t *testing.T) {
 			config.WithSizeLimit(1),
 			config.WithIdleThreshold(-1),
 		),
+		nil,
 	)
 	var (
 		expSession *session
@@ -887,6 +888,7 @@ func newClientWithStubBuilder(
 			cc:    balancer,
 		}).createSession,
 		config.New(options...),
+		nil,
 	)
 }
 

internal/table/session_test.go

@@ -378,6 +378,7 @@ func TestSessionOperationModeOnExecuteDataQuery(t *testing.T) {
 								),
 							),
 							config.New(),
+							nil,
 						)
 						ctx, cancel := context.WithTimeout(
 							context.Background(),
@@ -470,6 +471,7 @@ func TestCreateTableRegression(t *testing.T) {
 			),
 		),
 		config.New(),
+		nil,
 	)
 
 	ctx, cancel := context.WithTimeout(
@@ -564,6 +566,7 @@ func TestDescribeTableRegression(t *testing.T) {
 			),
 		),
 		config.New(),
+		nil,
 	)
 
 	ctx, cancel := context.WithTimeout(