github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/background/worker_test.go (about)

     1  package background
     2  
     3  import (
     4  	"context"
     5  	"runtime"
     6  	"sync"
     7  	"sync/atomic"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    15  )
    16  
    17  func TestWorkerContext(t *testing.T) {
    18  	t.Run("Empty", func(t *testing.T) {
    19  		w := Worker{}
    20  		require.NotNil(t, w.Context())
    21  		require.NotNil(t, w.ctx)
    22  		require.NotNil(t, w.stop)
    23  	})
    24  
    25  	t.Run("Dedicated", func(t *testing.T) {
    26  		type ctxkey struct{}
    27  		ctx := context.WithValue(context.Background(), ctxkey{}, "2")
    28  		w := NewWorker(ctx)
    29  		require.Equal(t, "2", w.Context().Value(ctxkey{}))
    30  	})
    31  
    32  	t.Run("Stop", func(t *testing.T) {
    33  		w := Worker{}
    34  		ctx := w.Context()
    35  		require.NoError(t, ctx.Err())
    36  
    37  		_ = w.Close(context.Background(), nil)
    38  		require.Error(t, ctx.Err())
    39  	})
    40  }
    41  
    42  func TestWorkerStart(t *testing.T) {
    43  	t.Run("Started", func(t *testing.T) {
    44  		w := NewWorker(xtest.Context(t))
    45  		started := make(empty.Chan)
    46  		w.Start("test", func(ctx context.Context) {
    47  			close(started)
    48  		})
    49  		xtest.WaitChannelClosed(t, started)
    50  	})
    51  	t.Run("Stopped", func(t *testing.T) {
    52  		ctx := xtest.Context(t)
    53  		w := NewWorker(ctx)
    54  		_ = w.Close(ctx, nil)
    55  
    56  		started := make(empty.Chan)
    57  		w.Start("test", func(ctx context.Context) {
    58  			close(started)
    59  		})
    60  
    61  		// expected: no close channel
    62  		time.Sleep(time.Second / 100)
    63  		select {
    64  		case <-started:
    65  			t.Fatal()
    66  		default:
    67  			// pass
    68  		}
    69  	})
    70  }
    71  
    72  func TestWorkerClose(t *testing.T) {
    73  	t.Run("StopBackground", func(t *testing.T) {
    74  		ctx := xtest.Context(t)
    75  		w := NewWorker(ctx)
    76  
    77  		started := make(empty.Chan)
    78  		stopped := atomic.Bool{}
    79  		w.Start("test", func(innerCtx context.Context) {
    80  			close(started)
    81  			<-innerCtx.Done()
    82  			stopped.Store(true)
    83  		})
    84  
    85  		xtest.WaitChannelClosed(t, started)
    86  		require.NoError(t, w.Close(ctx, nil))
    87  		require.True(t, stopped.Load())
    88  	})
    89  
    90  	t.Run("DoubleClose", func(t *testing.T) {
    91  		ctx := xtest.Context(t)
    92  		w := NewWorker(ctx)
    93  		require.NoError(t, w.Close(ctx, nil))
    94  		require.Error(t, w.Close(ctx, nil))
    95  	})
    96  }
    97  
    98  func TestWorkerConcurrentStartAndClose(t *testing.T) {
    99  	xtest.TestManyTimes(t, func(t testing.TB) {
   100  		targetClose := int64(10)
   101  
   102  		parallel := runtime.GOMAXPROCS(0)
   103  
   104  		var counter atomic.Int64
   105  
   106  		ctx := xtest.Context(t)
   107  		w := NewWorker(ctx)
   108  
   109  		stopNewStarts := atomic.Bool{}
   110  		var wgStarters sync.WaitGroup
   111  		for i := 0; i < parallel; i++ {
   112  			wgStarters.Add(1)
   113  			go func() {
   114  				defer wgStarters.Done()
   115  
   116  				for {
   117  					if stopNewStarts.Load() {
   118  						return
   119  					}
   120  
   121  					w.Start("test", func(ctx context.Context) {
   122  						counter.Add(1)
   123  					})
   124  				}
   125  			}()
   126  		}
   127  
   128  		// wait start some backgrounds - for ensure about process worked
   129  		xtest.SpinWaitCondition(t, nil, func() bool {
   130  			return counter.Load() > targetClose
   131  		})
   132  
   133  		require.NoError(t, w.Close(xtest.ContextWithCommonTimeout(ctx, t), nil))
   134  
   135  		stopNewStarts.Store(true)
   136  		xtest.WaitGroup(t, &wgStarters)
   137  
   138  		_, ok := <-w.tasks
   139  		require.False(t, ok)
   140  		require.True(t, w.closed)
   141  	})
   142  }
   143  
   144  func TestWorkerStartCompletedWhileLongWait(t *testing.T) {
   145  	xtest.TestManyTimes(t, func(t testing.TB) {
   146  		ctx := xtest.Context(t)
   147  		w := NewWorker(ctx)
   148  
   149  		allowStop := make(empty.Chan)
   150  		closeStarted := make(empty.Chan)
   151  		w.Start("test", func(ctx context.Context) {
   152  			<-ctx.Done()
   153  			close(closeStarted)
   154  
   155  			<-allowStop
   156  		})
   157  
   158  		closed := make(empty.Chan)
   159  
   160  		callStartFinished := make(empty.Chan)
   161  		go func() {
   162  			defer close(callStartFinished)
   163  			start := time.Now()
   164  
   165  			for time.Since(start) < time.Millisecond {
   166  				w.Start("test2", func(ctx context.Context) {
   167  					// pass
   168  				})
   169  			}
   170  		}()
   171  
   172  		go func() {
   173  			defer close(closed)
   174  
   175  			_ = w.Close(ctx, nil)
   176  		}()
   177  
   178  		xtest.WaitChannelClosed(t, callStartFinished)
   179  		runtime.Gosched()
   180  
   181  		select {
   182  		case <-closed:
   183  			t.Fatal()
   184  		default:
   185  			// pass
   186  		}
   187  
   188  		close(allowStop)
   189  		xtest.WaitChannelClosed(t, closed)
   190  	})
   191  }