add listening node status from session and close it early

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added listening node status from session and close session early if node not found in balancer
+
 ## v3.37.5
 * Refactoring of `xsql` errors checking
 

internal/balancer/connections_state.go

@@ -79,7 +79,10 @@ func (s *connectionsState) preferConnection(ctx context.Context) conn.Conn {
 	if e, hasPreferEndpoint := ContextEndpoint(ctx); hasPreferEndpoint {
 		c := s.connByNodeID[e.NodeID()]
 		if c != nil && isOkConnection(c, true) {
+			e.Choose(true)
 			return c
+		} else {
+			e.Choose(false)
 		}
 	}
 

internal/balancer/ctx.go

@@ -8,22 +8,15 @@ type (
 
 type Endpoint interface {
 	NodeID() uint32
+
+	// Choose calls from balancer if this endpoint to request chosen or not
+	Choose(chosen bool)
 }
 
 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/mock/conn.go

@@ -82,6 +82,9 @@ type Endpoint struct {
 	LocalDCField  bool
 }
 
+func (e *Endpoint) Choose(bool) {
+}
+
 func (e *Endpoint) NodeID() uint32 {
 	return e.NodeIDField
 }

internal/table/session.go

@@ -53,6 +53,15 @@ type session struct {
 	closeOnce sync.Once
 }
 
+func (s *session) Choose(chosen bool) {
+	if s == nil {
+		return
+	}
+	if !chosen {
+		s.SetStatus(options.SessionClosing)
+	}
+}
+
 func (s *session) NodeID() uint32 {
 	if s == nil {
 		return 0
@@ -165,7 +174,7 @@ func (s *session) Close(ctx context.Context) (err error) {
 		onDone := trace.TableOnSessionDelete(s.config.Trace(), &ctx, s)
 
 		_, err = s.tableService.DeleteSession(
-			balancer.WithEndpoint(ctx, s),
+			ctx,
 			&Ydb_Table.DeleteSessionRequest{
 				SessionId: s.id,
 				OperationParams: operation.Params(ctx,