Merge pull request #416 Fix race condition bugs for background worker

Author: rekby

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fix internal race-condition bugs in internal background worker
+
 ## v3.38.3
 * Added retries to initial discovering
 

internal/background/worker.go

@@ -5,7 +5,6 @@ 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/xcontext"
@@ -24,13 +23,19 @@ type Worker struct {
 	workers  sync.WaitGroup
 	onceInit sync.Once
 
+	tasksCompleted empty.Chan
+
 	m xsync.Mutex
 
-	closed      uint32
+	tasks chan backgroundTask
+
+	closed      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,54 +49,52 @@ 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)
-		return
-	}
-
+func (b *Worker) Start(name string, f CallbackFunc) {
 	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 {
+			return
+		}
 
-		pprof.Do(b.ctx, pprof.Labels("background", name), f)
-	}()
+		b.tasks <- backgroundTask{
+			callback: f,
+			name:     name,
+		}
+	})
 }
 
 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) {
-		return xerrors.WithStackTrace(ErrAlreadyClosed)
-	}
-
 	b.init()
 
-	b.m.Lock()
-	defer b.m.Unlock()
+	var resErr error
+	b.m.WithLock(func() {
+		if b.closed {
+			resErr = xerrors.WithStackTrace(ErrAlreadyClosed)
+			return
+		}
 
-	b.closeReason = err
-	if b.closeReason == nil {
-		b.closeReason = errClosedWithNilReason
+		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.stop(err)
+	<-b.tasksCompleted
 
 	bgCompleted := make(empty.Chan)
 
@@ -120,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(task backgroundTask) {
+			defer b.workers.Done()
+
+			pprof.Do(b.ctx, pprof.Labels("background", task.name), task.callback)
+		}(bgTask)
+	}
+}
+
+type backgroundTask struct {
+	callback CallbackFunc
+	name     string
+}

internal/background/worker_test.go

@@ -0,0 +1,191 @@
+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) {
+		type ctxkey struct{}
+		ctx := context.WithValue(context.Background(), ctxkey{}, "2")
+		w := NewWorker(ctx)
+		require.Equal(t, "2", w.Context().Value(ctxkey{}))
+	})
+
+	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) {
+	xtest.TestManyTimes(t, func(t testing.TB) {
+		targetClose := int64(100)
+		parallel := 10
+
+		var counter int64
+
+		ctx := xtest.Context(t)
+		w := NewWorker(ctx)
+
+		closeIndex := int64(0)
+		closed := make(empty.Chan)
+
+		go func() {
+			defer close(closed)
+
+			xtest.SpinWaitCondition(t, nil, func() bool {
+				return atomic.LoadInt64(&counter) > targetClose
+			})
+
+			require.NoError(t, w.Close(ctx, nil))
+			closeIndex = atomic.LoadInt64(&counter)
+		}()
+
+		stopNewStarts := xatomic.Bool{}
+		for i := 0; i < parallel; i++ {
+			go func() {
+				for {
+					if stopNewStarts.Load() {
+						return
+					}
+
+					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)
+
+			<-allowStop
+		})
+
+		closed := make(empty.Chan)
+
+		callStartFinished := make(empty.Chan)
+		go func() {
+			defer close(callStartFinished)
+			start := time.Now()
+
+			for time.Since(start) < time.Millisecond {
+				w.Start("test2", func(ctx context.Context) {
+					// pass
+				})
+			}
+		}()
+
+		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)
+	})
+}

internal/topic/topicreaderinternal/stream_reconnector_test.go

@@ -222,14 +222,13 @@ func TestTopicReaderReconnectorConnectionLoop(t *testing.T) {
 			},
 			{
 				callback: func(ctx context.Context) (batchedStreamReader, error) {
-					t.Error()
+					t.Fatal()
 					return nil, errors.New("unexpected call")
 				},
 			},
 		}...)
 
-		reconnector.background.Start("test-reconnectionLoop", reconnector.reconnectionLoop)
-		reconnector.reconnectFromBadStream <- nil
+		reconnector.start()
 
 		<-stream1Ready