Refactoring for better tracing

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added context to some internal methods for better tracing
+* Added `trace.FunctionID` helper and `FunctionID` field to trace start info's 
 * Replaced lazy initialization of ydb clients (table, topic, etc.) to explicit initialization on `ydb.Open` step
 
 ## v3.54.1

driver.go

@@ -26,6 +26,7 @@ import (
 	schemeConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/scheme/config"
 	internalScripting "github.com/ydb-platform/ydb-go-sdk/v3/internal/scripting"
 	scriptingConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/scripting/config"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	internalTable "github.com/ydb-platform/ydb-go-sdk/v3/internal/table"
 	tableConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/table/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/topic/topicclientinternal"
@@ -93,7 +94,12 @@ type Driver struct { //nolint:maligned
 }
 
 // Close closes Driver and clear resources
-func (d *Driver) Close(ctx context.Context) error {
+func (d *Driver) Close(ctx context.Context) (finalErr error) {
+	onDone := trace.DriverOnClose(d.config.Trace(), &ctx, stack.FunctionID(0))
+	defer func() {
+		onDone(finalErr)
+	}()
+
 	d.mtx.Lock()
 	defer d.mtx.Unlock()
 
@@ -103,16 +109,16 @@ func (d *Driver) Close(ctx context.Context) error {
 		}
 	}()
 
-	closers := make([]func(context.Context) error, 0)
+	closes := make([]func(context.Context) error, 0)
 	d.childrenMtx.WithLock(func() {
 		for _, child := range d.children {
-			closers = append(closers, child.Close)
+			closes = append(closes, child.Close)
 		}
 		d.children = nil
 	})
 
-	closers = append(
-		closers,
+	closes = append(
+		closes,
 		d.ratelimiter.Close,
 		d.coordination.Close,
 		d.scheme.Close,
@@ -124,8 +130,8 @@ func (d *Driver) Close(ctx context.Context) error {
 	)
 
 	var issues []error
-	for _, closer := range closers {
-		if err := closer(ctx); err != nil {
+	for _, f := range closes {
+		if err := f(ctx); err != nil {
 			issues = append(issues, err)
 		}
 	}
@@ -292,11 +298,7 @@ func connect(ctx context.Context, d *Driver) error {
 	}
 
 	onDone := trace.DriverOnInit(
-		d.config.Trace(),
-		&ctx,
-		d.config.Endpoint(),
-		d.config.Database(),
-		d.config.Secure(),
+		d.config.Trace(), &ctx, stack.FunctionID(2), d.config.Endpoint(), d.config.Database(), d.config.Secure(),
 	)
 	defer func() {
 		onDone(err)
@@ -313,7 +315,7 @@ func connect(ctx context.Context, d *Driver) error {
 	}
 
 	if d.pool == nil {
-		d.pool = conn.NewPool(d.config)
+		d.pool = conn.NewPool(ctx, d.config)
 	}
 
 	d.balancer, err = balancer.New(ctx, d.config, d.pool, d.discoveryOptions...)

internal/balancer/balancer.go

@@ -15,6 +15,7 @@ import (
 	discoveryConfig "github.com/ydb-platform/ydb-go-sdk/v3/internal/discovery/config"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/repeater"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"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/xsync"
@@ -91,7 +92,7 @@ func (b *Balancer) clusterDiscoveryAttempt(ctx context.Context) (err error) {
 	var (
 		address = "ydb:///" + b.driverConfig.Endpoint()
 		onDone  = trace.DriverOnBalancerClusterDiscoveryAttempt(
-			b.driverConfig.Trace(), &ctx, address,
+			b.driverConfig.Trace(), &ctx, stack.FunctionID(0), address,
 		)
 		endpoints []endpoint.Endpoint
 		localDC   string
@@ -127,9 +128,7 @@ func (b *Balancer) clusterDiscoveryAttempt(ctx context.Context) (err error) {
 
 func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, endpoints []endpoint.Endpoint, localDC string) {
 	onDone := trace.DriverOnBalancerUpdate(
-		b.driverConfig.Trace(),
-		&ctx,
-		b.balancerConfig.DetectlocalDC,
+		b.driverConfig.Trace(), &ctx, stack.FunctionID(0), b.balancerConfig.DetectlocalDC,
 	)
 	defer func() {
 		nodes := make([]trace.EndpointInfo, 0, len(endpoints))
@@ -163,8 +162,7 @@ func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, endpoints []end
 
 func (b *Balancer) Close(ctx context.Context) (err error) {
 	onDone := trace.DriverOnBalancerClose(
-		b.driverConfig.Trace(),
-		&ctx,
+		b.driverConfig.Trace(), &ctx, stack.FunctionID(0),
 	)
 	defer func() {
 		onDone(err)
@@ -189,8 +187,7 @@ func New(
 ) (b *Balancer, finalErr error) {
 	var (
 		onDone = trace.DriverOnBalancerInit(
-			driverConfig.Trace(),
-			&ctx,
+			driverConfig.Trace(), &ctx, stack.FunctionID(0),
 		)
 		discoveryConfig = discoveryConfig.New(append(opts,
 			discoveryConfig.With(driverConfig.Common),
@@ -235,7 +232,8 @@ func New(
 		}
 		// run background discovering
 		if d := discoveryConfig.Interval(); d > 0 {
-			b.discoveryRepeater = repeater.New(d, b.clusterDiscoveryAttempt,
+			b.discoveryRepeater = repeater.New(xcontext.WithoutDeadline(ctx),
+				d, b.clusterDiscoveryAttempt,
 				repeater.WithName("discovery"),
 				repeater.WithTrace(b.driverConfig.Trace()),
 			)
@@ -320,8 +318,7 @@ func (b *Balancer) connections() *connectionsState {
 
 func (b *Balancer) getConn(ctx context.Context) (c conn.Conn, err error) {
 	onDone := trace.DriverOnBalancerChooseEndpoint(
-		b.driverConfig.Trace(),
-		&ctx,
+		b.driverConfig.Trace(), &ctx, stack.FunctionID(0),
 	)
 	defer func() {
 		if err == nil {

internal/balancer/local_dc_test.go

@@ -136,7 +136,7 @@ func TestLocalDCDiscovery(t *testing.T) {
 	r := &Balancer{
 		driverConfig:   cfg,
 		balancerConfig: *cfg.Balancer(),
-		pool:           conn.NewPool(cfg),
+		pool:           conn.NewPool(context.Background(), cfg),
 		discoveryClient: discoveryMock{endpoints: []endpoint.Endpoint{
 			&mock.Endpoint{AddrField: "a:123", LocationField: "a"},
 			&mock.Endpoint{AddrField: "b:234", LocationField: "b"},

internal/conn/conn.go

@@ -15,6 +15,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/response"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xatomic"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
@@ -42,8 +43,8 @@ type Conn interface {
 	Ping(ctx context.Context) error
 	IsState(states ...State) bool
 	GetState() State
-	SetState(State) State
-	Unban() State
+	SetState(context.Context, State) State
+	Unban(context.Context) State
 }
 
 type conn struct {
@@ -93,9 +94,7 @@ func (c *conn) IsState(states ...State) bool {
 
 func (c *conn) park(ctx context.Context) (err error) {
 	onDone := trace.DriverOnConnPark(
-		c.config.Trace(),
-		&ctx,
-		c.Endpoint(),
+		c.config.Trace(), &ctx, stack.FunctionID(0), c.Endpoint(),
 	)
 	defer func() {
 		onDone(err)
@@ -112,7 +111,7 @@ func (c *conn) park(ctx context.Context) (err error) {
 		return nil
 	}
 
-	err = c.close()
+	err = c.close(ctx)
 
 	if err != nil {
 		return c.wrapError(err)
@@ -135,20 +134,20 @@ func (c *conn) Endpoint() endpoint.Endpoint {
 	return nil
 }
 
-func (c *conn) SetState(s State) State {
-	return c.setState(s)
+func (c *conn) SetState(ctx context.Context, s State) State {
+	return c.setState(ctx, s)
 }
 
-func (c *conn) setState(s State) State {
+func (c *conn) setState(ctx context.Context, s State) State {
 	if state := State(c.state.Swap(uint32(s))); state != s {
 		trace.DriverOnConnStateChange(
-			c.config.Trace(), c.endpoint.Copy(), state,
+			c.config.Trace(), &ctx, stack.FunctionID(0), c.endpoint.Copy(), state,
 		)(s)
 	}
 	return s
 }
 
-func (c *conn) Unban() State {
+func (c *conn) Unban(ctx context.Context) State {
 	var newState State
 	c.mtx.RLock()
 	cc := c.cc
@@ -159,7 +158,7 @@ func (c *conn) Unban() State {
 		newState = Offline
 	}
 
-	c.setState(newState)
+	c.setState(ctx, newState)
 	return newState
 }
 
@@ -186,9 +185,7 @@ func (c *conn) realConn(ctx context.Context) (cc *grpc.ClientConn, err error) {
 	}
 
 	onDone := trace.DriverOnConnDial(
-		c.config.Trace(),
-		&ctx,
-		c.endpoint.Copy(),
+		c.config.Trace(), &ctx, stack.FunctionID(0), c.endpoint.Copy(),
 	)
 
 	defer func() {
@@ -221,7 +218,7 @@ func (c *conn) realConn(ctx context.Context) (cc *grpc.ClientConn, err error) {
 	}
 
 	c.cc = cc
-	c.setState(Online)
+	c.setState(ctx, Online)
 
 	return c.cc, nil
 }
@@ -243,13 +240,13 @@ func isAvailable(raw *grpc.ClientConn) bool {
 }
 
 // conn must be locked
-func (c *conn) close() (err error) {
+func (c *conn) close(ctx context.Context) (err error) {
 	if c.cc == nil {
 		return nil
 	}
 	err = c.cc.Close()
 	c.cc = nil
-	c.setState(Offline)
+	c.setState(ctx, Offline)
 	return c.wrapError(err)
 }
 
@@ -268,19 +265,17 @@ func (c *conn) Close(ctx context.Context) (err error) {
 	}
 
 	onDone := trace.DriverOnConnClose(
-		c.config.Trace(),
-		&ctx,
-		c.Endpoint(),
+		c.config.Trace(), &ctx, stack.FunctionID(0), c.Endpoint(),
 	)
 	defer func() {
 		onDone(err)
 	}()
 
 	c.closed = true
 
-	err = c.close()
+	err = c.close(ctx)
 
-	c.setState(Destroyed)
+	c.setState(ctx, Destroyed)
 
 	for _, onClose := range c.onClose {
 		onClose(c)
@@ -301,10 +296,7 @@ func (c *conn) Invoke(
 		issues      []trace.Issue
 		useWrapping = UseWrapping(ctx)
 		onDone      = trace.DriverOnConnInvoke(
-			c.config.Trace(),
-			&ctx,
-			c.endpoint,
-			trace.Method(method),
+			c.config.Trace(), &ctx, stack.FunctionID(0), c.endpoint, trace.Method(method),
 		)
 		cc *grpc.ClientConn
 		md = metadata.MD{}
@@ -386,10 +378,7 @@ func (c *conn) NewStream(
 ) (_ grpc.ClientStream, err error) {
 	var (
 		streamRecv = trace.DriverOnConnNewStream(
-			c.config.Trace(),
-			&ctx,
-			c.endpoint.Copy(),
-			trace.Method(method),
+			c.config.Trace(), &ctx, stack.FunctionID(0), c.endpoint.Copy(), trace.Method(method),
 		)
 		useWrapping = UseWrapping(ctx)
 		cc          *grpc.ClientConn