Merge branch 'release-3.44'

asmyasnikov · asmyasnikov · commit 05257fb44ddb · 2023-04-25T19:35:30.000+03:00
diff --git a/.github/workflows/check-codegen.yml b/.github/workflows/check-codegen.yml
@@ -3,6 +3,7 @@ on:
   push:
     branches:
       - master
+      - release-*
   pull_request:
   workflow_dispatch:
 jobs:
diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml
@@ -13,10 +13,14 @@ name: "CodeQL"
 
 on:
   push:
-    branches: [ master ]
+    branches:
+      - master
+      - release-*
   pull_request:
     # The branches below must be a subset of the branches above
-    branches: [ master ]
+    branches:
+      - master
+      - release-*
   schedule:
     - cron: '36 3 * * 0'
 
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
@@ -3,6 +3,7 @@ on:
   push:
     branches:
       - master
+      - release-*
   pull_request:
   workflow_dispatch:
 jobs:
diff --git a/.github/workflows/slo.yml b/.github/workflows/slo.yml
@@ -3,6 +3,7 @@ on:
   push:
     branches:
       - master
+      - release-*
   pull_request:
   workflow_dispatch:
 jobs:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -3,6 +3,7 @@ on:
   push:
     branches:
       - master
+      - release-*
   pull_request:
   workflow_dispatch:
 jobs:
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,10 @@
 * Added `table/options.WithPartitions` for configure partitioning policy
 * Marked as deprecated `table/options.WithPartitioningPolicy{UniformPartitions,ExplicitPartitions}` (use `With{UniformPartitions,ExplicitPartitions}` instead)
 
+## v3.44.3
+* Fixed bug of processing endpoint with `node_id=0`
+* Refactored of checking node ID in cluster discovery before `Get` and during in `Put` of session into session pool
+
 ## v3.44.2
 * Removed debug print
 
diff --git a/connection.go b/connection.go
@@ -206,7 +206,6 @@ func (c *Driver) Table() table.Client {
 					c.tableOptions...,
 				)...,
 			),
-			internalTable.WithNodeIDs(c.balancer.Nodes()),
 		)
 		return c.table.Close
 	})
