move implementation to internal/ directory (#828)

internal/replication/driver/replication_stepqueue.go

@@ -0,0 +1,166 @@
+package driver
+
+import (
+	"container/heap"
+	"context"
+	"time"
+
+	"github.com/zrepl/zrepl/internal/daemon/logging/trace"
+	"github.com/zrepl/zrepl/internal/util/chainlock"
+)
+
+type stepQueueRec struct {
+	ident      interface{}
+	targetDate time.Time
+	wakeup     chan StepCompletedFunc
+}
+
+type stepQueue struct {
+	stop chan struct{}
+	reqs chan stepQueueRec
+}
+
+type stepQueueHeapItem struct {
+	idx int
+	req stepQueueRec
+}
+type stepQueueHeap []*stepQueueHeapItem
+
+func (h stepQueueHeap) Less(i, j int) bool {
+	return h[i].req.targetDate.Before(h[j].req.targetDate)
+}
+
+func (h stepQueueHeap) Swap(i, j int) {
+	h[i], h[j] = h[j], h[i]
+	h[i].idx = i
+	h[j].idx = j
+}
+
+func (h stepQueueHeap) Len() int {
+	return len(h)
+}
+
+func (h *stepQueueHeap) Push(elem interface{}) {
+	hitem := elem.(*stepQueueHeapItem)
+	hitem.idx = h.Len()
+	*h = append(*h, hitem)
+}
+
+func (h *stepQueueHeap) Pop() interface{} {
+	elem := (*h)[h.Len()-1]
+	elem.idx = -1
+	*h = (*h)[:h.Len()-1]
+	return elem
+}
+
+// returned stepQueue must be closed with method Close
+func newStepQueue() *stepQueue {
+	q := &stepQueue{
+		stop: make(chan struct{}),
+		reqs: make(chan stepQueueRec),
+	}
+	return q
+}
+
+// the returned done function must be called to free resources
+// allocated by the call to Start
+//
+// No WaitReady calls must be active at the time done is called
+// The behavior of calling WaitReady after done was called is undefined
+func (q *stepQueue) Start(concurrency int) (done func()) {
+	if concurrency < 1 {
+		panic("concurrency must be >= 1")
+	}
+	// l protects pending and queueItems
+	l := chainlock.New()
+	pendingCond := l.NewCond()
+	// priority queue
+	pending := &stepQueueHeap{}
+	// ident => queueItem
+	queueItems := make(map[interface{}]*stepQueueHeapItem)
+	// stopped is used for cancellation of "wake" goroutine
+	stopped := false
+	active := 0
+	go func() { // "stopper" goroutine
+		<-q.stop
+		defer l.Lock().Unlock()
+		stopped = true
+		pendingCond.Broadcast()
+	}()
+	go func() { // "reqs" goroutine
+		for {
+			select {
+			case <-q.stop:
+				select {
+				case <-q.reqs:
+					panic("WaitReady call active while calling Close")
+				default:
+					return
+				}
+			case req := <-q.reqs:
+				func() {
+					defer l.Lock().Unlock()
+					if _, ok := queueItems[req.ident]; ok {
+						panic("WaitReady must not be called twice for the same ident")
+					}
+					qitem := &stepQueueHeapItem{
+						req: req,
+					}
+					queueItems[req.ident] = qitem
+					heap.Push(pending, qitem)
+					pendingCond.Broadcast()
+				}()
+			}
+		}
+	}()
+	go func() { // "wake" goroutine
+		defer l.Lock().Unlock()
+		for {
+
+			for !stopped && (active >= concurrency || pending.Len() == 0) {
+				pendingCond.Wait()
+			}
+			if stopped {
+				return
+			}
+			if pending.Len() <= 0 {
+				return
+			}
+			active++
+			next := heap.Pop(pending).(*stepQueueHeapItem).req
+			delete(queueItems, next.ident)
+
+			next.wakeup <- func() {
+				defer l.Lock().Unlock()
+				active--
+				pendingCond.Broadcast()
+			}
+		}
+	}()
+
+	done = func() {
+		close(q.stop)
+	}
+	return done
+}
+
+type StepCompletedFunc func()
+
+func (q *stepQueue) sendAndWaitForWakeup(ident interface{}, targetDate time.Time) StepCompletedFunc {
+	req := stepQueueRec{
+		ident,
+		targetDate,
+		make(chan StepCompletedFunc),
+	}
+	q.reqs <- req
+	return <-req.wakeup
+}
+
+// Wait for the ident with targetDate to be selected to run.
+func (q *stepQueue) WaitReady(ctx context.Context, ident interface{}, targetDate time.Time) StepCompletedFunc {
+	defer trace.WithSpanFromStackUpdateCtx(&ctx)()
+	if targetDate.IsZero() {
+		panic("targetDate of zero is reserved for marking Done")
+	}
+	return q.sendAndWaitForWakeup(ident, targetDate)
+}