waitlist.go

package waitlist

import "sync"

type waitChanList []chan struct{}

type WaitList struct {
	m map[string]waitChanList
	sync.RWMutex
}

// NewWaitList - an object to keep track of (labeled) wait queues
func NewWaitList() *WaitList {
	return &WaitList{
		m: make(map[string]waitChanList),
	}
}

// WaitIfBusy does one of two things:
// 1. If there is no current wait create an empty wait list to signal that something is no running
// 2. If there is a wait list create a chan and add it to the channel, pull to block from that channel
func (wm *WaitList) WaitIfBusy(label string) {
	wm.Lock()
	wcs, ok := wm.m[label]
	if ok {
		// a request is in flight, ask this one to wait
		ch := make(chan struct{})
		wm.m[label] = append(wcs, ch)
		wm.Unlock()
		<-ch
	} else {
		// this is currently the only update in flight, make a list for others
		wm.Lock()
		wm.m[label] = make(waitChanList, 0)
		wm.Unlock()
	}
}

// ReleaseAll unblocks all waiting channels. Return how many waiters were released
func (wm *WaitList) ReleaseAll(label string) int {
	wm.Lock()
	defer wm.Unlock()
	w := wm.m[label]
	for _, ch := range w {
		ch <- struct{}{}
	}
	delete(wm.m, label)
	return len(w)
}

// ReleaseOne unblocks the first channel in the wait list and
func (wm *WaitList) ReleaseOne(label string) int {
	wm.Lock()
	defer wm.Unlock()
	w := wm.m[label]
	if len(w) > 0 {
		w[0] <- struct{}{}
		wm.m[label] = w[1:]
	} else {
		delete(wm.m, label)
	}
	return len(w) - 1
}