diff --git a/internal/balancer/balancer.go b/internal/balancer/balancer.go
@@ -43,14 +43,13 @@ 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) HasNode(id uint32) bool {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+	if _, has := b.connectionsState.connByNodeID[id]; has {
+		return true
+	}
+	return false
 }
 
 func (b *Balancer) OnUpdate(onApplyDiscoveredEndpoints func(ctx context.Context, endpoints []endpoint.Info)) {
diff --git a/internal/balancer/connections_state.go b/internal/balancer/connections_state.go
@@ -118,19 +118,15 @@ func (s *connectionsState) selectRandomConnection(conns []conn.Conn, allowBanned
 	return nil, failedConns
 }
 
-func connsToNodeIDMap(conns []conn.Conn) (res map[uint32]conn.Conn) {
+func connsToNodeIDMap(conns []conn.Conn) (nodes map[uint32]conn.Conn) {
+	if len(conns) == 0 {
+		return nil
+	}
+	nodes = make(map[uint32]conn.Conn, len(conns))
 	for _, c := range conns {
-		nodeID := c.Endpoint().NodeID()
-
-		if nodeID == 0 {
-			continue
-		}
-		if res == nil {
-			res = make(map[uint32]conn.Conn, len(conns))
-		}
-		res[nodeID] = c
+		nodes[c.Endpoint().NodeID()] = c
 	}
-	return res
+	return nodes
 }
 
 func sortPreferConnections(
diff --git a/internal/balancer/connections_state_test.go b/internal/balancer/connections_state_test.go
@@ -28,7 +28,9 @@ func TestConnsToNodeIDMap(t *testing.T) {
 			source: []conn.Conn{
 				&mock.Conn{NodeIDField: 0},
 			},
-			res: nil,
+			res: map[uint32]conn.Conn{
+				0: &mock.Conn{NodeIDField: 0},
+			},
 		},
 		{
 			name: "NonZero",
@@ -49,6 +51,7 @@ func TestConnsToNodeIDMap(t *testing.T) {
 				&mock.Conn{NodeIDField: 10},
 			},
 			res: map[uint32]conn.Conn{
+				0:  &mock.Conn{NodeIDField: 0},
 				1:  &mock.Conn{NodeIDField: 1},
 				10: &mock.Conn{NodeIDField: 10},
 			},
@@ -247,70 +250,83 @@ func TestNewState(t *testing.T) {
 		{
 			name: "NoFilter",
 			state: newConnectionsState([]conn.Conn{
-				&mock.Conn{AddrField: "1"},
-				&mock.Conn{AddrField: "2"},
+				&mock.Conn{AddrField: "1", NodeIDField: 1},
+				&mock.Conn{AddrField: "2", NodeIDField: 2},
 			}, nil, balancerConfig.Info{}, false),
 			res: &connectionsState{
-				connByNodeID: nil,
+				connByNodeID: map[uint32]conn.Conn{
+					1: &mock.Conn{AddrField: "1", NodeIDField: 1},
+					2: &mock.Conn{AddrField: "2", NodeIDField: 2},
+				},
 				prefer: []conn.Conn{
-					&mock.Conn{AddrField: "1"},
-					&mock.Conn{AddrField: "2"},
+					&mock.Conn{AddrField: "1", NodeIDField: 1},
+					&mock.Conn{AddrField: "2", NodeIDField: 2},
 				},
 				fallback: nil,
 				all: []conn.Conn{
-					&mock.Conn{AddrField: "1"},
-					&mock.Conn{AddrField: "2"},
+					&mock.Conn{AddrField: "1", NodeIDField: 1},
+					&mock.Conn{AddrField: "2", NodeIDField: 2},
 				},
 			},
 		},
 		{
 			name: "FilterDenyFallback",
 			state: newConnectionsState([]conn.Conn{
-				&mock.Conn{AddrField: "t1", LocationField: "t"},
-				&mock.Conn{AddrField: "f1", LocationField: "f"},
-				&mock.Conn{AddrField: "t2", LocationField: "t"},
-				&mock.Conn{AddrField: "f2", LocationField: "f"},
+				&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
+				&mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
+				&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
+				&mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
 			}, func(info balancerConfig.Info, c conn.Conn) bool {
 				return info.SelfLocation == c.Endpoint().Location()
 			}, balancerConfig.Info{SelfLocation: "t"}, false),
 			res: &connectionsState{
-				connByNodeID: nil,
+				connByNodeID: map[uint32]conn.Conn{
+					1: &mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
+					2: &mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
+					3: &mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
+					4: &mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
+				},
 				prefer: []conn.Conn{
-					&mock.Conn{AddrField: "t1", LocationField: "t"},
-					&mock.Conn{AddrField: "t2", LocationField: "t"},
+					&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
+					&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
 				},
 				fallback: nil,
 				all: []conn.Conn{
-					&mock.Conn{AddrField: "t1", LocationField: "t"},
-					&mock.Conn{AddrField: "t2", LocationField: "t"},
+					&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
+					&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
 				},
 			},
 		},
 		{
 			name: "FilterAllowFallback",
 			state: newConnectionsState([]conn.Conn{
-				&mock.Conn{AddrField: "t1", LocationField: "t"},
-				&mock.Conn{AddrField: "f1", LocationField: "f"},
-				&mock.Conn{AddrField: "t2", LocationField: "t"},
-				&mock.Conn{AddrField: "f2", LocationField: "f"},
+				&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
+				&mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
+				&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
+				&mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
 			}, func(info balancerConfig.Info, c conn.Conn) bool {
 				return info.SelfLocation == c.Endpoint().Location()
 			}, balancerConfig.Info{SelfLocation: "t"}, true),
 			res: &connectionsState{
-				connByNodeID: nil,
+				connByNodeID: map[uint32]conn.Conn{
+					1: &mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
+					2: &mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
+					3: &mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
+					4: &mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
+				},
 				prefer: []conn.Conn{
-					&mock.Conn{AddrField: "t1", LocationField: "t"},
-					&mock.Conn{AddrField: "t2", LocationField: "t"},
+					&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
+					&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
 				},
 				fallback: []conn.Conn{
-					&mock.Conn{AddrField: "f1", LocationField: "f"},
-					&mock.Conn{AddrField: "f2", LocationField: "f"},
+					&mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
+					&mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
 				},
 				all: []conn.Conn{
-					&mock.Conn{AddrField: "t1", LocationField: "t"},
-					&mock.Conn{AddrField: "f1", LocationField: "f"},
-					&mock.Conn{AddrField: "t2", LocationField: "t"},
-					&mock.Conn{AddrField: "f2", LocationField: "f"},
+					&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
+					&mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
+					&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
+					&mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
 				},
 			},
 		},
diff --git a/internal/conn/pool.go b/internal/conn/pool.go
@@ -14,12 +14,17 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
+type connsKey struct {
+	address string
+	nodeID  uint32
+}
+
 type Pool struct {
 	usages int64
 	config Config
 	mtx    xsync.RWMutex
 	opts   []grpc.DialOption
-	conns  map[string]*conn
+	conns  map[connsKey]*conn
 	done   chan struct{}
 }
 
@@ -33,7 +38,9 @@ func (p *Pool) Get(endpoint endpoint.Endpoint) Conn {
 		has     bool
 	)
 
-	if cc, has = p.conns[address]; has {
+	key := connsKey{address, endpoint.NodeID()}
+
+	if cc, has = p.conns[key]; has {
 		return cc
 	}
 
@@ -44,15 +51,15 @@ func (p *Pool) Get(endpoint endpoint.Endpoint) Conn {
 		withOnTransportError(p.Ban),
 	)
 
-	p.conns[address] = cc
+	p.conns[key] = cc
 
 	return cc
 }
 
 func (p *Pool) remove(c *conn) {
 	p.mtx.Lock()
 	defer p.mtx.Unlock()
-	delete(p.conns, c.Endpoint().Address())
+	delete(p.conns, connsKey{c.Endpoint().Address(), c.Endpoint().NodeID()})
 }
 
 func (p *Pool) isClosed() bool {
@@ -74,7 +81,7 @@ func (p *Pool) Ban(ctx context.Context, cc Conn, cause error) {
 	p.mtx.RLock()
 	defer p.mtx.RUnlock()
 
-	cc, ok := p.conns[e.Address()]
+	cc, ok := p.conns[connsKey{e.Address(), e.NodeID()}]
 	if !ok {
 		return
 	}
@@ -98,7 +105,7 @@ func (p *Pool) Allow(ctx context.Context, cc Conn) {
 	p.mtx.RLock()
 	defer p.mtx.RUnlock()
 
-	cc, ok := p.conns[e.Address()]
+	cc, ok := p.conns[connsKey{e.Address(), e.NodeID()}]
 	if !ok {
 		return
 	}
@@ -181,7 +188,7 @@ func NewPool(config Config) *Pool {
 		usages: 1,
 		config: config,
 		opts:   config.GrpcDialOptions(),
-		conns:  make(map[string]*conn),
+		conns:  make(map[connsKey]*conn),
 		done:   make(chan struct{}),
 	}
 	if ttl := config.ConnectionTTL(); ttl > 0 {
diff --git a/internal/meta/version.go b/internal/meta/version.go
@@ -3,7 +3,7 @@ package meta
 const (
 	VersionMajor = "3"
 	VersionMinor = "44"
-	VersionPatch = "2"
+	VersionPatch = "3"
 )
 
 const Version = VersionMajor + "." + VersionMinor + "." + VersionPatch
diff --git a/internal/table/client.go b/internal/table/client.go
diff --git a/internal/table/client_test.go b/internal/table/client_test.go
diff --git a/internal/table/session.go b/internal/table/session.go
diff --git a/testutil/driver.go b/testutil/driver.go

Original file line number	Diff line number	Diff line change
`@@ -206,7 +206,6 @@ func (c *Driver) Table() table.Client {`
`206`	`206`	`c.tableOptions...,`
`207`	`207`	`)...,`
`208`	`208`	`),`
`209`		`- internalTable.WithNodeIDs(c.balancer.Nodes()),`
`210`	`209`	`)`
`211`	`210`	`return c.table.Close`
`212`	`211`	`})`