trace.ClusterInsertDoneInfo.Inserted + trace.ClusterRemoveDoneInfo.Removed + refactoring closeSession() + keeper use closeSession

Author: asmyasnikov

.github/workflows/tests.yml

@@ -57,6 +57,7 @@ jobs:
       YDB_SSL_ROOT_CERTIFICATES_FILE: /tmp/ydb_certs/ca.pem
       YDB_ANONYMOUS_CREDENTIALS: 1
       YDB_SHUTDOWN_URLS: http://localhost:8765/actors/kqp_proxy?force_shutdown=all
+      HIDE_APPLICATION_OUTPUT: 1
     runs-on: ${{ matrix.os }}
     steps:
       - name: Install Go

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 3.11.12
+* Added `trace.ClusterInsertDoneInfo.Inserted` boolean flag for notify about success of insert endpoint into balancer
+* Added `trace.ClusterRemoveDoneInfo.Removed` boolean flag for notify about success of remove endpoint from balancer
+
 ## 3.11.11
 * Reverted usage of `math/rand` (instead `crypto/rand`)
 

internal/balancer/balancer.go

@@ -25,7 +25,7 @@ type Balancer interface {
 	Update(Element, info.Info)
 
 	// Remove removes previously inserted connection.
-	Remove(Element)
+	Remove(Element) bool
 
 	// Contains returns true if Balancer contains requested element.
 	Contains(Element) bool

internal/balancer/multi/multi.go

@@ -65,17 +65,25 @@ func (m *multi) Next() conn.Conn {
 }
 
 func (m *multi) Insert(conn conn.Conn) balancer.Element {
-	n := len(m.filter)
-	h := multiHandle{
-		elements: make([]balancer.Element, n),
-	}
+	var (
+		n = len(m.filter)
+		h = multiHandle{
+			elements: make([]balancer.Element, n),
+		}
+		inserted = false
+	)
+
 	for i, f := range m.filter {
 		if f(conn) {
 			x := m.balancer[i].Insert(conn)
 			h.elements[i] = x
+			inserted = true
 		}
 	}
-	return h
+	if inserted {
+		return h
+	}
+	return nil
 }
 
 func (m *multi) Update(x balancer.Element, info info.Info) {
@@ -86,10 +94,13 @@ func (m *multi) Update(x balancer.Element, info info.Info) {
 	}
 }
 
-func (m *multi) Remove(x balancer.Element) {
+func (m *multi) Remove(x balancer.Element) (removed bool) {
 	for i, x := range x.(multiHandle).elements {
 		if x != nil {
-			m.balancer[i].Remove(x)
+			if m.balancer[i].Remove(x) {
+				removed = true
+			}
 		}
 	}
+	return removed
 }

internal/balancer/rr/rr.go

@@ -103,11 +103,12 @@ func (r *roundRobin) Update(el balancer.Element, info info.Info) {
 	r.belt = r.distribute()
 }
 
-func (r *roundRobin) Remove(x balancer.Element) {
+func (r *roundRobin) Remove(x balancer.Element) bool {
 	el := x.(*list.Element)
 	r.conns.Remove(el)
 	r.inspectMinMax(el.Info)
 	r.belt = r.distribute()
+	return true
 }
 
 func (r *roundRobin) Contains(x balancer.Element) bool {

internal/balancer/single/single.go

@@ -30,11 +30,12 @@ func (b *single) Insert(conn conn.Conn) balancer.Element {
 	return conn
 }
 
-func (b *single) Remove(x balancer.Element) {
+func (b *single) Remove(x balancer.Element) bool {
 	if b.conn != x.(conn.Conn) {
 		panic("ydb: single Conn Balancer: Remove() unknown Conn")
 	}
 	b.conn = nil
+	return true
 }
 
 func (b *single) Update(balancer.Element, info.Info) {}

internal/balancer/stub/stub.go

@@ -13,7 +13,7 @@ type stubBalancer struct {
 	OnNext      func() conn.Conn
 	OnInsert    func(conn.Conn) balancer.Element
 	OnUpdate    func(balancer.Element, info.Info)
-	OnRemove    func(balancer.Element)
+	OnRemove    func(balancer.Element) bool
 	OnPessimize func(context.Context, balancer.Element) error
 	OnContains  func(balancer.Element) bool
 }
@@ -34,9 +34,10 @@ func Balancer() (*list.List, balancer.Balancer) {
 		OnInsert: func(conn conn.Conn) balancer.Element {
 			return cs.Insert(conn)
 		},
-		OnRemove: func(x balancer.Element) {
+		OnRemove: func(x balancer.Element) bool {
 			e := x.(*list.Element)
 			cs.Remove(e)
+			return true
 		},
 		OnUpdate: func(x balancer.Element, info info.Info) {
 			e := x.(*list.Element)
@@ -74,10 +75,11 @@ func (s stubBalancer) Update(el balancer.Element, i info.Info) {
 	}
 }
 
-func (s stubBalancer) Remove(el balancer.Element) {
+func (s stubBalancer) Remove(el balancer.Element) bool {
 	if f := s.OnRemove; f != nil {
-		f(el)
+		return f(el)
 	}
+	return true
 }
 
 func (s stubBalancer) Pessimize(ctx context.Context, el balancer.Element) error {

internal/cluster/cluster.go

@@ -236,13 +236,12 @@ func (c *cluster) Get(ctx context.Context, opts ...crudOption) (cc conn.Conn, er
 
 // Insert inserts new connection into the cluster.
 func (c *cluster) Insert(ctx context.Context, e endpoint.Endpoint, opts ...crudOption) (cc conn.Conn) {
-	onDone := trace.DriverOnClusterInsert(c.config.Trace(), &ctx, e.Copy())
+	var (
+		onDone   = trace.DriverOnClusterInsert(c.config.Trace(), &ctx, e.Copy())
+		inserted = false
+	)
 	defer func() {
-		if cc != nil {
-			onDone(cc.GetState())
-		} else {
-			onDone(conn.Unknown)
-		}
+		onDone(inserted, cc.GetState())
 	}()
 
 	options := parseOptions(opts...)
@@ -266,7 +265,7 @@ func (c *cluster) Insert(ctx context.Context, e endpoint.Endpoint, opts ...crudO
 
 	entry := entry.Entry{Conn: cc}
 
-	entry.InsertInto(c.balancer)
+	inserted = entry.InsertInto(c.balancer)
 
 	c.index[e.Address()] = entry
 
@@ -325,13 +324,12 @@ func (c *cluster) Update(ctx context.Context, e endpoint.Endpoint, opts ...crudO
 
 // Remove removes and closes previously inserted connection.
 func (c *cluster) Remove(ctx context.Context, e endpoint.Endpoint, opts ...crudOption) (cc conn.Conn) {
-	onDone := trace.DriverOnClusterRemove(c.config.Trace(), &ctx, e.Copy())
+	var (
+		onDone  = trace.DriverOnClusterRemove(c.config.Trace(), &ctx, e.Copy())
+		removed = false
+	)
 	defer func() {
-		if cc != nil {
-			onDone(cc.GetState())
-		} else {
-			onDone(conn.Unknown)
-		}
+		onDone(cc.GetState(), removed)
 	}()
 
 	options := parseOptions(opts...)
@@ -352,7 +350,8 @@ func (c *cluster) Remove(ctx context.Context, e endpoint.Endpoint, opts ...crudO
 		panic("ydb: can't remove not-existing endpoint")
 	}
 
-	entry.RemoveFrom(c.balancer)
+	removed = entry.RemoveFrom(c.balancer)
+
 	delete(c.index, e.Address())
 	delete(c.endpoints, e.NodeID())
 

internal/cluster/entry/entry.go

@@ -11,23 +11,22 @@ type Entry struct {
 	Handle balancer.Element
 }
 
-func (c *Entry) InsertInto(b balancer.Balancer) {
+func (c *Entry) InsertInto(b balancer.Balancer) bool {
 	if c.Handle != nil {
 		panic("ydb: Handle already exists")
 	}
 	if c.Conn == nil {
 		panic("ydb: can't insert nil Conn into balancer")
 	}
 	c.Handle = b.Insert(c.Conn)
-	if c.Handle == nil {
-		panic("ydb: balancer has returned nil Handle")
-	}
+	return c.Handle != nil
 }
 
-func (c *Entry) RemoveFrom(b balancer.Balancer) {
+func (c *Entry) RemoveFrom(b balancer.Balancer) bool {
 	if c.Handle == nil {
-		panic("ydb: no Handle to remove from balancer")
+		return false
 	}
 	b.Remove(c.Handle)
 	c.Handle = nil
+	return true
 }

internal/meta/version.go

@@ -1,5 +1,5 @@
 package meta
 
 const (
-	Version = "ydb-go-sdk/3.11.11"
+	Version = "ydb-go-sdk/3.11.12"
 )