A simple bucket based goroutine limiter

bucket.go

package main

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

type Limiter interface {
	Commit(func())
	Wait()
}

type tokenLimiter struct {
	tokenChan chan bool
	wg        sync.WaitGroup
	size      int
}

func (tl *tokenLimiter) Commit(fn func()) {
	// Acquire a token, blocking may happen
	<-tl.tokenChan

	go func() {
		tl.wg.Add(1)
		fn()
		tl.wg.Done()
		// We finish our work, produce a new token
		tl.tokenChan <- true
	}()
}

func (tl *tokenLimiter) Wait() {
	tl.wg.Wait()
}

func NewTokenLimiter(size int) Limiter {
	ch := make(chan bool, size)
	// Make default tokens
	for i := 0; i < size; i++ {
		ch <- true
	}
	return &tokenLimiter{
		tokenChan: ch,
		wg:        sync.WaitGroup{},
		size:      size,
	}
}

func main() {
	tl := NewTokenLimiter(2)

	for i := 0; i < 10; i++ {
		id := i
		tl.Commit(func() {
			fmt.Printf("Running task %d\n", id)
			time.Sleep(time.Duration(100*(id+1)) * time.Millisecond)
		})
	}

	tl.Wait()
}