fixed start func may lock while close wait stop backgrounds

Author: rekby

internal/background/worker.go

@@ -7,7 +7,6 @@ import (
 	"sync"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
-	"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"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
@@ -24,9 +23,13 @@ type Worker struct {
 	workers  sync.WaitGroup
 	onceInit sync.Once
 
+	tasksCompleted empty.Chan
+
 	m xsync.Mutex
 
-	closed      xatomic.Bool
+	tasks chan backgroundTask
+
+	closed      bool
 	stop        xcontext.CancelErrFunc
 	closeReason error
 }
@@ -47,27 +50,17 @@ func (b *Worker) Context() context.Context {
 }
 
 func (b *Worker) Start(name string, f CallbackFunc) {
-	if b.closed.Load() {
-		return
-	}
-
 	b.init()
 
-	if b.ctx.Err() != nil {
-		return
-	}
-
 	b.m.WithLock(func() {
-		if b.closed.Load() {
+		if b.closed {
 			return
 		}
-		b.workers.Add(1)
 
-		go func() {
-			defer b.workers.Done()
-
-			pprof.Do(b.ctx, pprof.Labels("background", name), f)
-		}()
+		b.tasks <- backgroundTask{
+			callback: f,
+			name:     name,
+		}
 	})
 }
 
@@ -78,27 +71,34 @@ func (b *Worker) Done() <-chan struct{} {
 }
 
 func (b *Worker) Close(ctx context.Context, err error) error {
-	if b.closed.Swap(true) {
-		return xerrors.WithStackTrace(ErrAlreadyClosed)
-	}
-
 	b.init()
 
+	var resErr error
 	b.m.WithLock(func() {
+		if b.closed {
+			resErr = xerrors.WithStackTrace(ErrAlreadyClosed)
+			return
+		}
+
+		b.closed = true
+
+		close(b.tasks)
 		b.closeReason = err
 		if b.closeReason == nil {
 			b.closeReason = errClosedWithNilReason
 		}
 
 		b.stop(err)
 	})
+	if resErr != nil {
+		return resErr
+	}
+
+	<-b.tasksCompleted
 
 	bgCompleted := make(empty.Chan)
 
 	go func() {
-		b.m.Lock()
-		defer b.m.Unlock()
-
 		b.workers.Wait()
 		close(bgCompleted)
 	}()
@@ -123,5 +123,27 @@ func (b *Worker) init() {
 		if b.ctx == nil {
 			b.ctx, b.stop = xcontext.WithErrCancel(context.Background())
 		}
+		b.tasks = make(chan backgroundTask)
+		b.tasksCompleted = make(empty.Chan)
+		go b.starterLoop()
 	})
 }
+
+func (b *Worker) starterLoop() {
+	defer close(b.tasksCompleted)
+
+	for bgTask := range b.tasks {
+		b.workers.Add(1)
+
+		go func() {
+			defer b.workers.Done()
+
+			pprof.Do(b.ctx, pprof.Labels("background", bgTask.name), bgTask.callback)
+		}()
+	}
+}
+
+type backgroundTask struct {
+	callback CallbackFunc
+	name     string
+}

internal/background/worker_test.go

@@ -95,44 +95,95 @@ func TestWorkerClose(t *testing.T) {
 }
 
 func TestWorkerConcurrentStartAndClose(t *testing.T) {
-	targetClose := int64(10000)
-	parallel := 10
+	xtest.TestManyTimes(t, func(t testing.TB) {
+		targetClose := int64(10000)
+		parallel := 10
 
-	var counter int64
+		var counter int64
 
-	ctx := xtest.Context(t)
-	w := NewWorker(ctx)
+		ctx := xtest.Context(t)
+		w := NewWorker(ctx)
+
+		closeIndex := int64(0)
+		closed := make(empty.Chan)
+		go func() {
+			xtest.SpinWaitCondition(t, nil, func() bool {
+				return atomic.LoadInt64(&counter) > targetClose
+			})
+			require.NoError(t, w.Close(ctx, nil))
+			closeIndex = atomic.LoadInt64(&counter)
+			close(closed)
+		}()
+
+		stopNewStarts := xatomic.Bool{}
+		for i := 0; i < parallel; i++ {
+			go func() {
+				for {
+					if stopNewStarts.Load() {
+						return
+					}
+
+					go func() {
+						w.Start("test", func(ctx context.Context) {
+							atomic.AddInt64(&counter, 1)
+						})
+					}()
+				}
+			}()
+		}
+
+		xtest.WaitChannelClosed(t, closed)
+		runtime.Gosched()
+		require.Equal(t, closeIndex, atomic.LoadInt64(&counter))
+		stopNewStarts.Store(true)
+	})
+}
+
+func TestWorkerStartCompletedWhileLongWait(t *testing.T) {
+	xtest.TestManyTimes(t, func(t testing.TB) {
+		ctx := xtest.Context(t)
+		w := NewWorker(ctx)
+
+		allowStop := make(empty.Chan)
+		closeStarted := make(empty.Chan)
+		w.Start("test", func(ctx context.Context) {
+			<-ctx.Done()
+			close(closeStarted)
 
-	closeIndex := int64(0)
-	closed := make(empty.Chan)
-	go func() {
-		xtest.SpinWaitCondition(t, nil, func() bool {
-			return atomic.LoadInt64(&counter) > targetClose
+			<-allowStop
 		})
-		require.NoError(t, w.Close(ctx, nil))
-		closeIndex = atomic.LoadInt64(&counter)
-		close(closed)
-	}()
 
-	stopNewStarts := xatomic.Bool{}
-	for i := 0; i < parallel; i++ {
+		closed := make(empty.Chan)
+
+		callStartFinished := make(empty.Chan)
 		go func() {
-			for {
-				if stopNewStarts.Load() {
-					return
-				}
+			defer close(callStartFinished)
+			start := time.Now()
 
-				go func() {
-					w.Start("test", func(ctx context.Context) {
-						atomic.AddInt64(&counter, 1)
-					})
-				}()
+			for time.Since(start) < time.Millisecond {
+				w.Start("test2", func(ctx context.Context) {
+					// pass
+				})
 			}
 		}()
-	}
 
-	xtest.WaitChannelClosed(t, closed)
-	runtime.Gosched()
-	require.Equal(t, closeIndex, atomic.LoadInt64(&counter))
-	stopNewStarts.Store(true)
+		go func() {
+			defer close(closed)
+
+			_ = w.Close(ctx, nil)
+		}()
+
+		xtest.WaitChannelClosed(t, callStartFinished)
+		runtime.Gosched()
+
+		select {
+		case <-closed:
+			t.Fatal()
+		default:
+			// pass
+		}
+
+		close(allowStop)
+		xtest.WaitChannelClosed(t, closed)
+	})
 }