fix worker start after close

Author: rekby

internal/background/worker.go

@@ -5,9 +5,9 @@ import (
 	"errors"
 	"runtime/pprof"
 	"sync"
-	"sync/atomic"
 
 	"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"
@@ -26,11 +26,13 @@ type Worker struct {
 
 	m xsync.Mutex
 
-	closed      uint32
+	closed      xatomic.Bool
 	stop        xcontext.CancelErrFunc
 	closeReason error
 }
 
+type CallbackFunc func(ctx context.Context)
+
 func NewWorker(parent context.Context) *Worker {
 	w := Worker{}
 	w.ctx, w.stop = xcontext.WithErrCancel(parent)
@@ -44,58 +46,59 @@ func (b *Worker) Context() context.Context {
 	return b.ctx
 }
 
-func (b *Worker) Start(name string, f func(ctx context.Context)) {
-	if atomic.LoadUint32(&b.closed) != 0 {
-		f(b.ctx)
+func (b *Worker) Start(name string, f CallbackFunc) {
+	if b.closed.Load() {
 		return
 	}
 
 	b.init()
 
-	b.m.Lock()
-	defer b.m.Unlock()
-
 	if b.ctx.Err() != nil {
 		return
 	}
 
-	b.workers.Add(1)
-	go func() {
-		defer b.workers.Done()
+	b.m.WithLock(func() {
+		if b.closed.Load() {
+			return
+		}
+		b.workers.Add(1)
 
-		pprof.Do(b.ctx, pprof.Labels("background", name), f)
-	}()
+		go func() {
+			defer b.workers.Done()
+
+			pprof.Do(b.ctx, pprof.Labels("background", name), f)
+		}()
+	})
 }
 
 func (b *Worker) Done() <-chan struct{} {
 	b.init()
 
-	b.m.Lock()
-	defer b.m.Unlock()
-
 	return b.ctx.Done()
 }
 
 func (b *Worker) Close(ctx context.Context, err error) error {
-	if !atomic.CompareAndSwapUint32(&b.closed, 0, 1) {
+	if b.closed.Swap(true) {
 		return xerrors.WithStackTrace(ErrAlreadyClosed)
 	}
 
 	b.init()
 
-	b.m.Lock()
-	defer b.m.Unlock()
-
-	b.closeReason = err
-	if b.closeReason == nil {
-		b.closeReason = errClosedWithNilReason
-	}
+	b.m.WithLock(func() {
+		b.closeReason = err
+		if b.closeReason == nil {
+			b.closeReason = errClosedWithNilReason
+		}
 
-	b.stop(err)
+		b.stop(err)
+	})
 
 	bgCompleted := make(empty.Chan)
 
 	go func() {
+		b.m.Lock()
+		defer b.m.Unlock()
+
 		b.workers.Wait()
 		close(bgCompleted)
 	}()

internal/background/worker_test.go

@@ -0,0 +1,138 @@
+package background
+
+import (
+	"context"
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"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/xtest"
+)
+
+func TestWorkerContext(t *testing.T) {
+	t.Run("Empty", func(t *testing.T) {
+		w := Worker{}
+		require.NotNil(t, w.Context())
+		require.NotNil(t, w.ctx)
+		require.NotNil(t, w.stop)
+	})
+
+	t.Run("Dedicated", func(t *testing.T) {
+		ctx := context.WithValue(context.Background(), "1", "2")
+		w := NewWorker(ctx)
+		require.Equal(t, "2", w.Context().Value("1"))
+	})
+
+	t.Run("Stop", func(t *testing.T) {
+		w := Worker{}
+		ctx := w.Context()
+		require.NoError(t, ctx.Err())
+
+		_ = w.Close(context.Background(), nil)
+		require.Error(t, ctx.Err())
+	})
+}
+
+func TestWorkerStart(t *testing.T) {
+	t.Run("Started", func(t *testing.T) {
+		w := NewWorker(xtest.Context(t))
+		started := make(empty.Chan)
+		w.Start("test", func(ctx context.Context) {
+			close(started)
+		})
+		xtest.WaitChannelClosed(t, started)
+	})
+	t.Run("Stopped", func(t *testing.T) {
+		ctx := xtest.Context(t)
+		w := NewWorker(ctx)
+		_ = w.Close(ctx, nil)
+
+		started := make(empty.Chan)
+		w.Start("test", func(ctx context.Context) {
+			close(started)
+		})
+
+		// expected: no close channel
+		time.Sleep(time.Second / 100)
+		select {
+		case <-started:
+			t.Fatal()
+		default:
+			// pass
+		}
+	})
+}
+
+func TestWorkerClose(t *testing.T) {
+	t.Run("StopBackground", func(t *testing.T) {
+		ctx := xtest.Context(t)
+		w := NewWorker(ctx)
+
+		started := make(empty.Chan)
+		stopped := xatomic.Bool{}
+		w.Start("test", func(innerCtx context.Context) {
+			close(started)
+			<-innerCtx.Done()
+			stopped.Store(true)
+		})
+
+		xtest.WaitChannelClosed(t, started)
+		require.NoError(t, w.Close(ctx, nil))
+		require.True(t, stopped.Load())
+	})
+
+	t.Run("DoubleClose", func(t *testing.T) {
+		ctx := xtest.Context(t)
+		w := NewWorker(ctx)
+		require.NoError(t, w.Close(ctx, nil))
+		require.Error(t, w.Close(ctx, nil))
+	})
+}
+
+func TestWorkerConcurrentStartAndClose(t *testing.T) {
+	targetClose := int64(10000)
+	parallel := 10
+
+	var counter int64
+
+	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)
+}