Merge branch 'master' of https://github.com/ydb-platform/ydb-go-sdk

Author: asmyasnikov

CHANGELOG.md

@@ -1,4 +1,5 @@
-* Moved `internal/deadline.ContextWithoutDeadline` to `internal/xcontext.WithoutDeadline`
+* Refactored `db.Table().CreateSession(ctx)` (maked retryable with internal create session timeout)
+* Refactored `internal/table/client.createSession(ctx)` (got rid of unnecessary goroutine)
 * Supported many user-agent records
 
 ## v3.30.0

internal/table/client.go

@@ -117,17 +117,71 @@ func (c *Client) CreateSession(ctx context.Context, opts ...table.Option) (_ tab
 	if c == nil {
 		return nil, xerrors.WithStackTrace(errNilClient)
 	}
+	createSession := func(ctx context.Context) (s Session, err error) {
+		type result struct {
+			s   Session
+			err error
+		}
+
+		ch := make(chan result)
+
+		go func() {
+			defer close(ch)
+
+			var (
+				s   Session
+				err error
+			)
+
+			createSessionCtx := xcontext.WithoutDeadline(ctx)
+
+			if timeout := c.config.CreateSessionTimeout(); timeout > 0 {
+				var cancel context.CancelFunc
+				createSessionCtx, cancel = context.WithTimeout(createSessionCtx, timeout)
+				defer cancel()
+			}
+
+			s, err = c.build(createSessionCtx)
+
+			select {
+			case ch <- result{
+				s:   s,
+				err: err,
+			}: // nop
+			case <-ctx.Done():
+				if s != nil {
+					_ = s.Close(ctx)
+				}
+			}
+		}()
+
+		select {
+		case r := <-ch:
+			if r.err != nil {
+				return nil, xerrors.WithStackTrace(r.err)
+			}
+			return r.s, nil
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	}
 	var s Session
 	if !c.config.AutoRetry() {
-		s, err = c.build(ctx)
-		return s, xerrors.WithStackTrace(err)
+		s, err = createSession(ctx)
+		if err != nil {
+			return nil, xerrors.WithStackTrace(err)
+		}
+		return s, nil
 	}
 	options := retryOptions(c.config.Trace(), opts...)
 	err = retry.Retry(
 		ctx,
 		func(ctx context.Context) (err error) {
 			s, err = c.build(ctx)
-			return xerrors.WithStackTrace(err)
+			if err != nil {
+				return xerrors.WithStackTrace(err)
+			}
+			return nil
 		},
 		retry.WithIdempotent(true),
 		retry.WithID("CreateSession"),
@@ -145,7 +199,10 @@ func (c *Client) CreateSession(ctx context.Context, opts ...table.Option) (_ tab
 			},
 		}),
 	)
