remove trace.ContextTable and trace.WithTable

Author: asmyasnikov

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/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)
 	}()