Merge pull request #1373 from ydb-platform/create-session-error

added throwing create session error to result error from `internal/table.Client.internalPoolGet`

Author: asmyasnikov

internal/table/client.go

@@ -80,7 +80,7 @@ type Client struct {
 	clock  clockwork.Clock
 
 	// read-write fields
-	mu                xsync.Mutex
+	mu                xsync.RWMutex
 	index             map[*session]sessionInfo
 	createInProgress  int // KIKIMR-9163: in-create-process counter
 	limit             int // Upper bound for Client size.
@@ -352,45 +352,53 @@ func (c *Client) internalPoolCreateSession(ctx context.Context) (s *session, err
 	return s, nil
 }
 
-type getOptions struct {
-	t *trace.Table
+type getSettings struct {
+	t           *trace.Table
+	maxAttempts int
 }
 
-type getOption func(o *getOptions)
+type getOption func(o *getSettings)
 
 func withTrace(t *trace.Table) getOption {
-	return func(o *getOptions) {
+	return func(o *getSettings) {
 		o.t = o.t.Compose(t)
 	}
 }
 
-func (c *Client) internalPoolGet(ctx context.Context, opts ...getOption) (s *session, err error) {
+const maxAttempts = 100
+
+func (c *Client) internalPoolGet(ctx context.Context, opts ...getOption) (s *session, finalErr error) { //nolint:funlen
 	if c.isClosed() {
 		return nil, xerrors.WithStackTrace(errClosedClient)
 	}
 
-	const maxAttempts = 100
-
-	var (
-		start = time.Now()
-		i     = 0
-		o     = getOptions{t: c.config.Trace()}
-	)
+	settings := getSettings{t: c.config.Trace(), maxAttempts: maxAttempts}
 	for _, opt := range opts {
 		if opt != nil {
-			opt(&o)
+			opt(&settings)
 		}
 	}
 
-	onDone := trace.TableOnPoolGet(o.t, &ctx,
+	var (
+		start            = time.Now()
+		i                int
+		lastErr          error
+		createSessionErr error
+		waitFromChErr    error
+	)
+
+	onDone := trace.TableOnPoolGet(settings.t, &ctx,
 		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/table.(*Client).internalPoolGet"),
 	)
 	defer func() {
-		onDone(s, i, err)
+		onDone(s, i, finalErr)
 	}()
 
-	for s == nil && err == nil && i < maxAttempts && !c.isClosed() {
-		i++
+	for ; i < settings.maxAttempts; i++ {
+		if c.isClosed() {
+			return nil, xerrors.WithStackTrace(errClosedClient)
+		}
+
 		s = tryGetIdleSession(c)
 		if s != nil {
 			if !s.isReady() {
@@ -403,18 +411,36 @@ func (c *Client) internalPoolGet(ctx context.Context, opts ...getOption) (s *ses
 			return s, nil
 		}
 
-		s, err = tryCreateNewSession(ctx, c)
-		if s != nil || !isCreateSessionErrorRetriable(err) {
-			return s, xerrors.WithStackTrace(err)
+		s, createSessionErr = tryCreateNewSession(ctx, c)
+		if s != nil {
+			return s, nil
 		}
 
-		s, err = c.internalPoolWaitFromCh(ctx, o.t)
-		if err != nil {
-			err = xerrors.WithStackTrace(err)
+		if !isCreateSessionErrorRetriable(createSessionErr) {
+			return nil, xerrors.WithStackTrace(createSessionErr)
 		}
+
+		s, waitFromChErr = c.internalPoolWaitFromCh(ctx, settings.t)
+		if s != nil {
+			return s, nil
+		}
+
+		if waitFromChErr != nil && !isCreateSessionErrorRetriable(waitFromChErr) {
+			return nil, xerrors.WithStackTrace(waitFromChErr)
+		}
+
+		lastErr = xerrors.WithStackTrace(xerrors.Join(createSessionErr, waitFromChErr))
 	}
 
-	return handleNoProgress(s, err, start, c, i)
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	return nil, xerrors.WithStackTrace(
+		fmt.Errorf("failed to get session from pool ("+
+			"attempts: %d, latency: %v, pool has %d sessions (%d busy, %d idle, %d create_in_progress): %w",
+			i, time.Since(start), len(c.index), len(c.index)-c.idle.Len(), c.idle.Len(), c.createInProgress, lastErr,
+		),
+	)
 }
 
 func tryGetIdleSession(c *Client) *session {
@@ -442,38 +468,6 @@ func tryCreateNewSession(ctx context.Context, c *Client) (*session, error) {
 	return s, err
 }
 
-func handleNoProgress(s *session, err error, start time.Time, c *Client, attempts int) (*session, error) {
-	if s == nil && err == nil {
-		if c.isClosed() {
-			err = xerrors.WithStackTrace(errClosedClient)
-		} else {
-			err = xerrors.WithStackTrace(errNoProgress)
-		}
-	}
-
-	if err != nil {
-		var (
-			index            int
-			idle             int
-			createInProgress int
-		)
-		c.mu.WithLock(func() {
-			index = len(c.index)
-			idle = c.idle.Len()
-			createInProgress = c.createInProgress
-		})
-
-		err = xerrors.WithStackTrace(
-			fmt.Errorf("failed to get session from pool ("+
-				"attempts: %d, latency: %v, pool has %d sessions (%d busy, %d idle, %d create_in_progress): %w",
-				attempts, time.Since(start), index, index-idle, idle, createInProgress, err,
-			),
-		)
-	}
-
-	return s, err
-}
-
 // Get returns first idle session from the Client and removes it from
 // there. If no items stored in Client it creates new one returns it.
 func (c *Client) Get(ctx context.Context) (s *session, err error) {
@@ -708,7 +702,9 @@ func (c *Client) DoTx(ctx context.Context, op table.TxOperation, opts ...table.O
 		}
 
 		defer func() {
-			err = handleTransactionError(ctx, tx, err)
+			if err != nil {
+				_ = tx.Rollback(ctx)
+			}
 		}()
 
 		if err = executeTxOperation(ctx, c, op, tx); err != nil {
@@ -724,22 +720,6 @@ func (c *Client) DoTx(ctx context.Context, op table.TxOperation, opts ...table.O
 	}, config.RetryOptions...)
 }
 
-func handleTransactionError(ctx context.Context, tx table.Transaction, err error) error {
-	if err != nil {
-		errRollback := tx.Rollback(ctx)
-		if errRollback != nil {
-			return xerrors.NewWithIssues("",
-				xerrors.WithStackTrace(err),
-				xerrors.WithStackTrace(errRollback),
-			)
-		}
-
-		return xerrors.WithStackTrace(err)
-	}
-
-	return nil
-}
-
 func executeTxOperation(ctx context.Context, c *Client, op table.TxOperation, tx table.Transaction) (err error) {
 	if panicCallback := c.config.PanicCallback(); panicCallback != nil {
 		defer func() {

internal/table/client_test.go

@@ -108,8 +108,7 @@ func TestSessionPoolCloseWhenWaiting(t *testing.T) {
 				wait = make(chan struct{})
 				got  = make(chan error)
 			)
-			p := newClientWithStubBuilder(
-				t,
+			p := newClientWithStubBuilder(t,
 				testutil.NewBalancer(testutil.WithInvokeHandlers(testutil.InvokeHandlers{
 					testutil.TableCreateSession: func(interface{}) (proto.Message, error) {
 						return &Ydb_Table.CreateSessionResult{
@@ -119,15 +118,13 @@ func TestSessionPoolCloseWhenWaiting(t *testing.T) {
 				})),
 				1,
 				config.WithSizeLimit(1),
-				config.WithTrace(
-					&trace.Table{
-						OnPoolWait: func(trace.TablePoolWaitStartInfo) func(trace.TablePoolWaitDoneInfo) {
-							wait <- struct{}{}
+				config.WithTrace(&trace.Table{
+					OnPoolWait: func(trace.TablePoolWaitStartInfo) func(trace.TablePoolWaitDoneInfo) {
+						wait <- struct{}{}
 
-							return nil
-						},
+						return nil
 					},
-				),
+				}),
 			)
 			defer func() {
 				_ = p.Close(context.Background())
@@ -174,7 +171,7 @@ func TestSessionPoolCloseWhenWaiting(t *testing.T) {
 			case err := <-got:
 				if !xerrors.Is(err, errClosedClient) {
 					t.Fatalf(
-						"unexpected error: %v; want %v",
+						"unexpected error: %q; want %q'",
 						err, errClosedClient,
 					)
 				}

internal/table/errors.go

@@ -28,15 +28,15 @@ var (
 	// requested session is closed early.
 	errSessionClosed = xerrors.Wrap(errors.New("session closed early"))
 
-	// errNoProgress returned by a Client instance to indicate that
-	// operation could not be completed.
-	errNoProgress = xerrors.Wrap(errors.New("no progress"))
-
 	// errParamsRequired returned by a Client instance to indicate that required params is not defined
 	errParamsRequired = xerrors.Wrap(errors.New("params required"))
 )
 
 func isCreateSessionErrorRetriable(err error) bool {
+	if err == nil {
+		panic(err)
+	}
+
 	switch {
 	case
 		xerrors.Is(err, errSessionPoolOverflow),

internal/table/retry_test.go

@@ -15,12 +15,13 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xrand"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
 	"github.com/ydb-platform/ydb-go-sdk/v3/testutil"
 )
 
-func TestRetryerBackoffRetryCancelation(t *testing.T) {
+func TestDoBackoffRetryCancelation(t *testing.T) {
 	for _, testErr := range []error{
 		// Errors leading to Wait repeat.
 		xerrors.Transport(
@@ -78,7 +79,7 @@ func TestRetryerBackoffRetryCancelation(t *testing.T) {
 	}
 }
 
-func TestRetryerBadSession(t *testing.T) {
+func TestDoBadSession(t *testing.T) {
 	closed := make(map[table.Session]bool)
 	p := SessionProviderFunc{
 		OnGet: func(ctx context.Context) (*session, error) {
@@ -133,7 +134,36 @@ func TestRetryerBadSession(t *testing.T) {
 	}
 }
 
-func TestRetryerSessionClosing(t *testing.T) {
+func TestDoCreateSessionError(t *testing.T) {
+	p := SessionProviderFunc{
+		OnGet: func(ctx context.Context) (*session, error) {
+			return nil, xerrors.Operation(xerrors.WithStatusCode(Ydb.StatusIds_UNAVAILABLE))
+		},
+		OnPut: func(ctx context.Context, s *session) error {
+			if s.isClosing() {
+				return s.Close(ctx)
+			}
+
+			return nil
+		},
+	}
+	ctx, cancel := xcontext.WithTimeout(xtest.Context(t), time.Second)
+	defer cancel()
+	err := do(ctx, p, config.New(),
+		func(ctx context.Context, s table.Session) error {
+			return nil
+		},
+		nil,
+	)
+	if !xerrors.Is(err, context.DeadlineExceeded) {
+		t.Errorf("unexpected error: %v", err)
+	}
+	if !xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE) {
+		t.Errorf("unexpected error: %v", err)
+	}
+}
+
+func TestDoSessionClosing(t *testing.T) {
 	closed := make(map[table.Session]bool)
 	p := SessionProviderFunc{
 		OnGet: func(ctx context.Context) (*session, error) {
@@ -183,7 +213,7 @@ func TestRetryerSessionClosing(t *testing.T) {
 	}
 }
 
-func TestRetryerImmediateReturn(t *testing.T) {
+func TestDoImmediateReturn(t *testing.T) {
 	for _, testErr := range []error{
 		xerrors.Operation(
 			xerrors.WithStatusCode(Ydb.StatusIds_GENERIC_ERROR),
@@ -236,7 +266,7 @@ func TestRetryerImmediateReturn(t *testing.T) {
 
 // We are testing all suspentions of custom operation func against to all deadline
 // timeouts - all sub-tests must have latency less than timeouts (+tolerance)
-func TestRetryContextDeadline(t *testing.T) {
+func TestDoContextDeadline(t *testing.T) {
 	timeouts := []time.Duration{
 		50 * time.Millisecond,
 		100 * time.Millisecond,
@@ -372,7 +402,7 @@ func (e *CustomError) Unwrap() error {
 	return e.Err
 }
 
-func TestRetryWithCustomErrors(t *testing.T) {
+func TestDoWithCustomErrors(t *testing.T) {
 	var (
 		limit = 10
 		ctx   = context.Background()

internal/table/session.go

@@ -1297,8 +1297,7 @@ func (s *session) BulkUpsert(ctx context.Context, table string, rows value.Value
 	return nil
 }
 
-// BeginTransaction begins new transaction within given session with given
-// settings.
+// BeginTransaction begins new transaction within given session with given settings.
 func (s *session) BeginTransaction(
 	ctx context.Context,
 	txSettings *table.TransactionSettings,