wg.go

wg.go

package main

// Original code found at: https://play.golang.org/p/kWmjEuXn0ab

import (
	"fmt"
	"sync"
	"time"
)

func main() {
	g := newGroup(3)
	for i := 0; i < 10; i++ {
		g.Add(1)
		go func(i int) {
			time.Sleep(1 * time.Second)
			fmt.Println(i)
			g.Done()
		}(i)
	}
	g.Wait()
	fmt.Println("done")
}

type group struct {
	mtx sync.Mutex
	cnd *sync.Cond
	max int
	cur int
}

func newGroup(max int) *group {
	g := &group{
		max: max,
	}
	g.cnd = sync.NewCond(&g.mtx)
	return g
}

func (g *group) Add(n int) {
	g.mtx.Lock()
	defer g.mtx.Unlock()

	for (g.cur + n) > g.max {
		g.cnd.Wait()
	}

	g.cur += n
}

func (g *group) Done() {
	g.mtx.Lock()
	defer g.mtx.Unlock()

	g.cur--

	if g.cur < 0 {
		panic("improper use of group")
	}

	g.cnd.Broadcast()
}

func (g *group) Wait() {
	g.mtx.Lock()
	defer g.mtx.Unlock()

	for g.cur > 0 {
		g.cnd.Wait()
	}
}