merging with main

Author: Gleb Brozhe

internal/balancer/balancer.go

@@ -97,7 +97,8 @@ func (b *Balancer) clusterDiscoveryAttempt(ctx context.Context) (err error) {
 		address = "ydb:///" + b.driverConfig.Endpoint()
 		onDone  = trace.DriverOnBalancerClusterDiscoveryAttempt(
 			b.driverConfig.Trace(), &ctx,
-			stack.FunctionID(""),
+			stack.FunctionID(
+				"github.com/ydb-platform/ydb-go-sdk/3/internal/balancer.(*Balancer).clusterDiscoveryAttempt"),
 			address,
 		)
 		endpoints []endpoint.Endpoint
@@ -173,7 +174,8 @@ func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, endpoints []end
 	var (
 		onDone = trace.DriverOnBalancerUpdate(
 			b.driverConfig.Trace(), &ctx,
-			stack.FunctionID(""),
+			stack.FunctionID(
+				"github.com/ydb-platform/ydb-go-sdk/3/internal/balancer.(*Balancer).applyDiscoveredEndpoints"),
 			b.config.DetectLocalDC,
 		)
 		previousConns []conn.Conn
@@ -211,7 +213,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,
-		stack.FunctionID(""),
+		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/balancer.(*Balancer).Close"),
 	)
 	defer func() {
 		onDone(err)
@@ -234,10 +236,19 @@ func New(
 	pool *conn.Pool,
 	opts ...discoveryConfig.Option,
 ) (b *Balancer, finalErr error) {
-	onDone := trace.DriverOnBalancerInit(
-		driverConfig.Trace(), &ctx,
-		stack.FunctionID(""),
-		driverConfig.Balancer().String(),
+	var (
+		onDone = trace.DriverOnBalancerInit(
+			driverConfig.Trace(), &ctx,
+			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/balancer.New"),
+			driverConfig.Balancer().String(),
+		)
+		discoveryConfig = discoveryConfig.New(append(opts,
+			discoveryConfig.With(driverConfig.Common),
+			discoveryConfig.WithEndpoint(driverConfig.Endpoint()),
+			discoveryConfig.WithDatabase(driverConfig.Database()),
+			discoveryConfig.WithSecure(driverConfig.Secure()),
+			discoveryConfig.WithMeta(driverConfig.Meta()),
+		)...)
 	)
 	defer func() {
 		onDone(finalErr)
@@ -248,43 +259,28 @@ func New(
 		pool:            pool,
 		localDCDetector: detectLocalDC,
 	}
+	d := internalDiscovery.New(ctx, pool.Get(
+		endpoint.New(driverConfig.Endpoint()),
+	), discoveryConfig)
+
+	b.discoveryClient = d
 
 	if config := driverConfig.Balancer(); config == nil {
 		b.config = balancerConfig.Config{}
 	} else {
 		b.config = *config
 	}
 
-	err := startDiscovery(ctx, b, opts)
-	if err != nil {
-		return nil, err
-	}
-
-	return b, nil
-}
-
-func startDiscovery(ctx context.Context, b *Balancer, opts []discoveryConfig.Option) error {
-	discoveryConfig := discoveryConfig.New(append(opts,
-		discoveryConfig.With(b.driverConfig.Common),
-		discoveryConfig.WithEndpoint(b.driverConfig.Endpoint()),
-		discoveryConfig.WithDatabase(b.driverConfig.Database()),
-		discoveryConfig.WithSecure(b.driverConfig.Secure()),
-		discoveryConfig.WithMeta(b.driverConfig.Meta()),
-	)...)
-	d := internalDiscovery.New(ctx, b.pool.Get(
-		endpoint.New(b.driverConfig.Endpoint()),
-	), discoveryConfig)
-
-	b.discoveryClient = d
-
 	if b.config.SingleConn {
 		b.applyDiscoveredEndpoints(ctx, []endpoint.Endpoint{
-			endpoint.New(b.driverConfig.Endpoint()),
+			endpoint.New(driverConfig.Endpoint()),
 		}, "")
 	} else {
+		// initialization of balancer state
 		if err := b.clusterDiscovery(ctx); err != nil {
-			return err
+			return nil, xerrors.WithStackTrace(err)
 		}
+		// run background discovering
 		if d := discoveryConfig.Interval(); d > 0 {
 			b.discoveryRepeater = repeater.New(xcontext.ValueOnly(ctx),
 				d, b.clusterDiscoveryAttempt,
@@ -294,7 +290,7 @@ func startDiscovery(ctx context.Context, b *Balancer, opts []discoveryConfig.Opt
 		}
 	}
 
-	return nil
+	return b, nil
 }
 
 func (b *Balancer) Invoke(
@@ -377,7 +373,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,
-		stack.FunctionID(""),
+		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/balancer.(*Balancer).getConn"),
 	)
 	defer func() {
 		if err == nil {

internal/scripting/client.go

@@ -72,7 +72,7 @@ func (c *Client) execute(
 ) (r result.Result, err error) {
 	var (
 		onDone = trace.ScriptingOnExecute(c.config.Trace(), &ctx,
-			stack.FunctionID(""),
+			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/scripting.(*Client).execute"),
 			query, parameters,
 		)
 		a       = allocator.New()
@@ -151,7 +151,7 @@ func (c *Client) explain(
 ) (e table.ScriptingYQLExplanation, err error) {
 	var (
 		onDone = trace.ScriptingOnExplain(c.config.Trace(), &ctx,
-			stack.FunctionID(""),
+			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/scripting.(*Client).explain"),
 			query,
 		)
 		request = &Ydb_Scripting.ExplainYqlRequest{
@@ -218,53 +218,26 @@ func (c *Client) StreamExecute(
 	return r, xerrors.WithStackTrace(err)
 }
 
-func newReceiveStreamResultFunc(
-	stream Ydb_Scripting_V1.ScriptingService_StreamExecuteYqlClient,
-	traceErr func(error) func(error),
-) func(ctx context.Context) (*Ydb.ResultSet, *Ydb_TableStats.QueryStats, error) {
-	return func(ctx context.Context) (*Ydb.ResultSet, *Ydb_TableStats.QueryStats, error) {
-		defer func() {
-			err := recover()
-			if err != nil {
-				errError := err.(error)
-				traceErr(xerrors.HideEOF(errError))
-			}
-		}()
-		select {
-		case <-ctx.Done():
-			return nil, nil, xerrors.WithStackTrace(ctx.Err())
-		default:
-			var response *Ydb_Scripting.ExecuteYqlPartialResponse
-			response, err := stream.Recv()
-			if err != nil {
-				return nil, nil, xerrors.WithStackTrace(err)
-			}
-
-			result := response.GetResult()
-			if result == nil {
-				return nil, nil, xerrors.WithStackTrace(errors.New("no result set"))
-			}
-
-			return result.GetResultSet(), result.GetQueryStats(), nil
-		}
-	}
-}
-
 func (c *Client) streamExecute(
 	ctx context.Context,
 	query string,
 	parameters *params.Parameters,
-) (r scanner.StreamResult, err error) {
+) (r result.StreamResult, err error) {
 	var (
-		onIntermediate = trace.ScriptingOnStreamExecute(c.config.Trace(), &ctx, stack.FunctionID(""), query, parameters)
-		a              = allocator.New()
-		request        = &Ydb_Scripting.ExecuteYqlRequest{
+		onIntermediate = trace.ScriptingOnStreamExecute(c.config.Trace(), &ctx,
+			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/scripting.(*Client).streamExecute"),
+			query, parameters,
+		)
+		a       = allocator.New()
+		request = &Ydb_Scripting.ExecuteYqlRequest{
 			Script:     query,
 			Parameters: parameters.ToYDB(a),
 			OperationParams: operation.Params(
-				ctx, c.config.OperationTimeout(),
+				ctx,
+				c.config.OperationTimeout(),
 				c.config.OperationCancelAfter(),
-				operation.ModeSync),
+				operation.ModeSync,
+			),
 		}
 	)
 	defer func() {
@@ -275,33 +248,53 @@ func (c *Client) streamExecute(
 	}()
 
 	ctx, cancel := xcontext.WithCancel(ctx)
-	defer cancel()
 
 	stream, err := c.service.StreamExecuteYql(ctx, request)
 	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
-	}
-
-	recvFunc := newReceiveStreamResultFunc(stream, onIntermediate)
-	onCloseFunc := func(err error) error {
-		onIntermediate(xerrors.HideEOF(err))(xerrors.HideEOF(err))
-
-		return err
-	}
+		cancel()
 
-	result, err := scanner.NewStream(ctx, recvFunc, onCloseFunc)
-	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
 
-	return result, nil
+	return scanner.NewStream(ctx,
+		func(ctx context.Context) (
+			set *Ydb.ResultSet,
+			stats *Ydb_TableStats.QueryStats,
+			err error,
+		) {
+			defer func() {
+				onIntermediate(xerrors.HideEOF(err))
+			}()
+			select {
+			case <-ctx.Done():
+				return nil, nil, xerrors.WithStackTrace(ctx.Err())
+			default:
+				var response *Ydb_Scripting.ExecuteYqlPartialResponse
+				response, err = stream.Recv()
+				result := response.GetResult()
+				if result == nil || err != nil {
+					return nil, nil, xerrors.WithStackTrace(err)
+				}
+
+				return result.GetResultSet(), result.GetQueryStats(), nil
+			}
+		},
+		func(err error) error {
+			cancel()
+			onIntermediate(xerrors.HideEOF(err))(xerrors.HideEOF(err))
+
+			return err
+		},
+	)
 }
 
 func (c *Client) Close(ctx context.Context) (err error) {
 	if c == nil {
 		return xerrors.WithStackTrace(errNilClient)
 	}
-	onDone := trace.ScriptingOnClose(c.config.Trace(), &ctx, stack.FunctionID(""))
+	onDone := trace.ScriptingOnClose(c.config.Trace(), &ctx,
+		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/scripting.(*Client).Close"),
+	)
 	defer func() {
 		onDone(err)
 	}()