Pattern for allowing controlled concurrency.

multilock.go

package multilock

import (
	"context"
)

// Resource represents a pool that, when drained, locks until a corresponding
// unlock is called.
type Resource struct {
	c chan struct{}
}

func NewResource(concurrency int) *Resource {
	l := &Resource{
		c: make(chan struct{}, concurrency),
	}
	for i := 0; i < concurrency; i++ {
		l.c <- struct{}{}
	}
	return l
}

// Lock drains the resource by one. If the resource is fully drained, Lock will
// block until one becomes available.
func (l *Resource) Lock() {
	<-l.c
}

// LockContext drains the resource by one. If the resource is fully drained,
// LockContext will block until one becomes available, or else the context is
// cancelled. Returns true if successful, false if cancelled.
func (l *Resource) LockContext(ctx context.Context) bool {
	select {
	case <-l.c:
		return true
	case <-ctx.Done():
		return false
	}
}

// Unlock returns a resource to the pool. Calls to Lock should be matched with
// calls to Unlock.
func (l *Resource) Unlock() {
	l.c <- struct{}{}
}