diff --git a/group_options.go b/group_options.go index f10cc42..34d680b 100644 --- a/group_options.go +++ b/group_options.go @@ -3,11 +3,13 @@ package syncs import "context" type options struct { - ctx context.Context - cancel context.CancelFunc - preLock bool - termOnError bool - discardIfFull bool + ctx context.Context + cancel context.CancelFunc + preLock bool + termOnError bool + discardIfFull bool + waitQueue int + notBlockCaller bool } // GroupOption functional option type @@ -35,3 +37,15 @@ func Discard(o *options) { o.discardIfFull = true o.preLock = true // discard implies preemptive } + +// NonBlocking wait in caller or in goroutine +func NonBlocking(o *options) { + o.notBlockCaller = true +} + +// WaitQueue size sets maximum amount of waiters that will block +func WaitQueue(size int) GroupOption { + return func(o *options) { + o.waitQueue = size + } +} diff --git a/sizedgroup.go b/sizedgroup.go index b73f77a..87227b8 100644 --- a/sizedgroup.go +++ b/sizedgroup.go @@ -10,8 +10,9 @@ import ( // SizedGroup interface enforces constructor usage and doesn't allow direct creation of sizedGroup type SizedGroup struct { options - wg sync.WaitGroup - sema Locker + wg sync.WaitGroup + sema Locker + waiters *Locker } // NewSizedGroup makes wait group with limited size alive goroutines @@ -21,6 +22,10 @@ func NewSizedGroup(size int, opts ...GroupOption) *SizedGroup { for _, opt := range opts { opt(&res.options) } + if res.options.waitQueue != 0 { + var ptr = NewSemaphore(res.options.waitQueue) + res.waiters = &ptr + } return &res } @@ -40,14 +45,29 @@ func (g *SizedGroup) Go(fn func(ctx context.Context)) { return } + waitForWaiters := false if g.preLock { lockOk := g.sema.TryLock() - if !lockOk && g.discardIfFull { - // lock failed and discardIfFull is set, discard this goroutine - return - } - if !lockOk && !g.discardIfFull { - g.sema.Lock() // make sure we have block until lock is acquired + if !lockOk { + waiterOk := false + if g.waiters != nil { + waiterOk = (*g.waiters).TryLock() + } + if !waiterOk && g.discardIfFull { + // lock failed and discardIfFull is set, discard this goroutine + return + } + if !waiterOk && g.waiters != nil { + // wait for waiters and then for sema + (*g.waiters).Lock() + } + waitForWaiters = true + if !g.options.notBlockCaller { + g.sema.Lock() // make sure we have block until lock is acquired + if g.waiters != nil { + (*g.waiters).Unlock() + } + } } } @@ -55,6 +75,13 @@ func (g *SizedGroup) Go(fn func(ctx context.Context)) { go func() { defer g.wg.Done() + if waitForWaiters && g.preLock && g.options.notBlockCaller { + g.sema.Lock() // make sure we have block until lock is acquired + if g.waiters != nil { + (*g.waiters).Unlock() + } + } + if canceled() { return } diff --git a/sizedgroup_test.go b/sizedgroup_test.go index 581d7e0..31b1978 100644 --- a/sizedgroup_test.go +++ b/sizedgroup_test.go @@ -95,3 +95,78 @@ func ExampleSizedGroup_go() { grp.Wait() // wait for completion } + +func TestSizedGroupWaiters(t *testing.T) { + swg := NewSizedGroup(10, Preemptive, WaitQueue(10)) + var c uint32 + + timeNow := time.Now() + + for i := 0; i < 300; i++ { + swg.Go(func(ctx context.Context) { + time.Sleep(5 * time.Millisecond) + atomic.AddUint32(&c, 1) + }) + } + swg.Wait() + timeDiff := time.Since(timeNow) + assert.GreaterOrEqual(t, timeDiff, time.Millisecond*5*(300/10)) + assert.Equal(t, uint32(300), c, fmt.Sprintf("%d, not all routines have been executed", c)) +} + +func TestSizedGroupWaiters_NonBlocking(t *testing.T) { + swg := NewSizedGroup(10, NonBlocking, Preemptive, WaitQueue(10)) + var c uint32 + + timeNow := time.Now() + + for i := 0; i < 300; i++ { + swg.Go(func(ctx context.Context) { + time.Sleep(5 * time.Millisecond) + atomic.AddUint32(&c, 1) + }) + } + swg.Wait() + timeDiff := time.Since(timeNow) + assert.GreaterOrEqual(t, timeDiff, time.Millisecond*5*(300/10)) + assert.Equal(t, uint32(300), c, fmt.Sprintf("%d, not all routines have been executed", c)) +} + +func TestSizedGroupWaiters_Discard(t *testing.T) { + swg := NewSizedGroup(10, Preemptive, Discard, WaitQueue(5), NonBlocking) + var c uint32 + + timeNow := time.Now() + for i := 0; i < 10; i++ { + newI := i + swg.Go(func(ctx context.Context) { + if newI < 5 { + time.Sleep(3 * time.Second) + } else { + time.Sleep(1 * time.Second) + } + atomic.AddUint32(&c, 1) + }) + } + { + timeNow2 := time.Now() + for i := 0; i < 5; i++ { + swg.Go(func(ctx context.Context) { + time.Sleep(1 * time.Second) + atomic.AddUint32(&c, 1) + }) + } + timeDiff := time.Since(timeNow2) + assert.True(t, timeDiff < time.Second+time.Millisecond*100) + } + for i := 0; i < 10; i++ { + swg.Go(func(ctx context.Context) { + panic("wrong logic") + }) + } + + swg.Wait() + timeDiff := time.Since(timeNow) + assert.True(t, (timeDiff > time.Second*3) && (timeDiff < time.Second*3+time.Millisecond*100)) + assert.Equal(t, uint32(15), c, fmt.Sprintf("%d, not all routines have been executed", c)) +}