WIP

internal/conn/pool.go

@@ -21,11 +21,12 @@ import (
 )
 
 type Pool struct {
-	usages      int64
-	config      Config
-	dialOptions []grpc.DialOption
-	conns       xsync.Map[string, *conn]
-	done        chan struct{}
+	usages       int64
+	config       Config
+	dialOptions  []grpc.DialOption
+	conns        xsync.Map[string, *conn]
+	done         chan struct{}
+	parkerStoped chan struct{}
 }
 
 func (p *Pool) DialTimeout() time.Duration {
@@ -160,6 +161,7 @@ func (p *Pool) Release(ctx context.Context) (finalErr error) {
 	)
 
 	wg.Add(cap(errCh))
+
 	p.conns.Range(func(_ string, c *conn) bool {
 		go func(c closer.Closer) {
 			defer wg.Done()
@@ -170,6 +172,9 @@ func (p *Pool) Release(ctx context.Context) (finalErr error) {
 
 		return true
 	})
+
+	<-p.parkerStoped
+
 	wg.Wait()
 	close(errCh)
 
@@ -186,8 +191,11 @@ func (p *Pool) Release(ctx context.Context) (finalErr error) {
 }
 
 func (p *Pool) connParker(ctx context.Context, ttl, interval time.Duration) {
+	defer close(p.parkerStoped)
+
 	ticker := time.NewTicker(interval)
 	defer ticker.Stop()
+
 	for {
 		select {
 		case <-p.done:
@@ -216,10 +224,11 @@ func NewPool(ctx context.Context, config Config) *Pool {
 	defer onDone()
 
 	p := &Pool{
-		usages:      1,
-		config:      config,
-		dialOptions: config.GrpcDialOptions(),
-		done:        make(chan struct{}),
+		usages:       1,
+		config:       config,
+		dialOptions:  config.GrpcDialOptions(),
+		done:         make(chan struct{}),
+		parkerStoped: make(chan struct{}),
 	}
 
 	p.dialOptions = append(p.dialOptions,
@@ -248,6 +257,8 @@ func NewPool(ctx context.Context, config Config) *Pool {
 
 	if ttl := config.ConnectionTTL(); ttl > 0 {
 		go p.connParker(xcontext.ValueOnly(ctx), ttl, ttl/2) //nolint:gomnd
+	} else {
+		close(p.parkerStoped)
 	}
 
 	return p

internal/query/execute_query.go

@@ -2,6 +2,7 @@
 
 import (
 	"context"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xruntime"
 	"io"
 	"time"
 
@@ -146,6 +147,10 @@
 		return nil, xerrors.WithStackTrace(err)
 	}
 
+	xruntime.AddCleanup(r, func(cancelStream context.CancelFunc) {
+		cancelStream()
+	}, executeCancel)
+
 	return r, nil
 }
 

internal/query/result.go

@@ -228,20 +228,20 @@ func nextPart(stream Ydb_Query_V1.QueryService_ExecuteQueryClient) (
 func (r *streamResult) Close(ctx context.Context) (finalErr error) {
 	defer r.closeOnce()
 
-	if r.trace != nil {
-		onDone := trace.QueryOnResultClose(r.trace, &ctx,
-			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*streamResult).Close"),
-		)
-		defer func() {
-			onDone(finalErr)
-		}()
-	}
-
 	for {
 		select {
 		case <-r.closed:
 			return nil
 		default:
+			if r.trace != nil {
+				onDone := trace.QueryOnResultClose(r.trace, &ctx,
+					stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query.(*streamResult).Close"),
+				)
+				defer func() {
+					onDone(finalErr)
+				}()
+			}
+
 			_, err := r.nextPart(ctx)
 			if err != nil {
 				if xerrors.Is(err, io.EOF) {

internal/query/session.go

@@ -2,7 +2,6 @@ package query
 
 import (
 	"context"
-
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"

internal/query/session_core.go

@@ -3,8 +3,6 @@ package query
 import (
 	"context"
 	"os"
-	"runtime/pprof"
-	"strconv"
 	"sync/atomic"
 	"time"
 
@@ -38,10 +36,11 @@ type (
 		SetStatus(code Status)
 	}
 	sessionCore struct {
-		cc     grpc.ClientConnInterface
-		Client Ydb_Query_V1.QueryServiceClient
-		Trace  *trace.Query
-		done   chan struct{}
+		cc                 grpc.ClientConnInterface
+		Client             Ydb_Query_V1.QueryServiceClient
+		Trace              *trace.Query
+		done               chan struct{}
+		attachStreamExited chan struct{}
 
 		deleteTimeout  time.Duration
 		id             string
@@ -120,9 +119,10 @@ func Open(
 	ctx context.Context, client Ydb_Query_V1.QueryServiceClient, opts ...Option,
 ) (_ *sessionCore, finalErr error) {
 	core := &sessionCore{
-		Client: client,
-		Trace:  &trace.Query{},
-		done:   make(chan struct{}),
+		Client:             client,
+		Trace:              &trace.Query{},
+		done:               make(chan struct{}),
+		attachStreamExited: make(chan struct{}),
 	}
 
 	for _, opt := range opts {
@@ -200,6 +200,8 @@ func (core *sessionCore) attach(ctx context.Context) (finalErr error) {
 	core.closeOnce = xsync.OnceFunc(func(ctx context.Context) error {
 		defer cancelAttach()
 
+		<-core.attachStreamExited
+
 		core.SetStatus(StatusClosing)
 		defer core.SetStatus(StatusClosed)
 
@@ -210,31 +212,18 @@ func (core *sessionCore) attach(ctx context.Context) (finalErr error) {
 		return nil
 	})
 
-	if markGoroutineWithLabelNodeIDForAttachStream {
-		pprof.Do(ctx, pprof.Labels(
-			"node_id", strconv.Itoa(int(core.NodeID())),
-		), func(context.Context) {
-			go core.listenAttachStream(attachStream)
-		})
-	} else {
-		go core.listenAttachStream(attachStream)
-	}
+	go core.listenAttachStream(attachStream)
 
 	return nil
 }
 
 func (core *sessionCore) listenAttachStream(attachStream Ydb_Query_V1.QueryService_AttachSessionClient) {
 	defer func() {
-		select {
-		case <-core.done:
-			return
-		default:
-			close(core.done)
-		}
+		close(core.attachStreamExited)
 	}()
 
 	for core.IsAlive() {
-		if _, recvErr := attachStream.Recv(); recvErr != nil {
+		if s, recvErr := attachStream.Recv(); recvErr != nil || s.GetStatus() != Ydb.StatusIds_SUCCESS {
 			return
 		}
 	}

internal/repeater/repeater.go

@@ -8,7 +8,6 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/backoff"
 	"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/trace"
 )
 
@@ -29,8 +28,8 @@ type repeater struct {
 	// Task is a function that must be executed periodically.
 	task func(context.Context) error
 
-	cancel  context.CancelFunc
-	stopped chan struct{}
+	done          chan struct{}
+	workerStopped chan struct{}
 
 	force chan struct{}
 	clock clockwork.Clock
@@ -96,16 +95,14 @@ func New(
 	task func(ctx context.Context) (err error),
 	opts ...option,
 ) *repeater {
-	ctx, cancel := xcontext.WithCancel(ctx)
-
 	r := &repeater{
-		interval: interval,
-		task:     task,
-		cancel:   cancel,
-		stopped:  make(chan struct{}),
-		force:    make(chan struct{}, 1),
-		clock:    clockwork.NewRealClock(),
-		trace:    &trace.Driver{},
+		interval:      interval,
+		task:          task,
+		done:          make(chan struct{}),
+		workerStopped: make(chan struct{}),
+		force:         make(chan struct{}, 1),
+		clock:         clockwork.NewRealClock(),
+		trace:         &trace.Driver{},
 	}
 
 	for _, opt := range opts {
@@ -114,17 +111,17 @@ func New(
 		}
 	}
 
-	go r.worker(ctx, r.clock.NewTicker(interval))
+	go r.worker(r.clock.NewTicker(interval))
 
 	return r
 }
 
 func (r *repeater) stop(onCancel func()) {
-	r.cancel()
+	close(r.done)
 	if onCancel != nil {
 		onCancel()
 	}
-	<-r.stopped
+	<-r.workerStopped
 }
 
 // Stop stops to execute its task.
@@ -162,11 +159,9 @@ func (r *repeater) wakeUp(e Event) (err error) {
 	return r.task(ctx)
 }
 
-func (r *repeater) worker(ctx context.Context, tick clockwork.Ticker) {
-	defer func() {
-		close(r.stopped)
-		tick.Stop()
-	}()
+func (r *repeater) worker(tick clockwork.Ticker) {
+	defer close(r.workerStopped)
+	defer tick.Stop()
 
 	// force returns backoff with delays [500ms...32s]
 	force := backoff.New(
@@ -187,7 +182,7 @@ func (r *repeater) worker(ctx context.Context, tick clockwork.Ticker) {
 		defer force.Stop()
 
 		select {
-		case <-ctx.Done():
+		case <-r.done:
 			return EventCancel
 		case <-tick.Chan():
 			return EventTick
@@ -210,7 +205,7 @@ func (r *repeater) worker(ctx context.Context, tick clockwork.Ticker) {
 
 	for {
 		select {
-		case <-ctx.Done():
+		case <-r.done:
 			return
 
 		case <-tick.Chan():

internal/xtest/context.go

@@ -2,10 +2,9 @@ package xtest
 
 import (
 	"context"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"runtime/pprof"
 	"testing"
-
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 )
 
 func Context(t testing.TB) context.Context {