Generic interface for greedily switching amongst arbitrary numbers of exclusive modes.

exclusive_mode_switcher.go

// ExclusiveModeSwitcher implements a generic interface for having an arbitrary number of exclusive modes
// that should only be activated in exclusion to other modes. The specific assumption is that multiple users
// of the same mode should be executed greedily until enough users finish.
type ExclusiveModeSwitcher struct {
	cond *sync.Cond
	users map[string]uint
}

// NewExclusiveModeSwitcher constructs a new mode switcher backend.
func NewExclusiveModeSwitcher() *ExclusiveModeSwitcher {
	return &ExclusiveModeSwitcher{
		cond:    sync.NewCond(new(sync.Mutex)),
		users: make(map[string]uint),
	}
}

// With executes the supplied function once the given mode can be switched to. The supplied context
// function can be used to cancel waiting for the mode.
func (ms *ExclusiveModeSwitcher) With(ctx context.Context, mode string, fn func()) {
	// Need this go-routine to broadcast so we can exit expediently if the context is cancelled.
	go func() {
		select {
		case <-ctx.Done():
			ms.cond.Broadcast()
		}
	}()

	ms.cond.L.Lock()
	// Check if any other modes are in use
	keepWaiting := false
	for {
		keepWaiting = false
		for k, v := range ms.users {
			if k == mode {
				continue
			}
			if v > 0 {
				keepWaiting = true
				break
			}
		}
		// Check we haven't cancelled since last check.
		select {
		case <- ctx.Done():
			// Context is finished. We're not going to keep waiting.
			// Unlock and exit the function.
			ms.cond.L.Unlock()
			return
		default:
		}
		if keepWaiting {
			// Another mode is in use, keep waiting.
			ms.cond.Wait()
		} else {
			// No other modes in use. Assert this mode.
			break
		}
	}
	// Add a user to this mode while we have the lock. This means the mode becomes
	// active.
	ms.users[mode]++
	ms.cond.L.Unlock()
	// We need to wake the other go-routines so any waiters for this mode become active.
	ms.cond.Broadcast()
	fn()
	// When we finish, we need the lock so we can safely remove our mode.
	ms.cond.L.Lock()
	ms.users[mode]--
	ms.cond.L.Unlock()
	// Broadcast so once the mode users reach 0, another mode can run.
	ms.cond.Broadcast()
}