ctxwg.go

package internal

import (
	"context"
	"sync"
	"sync/atomic"
	"time"
)

type CtxWG struct {
	ctx   context.Context
	count atomic.Value
	mu    sync.Mutex

	indef <-chan time.Time
}

func NewChanWG(ctx context.Context) *CtxWG {
	wg := &CtxWG{}
	wg.count.Store(0)
	wg.ctx = ctx
	return wg
}

func (c *CtxWG) cnt(n int) int {
	c.mu.Lock()
	v := c.count.Load().(int) + n
	c.count.Store(v)
	c.mu.Unlock()
	return v
}

func (c *CtxWG) Add(n int) {
	if c.cnt(0)+n >= 0 {
		c.cnt(n)
	}
}

func (c *CtxWG) Done() {
	c.Add(-1)
}

func (c *CtxWG) Wait(poll time.Duration, timeout time.Duration) <-chan bool {
	readyCh := make(chan bool)

	if poll == 0 {
		poll = 100 * time.Millisecond
	}

	timeoutFN := func() <-chan time.Time {
		return time.After(timeout)
	}
	if timeout == 0 {
		timeoutFN = func() <-chan time.Time {
			return c.indef
		}
	}

	go func() {
		defer close(readyCh)

		for {
			select {
			case <-timeoutFN():
				readyCh <- false
				return
			case <-c.ctx.Done():
				readyCh <- false
				return
			case <-time.After(poll):
				cnt := c.cnt(0)
				if cnt == 0 {
					readyCh <- true
					return
				}
			}
		}
	}()

	return readyCh
}

ctxwg_test.go

package internal

import (
	"context"
	"fmt"
	"testing"
	"time"
)

func TestNewChanWG(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())

	wg := NewChanWG(ctx)
	wg.Add(10)
	go func() {
		for i := 0; i <= 10; i++ {
			func() {
				defer wg.Done()

				fmt.Println("WORKING", i)
				<-time.After(1 * time.Second)
			}()

			if i == 5 {
				fmt.Println("CANCELLING")
				cancel()
			}
		}
	}()

	go func() {
		<-time.After(4 * time.Second)
		cancel()
	}()

	select {
	case <-wg.Wait(100*time.Millisecond, 0):
		fmt.Println("DONE")
	}
}