-	return s, xerrors.WithStackTrace(err)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
+	}
+	return s, nil
 }
 
 func (c *Client) isClosed() bool {
@@ -175,17 +232,12 @@ type Session interface {
 	table.ClosableSession
 
 	Status() string
-	OnClose(f func(ctx context.Context))
+	OnClose(f func())
 
 	isClosed() bool
 	isClosing() bool
 }
 
-type createSessionResult struct {
-	s   Session
-	err error
-}
-
 // p.mu must NOT be held.
 func (c *Client) createSession(ctx context.Context) (s Session, err error) {
 	// pre-check the Client size
@@ -201,28 +253,37 @@ func (c *Client) createSession(ctx context.Context) (s Session, err error) {
 		return nil, xerrors.WithStackTrace(errSessionPoolOverflow)
 	}
 
-	resCh := make(chan createSessionResult, 1) // for non-block write
+	type result struct {
+		s   Session
+		err error
+	}
+
+	ch := make(chan result)
 
 	go func() {
+		defer close(ch)
+
 		var (
 			s   Session
 			err error
 		)
 
-		createSessionCtx, cancel := context.WithTimeout(
-			meta.WithAllowFeatures(
-				xcontext.WithoutDeadline(ctx),
-				meta.HintSessionBalancer,
-			),
-			c.config.CreateSessionTimeout(),
+		createSessionCtx := xcontext.WithoutDeadline(ctx)
+
+		createSessionCtx = meta.WithAllowFeatures(createSessionCtx,
+			meta.HintSessionBalancer,
 		)
 
-		onDone := trace.TableOnPoolSessionNew(c.config.Trace(), &ctx)
+		if timeout := c.config.CreateSessionTimeout(); timeout > 0 {
+			var cancel context.CancelFunc
+			createSessionCtx, cancel = context.WithTimeout(createSessionCtx, timeout)
+			defer cancel()
+		}
+
+		onDone := trace.TableOnPoolSessionNew(c.config.Trace(), &createSessionCtx)
 
 		defer func() {
 			onDone(s, err)
-			cancel()
-			close(resCh)
 		}()
 
 		s, err = c.build(createSessionCtx)
@@ -231,7 +292,7 @@ func (c *Client) createSession(ctx context.Context) (s Session, err error) {
 		}
 
 		if s != nil {
-			s.OnClose(func(ctx context.Context) {
+			s.OnClose(func() {
 				c.mu.Lock()
 				defer c.mu.Unlock()
 
@@ -258,32 +319,30 @@ func (c *Client) createSession(ctx context.Context) (s Session, err error) {
 
 		c.mu.WithLock(func() {
 			c.createInProgress--
-			if s != nil {
+			if s != nil && !c.closed {
 				c.index[s] = sessionInfo{}
 				trace.TableOnPoolStateChange(c.config.Trace(), len(c.index), "append")
 			}
 		})
 
-		resCh <- createSessionResult{
+		select {
+		case ch <- result{
 			s:   s,
 			err: err,
+		}: // nop
+		case <-ctx.Done():
+			// nop
 		}
 	}()
 
 	select {
-	case r := <-resCh:
-		if r.s == nil && r.err == nil {
-			panic("ydb: abnormal result of createSession()")
+	case r := <-ch:
+		if r.err != nil {
+			return nil, xerrors.WithStackTrace(r.err)
 		}
-		return r.s, xerrors.WithStackTrace(r.err)
+		return r.s, nil
 	case <-ctx.Done():
-		// read result from resCh for prevention of forgetting session
-		go func() {
-			if r, ok := <-resCh; ok && r.s != nil {
-				_ = r.s.Close(ctx)
-			}
-		}()
-		return nil, ctx.Err()
+		return nil, xerrors.WithStackTrace(ctx.Err())
 	}
 }
 

internal/table/client_test.go

@@ -289,9 +289,9 @@ func TestSessionPoolClose(t *testing.T) {
 		closed3 = false
 	)
 
-	s1.OnClose(func(context.Context) { closed1 = true })
-	s2.OnClose(func(context.Context) { closed2 = true })
-	s3.OnClose(func(context.Context) { closed3 = true })
+	s1.OnClose(func() { closed1 = true })
+	s2.OnClose(func() { closed2 = true })
+	s3.OnClose(func() { closed3 = true })
 
 	mustPutSession(t, p, s1)
 	mustPutSession(t, p, s2)
@@ -524,7 +524,7 @@ func TestSessionPoolRacyGet(t *testing.T) {
 	// Release the first create session request.
 	// Created session must be stored in the Client.
 	expSession = r1.session
-	expSession.OnClose(func(context.Context) {
+	expSession.OnClose(func() {
 		t.Fatalf("unexpected first session close")
 	})
 	close(r1.release)
@@ -1057,7 +1057,7 @@ func TestSessionPoolKeepAliveMinSize(t *testing.T) {
 	sessionBuilder := func(t *testing.T, pool *Client) (Session, chan bool) {
 		s := mustCreateSession(t, pool)
 		closed := make(chan bool)
-		s.OnClose(func(context.Context) {
+		s.OnClose(func() {
 			close(closed)
 		})
 		return s, closed
@@ -1124,7 +1124,7 @@ func TestSessionPoolKeepAliveWithBadSession(t *testing.T) {
 	}()
 	s := mustCreateSession(t, p)
 	closed := make(chan bool)
-	s.OnClose(func(context.Context) {
+	s.OnClose(func() {
 		close(closed)
 	})
 	mustPutSession(t, p, s)

internal/table/retry_test.go

@@ -80,7 +80,7 @@ func TestRetryerBadSession(t *testing.T) {
 	p := SessionProviderFunc{
 		OnGet: func(ctx context.Context) (Session, error) {
 			s := simpleSession(t)
-			s.OnClose(func(context.Context) {
+			s.OnClose(func() {
 				closed[s] = true
 			})
 			return s, nil
@@ -129,7 +129,7 @@ func TestRetryerSessionClosing(t *testing.T) {
 	p := SessionProviderFunc{
 		OnGet: func(ctx context.Context) (Session, error) {
 			s := simpleSession(t)
-			s.OnClose(func(context.Context) {
+			s.OnClose(func() {
 				closed[s] = true
 			})
 			return s, nil

internal/table/session.go

@@ -51,7 +51,7 @@ type session struct {
 	status    options.SessionStatus
 
 	onCloseMtx xsync.RWMutex
-	onClose    []func(ctx context.Context)
+	onClose    []func()
 }
 
 func (s *session) NodeID() uint32 {
@@ -142,7 +142,7 @@ func (s *session) ID() string {
 	return s.id
 }
 
-func (s *session) OnClose(cb func(ctx context.Context)) {
+func (s *session) OnClose(cb func()) {
 	if s.isClosed() {
 		return
 	}
@@ -178,8 +178,8 @@ func (s *session) Close(ctx context.Context) (err error) {
 	// call all close listeners before doing request
 	// firstly this need to clear Client from this session
 	s.onCloseMtx.WithRLock(func() {
-		for _, cb := range s.onClose {
-			cb(ctx)
+		for _, onClose := range s.onClose {
+			onClose()
 		}
 	})
 

table/table.go

@@ -29,15 +29,20 @@ type TxOperation func(ctx context.Context, tx TransactionActor) error
 
 type ClosableSession interface {
 	closer.Closer
+
 	Session
+
+	OnClose(f func())
 }
 
 type Client interface {
 	closer.Closer
 
 	// CreateSession returns session or error for manually control of session lifecycle
-	// CreateSession do not provide retry loop for failed create session requests.
-	// Best effort policy may be implements with outer retry loop includes CreateSession call
+	//
+	// CreateSession implements internal busy loop until one of the following conditions is met:
+	// - context was canceled or deadlined
+	// - session was created
 	CreateSession(ctx context.Context, opts ...Option) (s ClosableSession, err error)
 
 	// Do provide the best effort for execute operation.