remove trace.ContextDriver and trace.WithDriver

Author: asmyasnikov

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

trace/context.go

@@ -2,24 +2,6 @@ package trace
 
 import "context"
 
-type driverContextKey struct{}
-
-// WithDriver returns deadline which has associated Driver with it.
-func WithDriver(ctx context.Context, t Driver) context.Context {
-	return context.WithValue(ctx,
-		driverContextKey{},
-		ContextDriver(ctx).Compose(t),
-	)
-}
-
-// ContextDriver returns Driver associated with ctx.
-// If there is no Driver associated with ctx then zero value
-// of Driver is returned.
-func ContextDriver(ctx context.Context) Driver {
-	t, _ := ctx.Value(driverContextKey{}).(Driver)
-	return t
-}
-
 type retryContextKey struct{}
 
 // WithRetry returns deadline which has associated Retry with it.