Merge pull request #187 from ydb-platform/remove-context-trace

Remove context trace

Author: asmyasnikov

CHANGELOG.md

@@ -1,4 +1,5 @@
-* Refactored gtrace tool for generate `Compose` options
+* Removed redundant `trace.With{Table,Driver,Retry}` and `trace.Context{Table,Driver,Retry}` funcs 
+* Refactored tool `gtrace` for generate `Compose` options
 * Added panic recover on trace calls in `Compose` call step
 * Added `trace.With{Discovery,Driver,Coordination,Ratelimiter,Table,Scheme,Scripting}PanicCallback` options
 * Added `ydb.WithPanicCallback` option

config/config.go

@@ -342,7 +342,7 @@ func defaultConfig() (c *config) {
 					return builder.New(
 						ctx,
 						address,
-						trace.ContextDriver(ctx).Compose(c.trace),
+						c.trace,
 					)
 				},
 			),

internal/conn/conn.go

@@ -176,7 +176,7 @@ func (c *conn) GetState() (s State) {
 
 func (c *conn) take(ctx context.Context) (cc *grpc.ClientConn, err error) {
 	onDone := trace.DriverOnConnTake(
-		trace.ContextDriver(ctx).Compose(c.config.Trace()),
+		c.config.Trace(),
 		&ctx,
 		c.endpoint.Copy(),
 	)
@@ -328,7 +328,7 @@ func (c *conn) Invoke(
 		issues   []trace.Issue
 		wrapping = needWrapping(ctx)
 		onDone   = trace.DriverOnConnInvoke(
-			trace.ContextDriver(ctx).Compose(c.config.Trace()),
+			c.config.Trace(),
 			&ctx,
 			c.endpoint,
 			trace.Method(method),
@@ -395,7 +395,7 @@ func (c *conn) NewStream(
 ) (_ grpc.ClientStream, err error) {
 	var (
 		streamRecv = trace.DriverOnConnNewStream(
-			trace.ContextDriver(ctx).Compose(c.config.Trace()),
+			c.config.Trace(),
 			&ctx,
 			c.endpoint.Copy(),
 			trace.Method(method),

internal/conn/pool.go

@@ -91,7 +91,7 @@ func (p *pool) Pessimize(ctx context.Context, cc Conn, cause error) {
 	}
 
 	trace.DriverOnPessimizeNode(
-		trace.ContextDriver(ctx).Compose(p.config.Trace()),
+		p.config.Trace(),
 		&ctx,
 		e,
 		cc.GetState(),

internal/db/database.go

@@ -51,7 +51,7 @@ func New(
 	opts ...discoveryConfig.Option,
 ) (_ Connection, err error) {
 	onDone := trace.DriverOnInit(
-		trace.ContextDriver(ctx).Compose(c.Trace()),
+		c.Trace(),
 		&ctx,
 		c.Endpoint(),
 		c.Database(),

internal/meta/meta.go

@@ -120,7 +120,7 @@ func (m *meta) meta(ctx context.Context) (_ metadata.MD, err error) {
 
 	var token string
 
-	getCredentialsDone := trace.DriverOnGetCredentials(trace.ContextDriver(ctx).Compose(m.trace), &ctx)
+	getCredentialsDone := trace.DriverOnGetCredentials(m.trace, &ctx)
 	defer func() {
 		getCredentialsDone(token, err)
 	}()

internal/table/client.go

@@ -62,7 +62,7 @@ func newClient(
 	builder SessionBuilder,
 	config config.Config,
 ) *client {
-	onDone := trace.TableOnInit(config.Trace().Compose(trace.ContextTable(ctx)), &ctx)
+	onDone := trace.TableOnInit(config.Trace(), &ctx)
 	if builder == nil {
 		builder = func(ctx context.Context) (s Session, err error) {
 			return newSession(ctx, cc, config)
@@ -213,7 +213,7 @@ func (c *client) createSession(ctx context.Context) (s Session, err error) {
 			c.config.CreateSessionTimeout(),
 		)
 
-		onDone := trace.TableOnPoolSessionNew(c.config.Trace().Compose(trace.ContextTable(ctx)), &ctx)
+		onDone := trace.TableOnPoolSessionNew(c.config.Trace(), &ctx)
 
 		defer func() {
 			onDone(s, err)
@@ -285,15 +285,28 @@ func (c *client) createSession(ctx context.Context) (s Session, err error) {
 	}
 }
 
-// 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) {
+type getOptions struct {
+	t trace.Table
+}
+
+type getOption func(o *getOptions)
+
+func withTrace(t trace.Table) getOption {
+	return func(o *getOptions) {
+		o.t = o.t.Compose(t)
+	}
+}
+
+func (c *client) get(ctx context.Context, opts ...getOption) (s Session, err error) {
 	var (
 		i = 0
-		t = c.config.Trace().Compose(trace.ContextTable(ctx))
+		o = getOptions{t: c.config.Trace()}
 	)
+	for _, opt := range opts {
+		opt(&o)
+	}
 
-	onDone := trace.TableOnPoolGet(t, &ctx)
+	onDone := trace.TableOnPoolGet(o.t, &ctx)
 	defer func() {
 		onDone(s, i, err)
 	}()
@@ -329,7 +342,7 @@ func (c *client) Get(ctx context.Context) (s Session, err error) {
 		// are less than maximum amount of touched session. That is, we want to
 		// be fair here and not to lock more goroutines than we could ship
 		// session to.
-		s, err = c.waitFromCh(ctx, t)
+		s, err = c.waitFromCh(ctx, o.t)
 		if err != nil {
 			err = errors.WithStackTrace(err)
 		}
@@ -345,6 +358,12 @@ func (c *client) Get(ctx context.Context) (s Session, err error) {
 	return s, nil
 }
 
+// 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) {
+	return c.get(ctx)
+}
+
 func (c *client) waitFromCh(ctx context.Context, t trace.Table) (s Session, err error) {
 	var (
 		ch *chan Session
@@ -403,7 +422,7 @@ func (c *client) waitFromCh(ctx context.Context, t trace.Table) (s Session, err
 // Get() or Take() calls. In other way it will produce unexpected behavior or
 // panic.
 func (c *client) Put(ctx context.Context, s Session) (err error) {
-	onDone := trace.TableOnPoolPut(c.config.Trace().Compose(trace.ContextTable(ctx)), &ctx, s)
+	onDone := trace.TableOnPoolPut(c.config.Trace(), &ctx, s)
 	defer func() {
 		onDone(err)
 	}()
@@ -447,7 +466,7 @@ func (c *client) Put(ctx context.Context, s Session) (err error) {
 // It returns first error occurred during stale sessions' deletion.
 // Note that even on error it calls Close() on each session.
 func (c *client) Close(ctx context.Context) (err error) {
-	onDone := trace.TableOnClose(c.config.Trace().Compose(trace.ContextTable(ctx)), &ctx)
+	onDone := trace.TableOnClose(c.config.Trace(), &ctx)
 	defer func() {
 		onDone(err)
 	}()
@@ -851,7 +870,7 @@ func (c *client) closeSession(ctx context.Context, s Session, opts ...closeSessi
 	}
 
 	if h.withTrace {
-		onDone := trace.TableOnPoolSessionClose(c.config.Trace().Compose(trace.ContextTable(ctx)), &ctx, s)
+		onDone := trace.TableOnPoolSessionClose(c.config.Trace(), &ctx, s)
 		defer onDone()
 	}
 

internal/table/client_test.go

@@ -195,16 +195,14 @@ func TestSessionPoolCloseWhenWaiting(t *testing.T) {
 			mustGetSession(t, p)
 
 			go func() {
-				_, err := p.Get(
-					trace.WithTable(
-						context.Background(),
-						trace.Table{
-							OnPoolGet: func(trace.TablePoolGetStartInfo) func(trace.TablePoolGetDoneInfo) {
-								get <- struct{}{}
-								return nil
-							},
+				_, err := p.get(
+					context.Background(),
+					withTrace(trace.Table{
+						OnPoolGet: func(trace.TablePoolGetStartInfo) func(trace.TablePoolGetDoneInfo) {
+							get <- struct{}{}
+							return nil
 						},
-					),
+					}),
 				)
 				got <- err
 			}()
@@ -326,7 +324,7 @@ func TestSessionPoolClose(t *testing.T) {
 }
 
 func TestRaceWgClosed(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 55*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 	defer cancel()
 
 	defer func() {
@@ -622,20 +620,18 @@ func TestSessionPoolSizeLimitOverflow(t *testing.T) {
 				}
 			}
 			go func() {
-				session, err := p.Get(
-					trace.WithTable(
-						context.Background(),
-						trace.Table{
-							OnPoolGet: func(trace.TablePoolGetStartInfo) func(trace.TablePoolGetDoneInfo) {
-								get <- struct{}{}
-								return nil
-							},
-							OnPoolWait: func(trace.TablePoolWaitStartInfo) func(trace.TablePoolWaitDoneInfo) {
-								wait <- struct{}{}
-								return nil
-							},
+				session, err := p.get(
+					context.Background(),
+					withTrace(trace.Table{
+						OnPoolGet: func(trace.TablePoolGetStartInfo) func(trace.TablePoolGetDoneInfo) {
+							get <- struct{}{}
+							return nil
+						},
+						OnPoolWait: func(trace.TablePoolWaitStartInfo) func(trace.TablePoolWaitDoneInfo) {
+							wait <- struct{}{}
+							return nil
 						},
-					),
+					}),
 				)
 				got <- sessionAndError{session, err}
 			}()
@@ -1239,17 +1235,7 @@ func mustResetTimer(t *testing.T, ch <-chan time.Duration, exp time.Duration) {
 func mustCreateSession(t *testing.T, p *client) Session {
 	wg := sync.WaitGroup{}
 	defer wg.Wait()
-	s, err := p.createSession(trace.WithTable(
-		context.Background(),
-		trace.Table{
-			OnPoolSessionNew: func(info trace.TablePoolSessionNewStartInfo) func(trace.TablePoolSessionNewDoneInfo) {
-				wg.Add(1)
-				return func(info trace.TablePoolSessionNewDoneInfo) {
-					wg.Done()
-				}
-			},
-		},
-	))
+	s, err := p.createSession(context.Background())
 	if err != nil {
 		t.Helper()
 		t.Fatalf("%s: %v", caller(), err)
@@ -1272,27 +1258,7 @@ func mustPutSession(t *testing.T, p *client, s Session) {
 	wg := sync.WaitGroup{}
 	defer wg.Wait()
 	if err := p.Put(
-		trace.WithTable(
-			context.Background(),
-			trace.Table{
-				OnPoolPut: func(info trace.TablePoolPutStartInfo) func(trace.TablePoolPutDoneInfo) {
-					wg.Add(1)
-					return func(info trace.TablePoolPutDoneInfo) {
-						wg.Done()
-					}
-				},
-				OnPoolSessionClose: func(
-					info trace.TablePoolSessionCloseStartInfo,
-				) func(
-					doneInfo trace.TablePoolSessionCloseDoneInfo,
-				) {
-					wg.Add(1)
-					return func(info trace.TablePoolSessionCloseDoneInfo) {
-						wg.Done()
-					}
-				},
-			},
-		),
+		context.Background(),
 		s,
 	); err != nil {
 		t.Helper()
@@ -1303,23 +1269,7 @@ func mustPutSession(t *testing.T, p *client, s Session) {
 func mustClose(t *testing.T, p *client) {
 	wg := sync.WaitGroup{}
 	defer wg.Wait()
-	if err := p.Close(
-		trace.WithTable(
-			context.Background(),
-			trace.Table{
-				OnPoolSessionClose: func(
-					info trace.TablePoolSessionCloseStartInfo,
-				) func(
-					doneInfo trace.TablePoolSessionCloseDoneInfo,
-				) {
-					wg.Add(1)
-					return func(info trace.TablePoolSessionCloseDoneInfo) {
-						wg.Done()
-					}
-				},
-			},
-		),
-	); err != nil {
+	if err := p.Close(context.Background()); err != nil {
 		t.Helper()
 		t.Fatalf("%s: %v", caller(), err)
 	}

internal/table/retry_test.go

@@ -17,7 +17,6 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/table"
 	"github.com/ydb-platform/ydb-go-sdk/v3/table/options"
 	"github.com/ydb-platform/ydb-go-sdk/v3/testutil"
-	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
 func TestRetryerBackoffRetryCancelation(t *testing.T) {
@@ -212,7 +211,6 @@ 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) {
-	tolerance := 10 * time.Millisecond
 	timeouts := []time.Duration{
 		50 * time.Millisecond,
 		100 * time.Millisecond,
@@ -318,31 +316,8 @@ func TestRetryContextDeadline(t *testing.T) {
 			t.Run(fmt.Sprintf("Timeout=%v,Sleep=%v", timeout, sleep), func(t *testing.T) {
 				ctx, cancel := context.WithTimeout(context.Background(), timeout)
 				defer cancel()
-				start := time.Now()
 				_ = do(
-					trace.WithRetry(
-						ctx,
-						trace.Retry{
-							OnRetry: func(
-								info trace.RetryLoopStartInfo,
-							) func(
-								intermediateInfo trace.RetryLoopIntermediateInfo,
-							) func(
-								trace.RetryLoopDoneInfo,
-							) {
-								return func(
-									info trace.RetryLoopIntermediateInfo,
-								) func(trace.RetryLoopDoneInfo) {
-									return func(info trace.RetryLoopDoneInfo) {
-										latency := time.Since(start)
-										if latency-timeouts[i] > tolerance {
-											t.Errorf("unexpected latency: %v", latency)
-										}
-									}
-								}
-							},
-						},
-					),
+					ctx,
 					p,
 					func(ctx context.Context, _ table.Session) error {
 						select {

internal/table/session.go

@@ -89,7 +89,7 @@ func (s *session) isClosing() bool {
 }
 
 func newSession(ctx context.Context, cc grpc.ClientConnInterface, config config.Config) (s Session, err error) {
-	onDone := trace.TableOnSessionNew(config.Trace().Compose(trace.ContextTable(ctx)), &ctx)
+	onDone := trace.TableOnSessionNew(config.Trace(), &ctx)
 	defer func() {
 		onDone(s, err)
 	}()