New scheduled queue, for scheduling calls to callables

common/sched_queue.go

@@ -0,0 +1,145 @@
+package common
+
+import (
+	"container/heap"
+	"math"
+	"time"
+
+	"github.com/benbjohnson/clock"
+	"runtime"
+	"sync/atomic"
+)
+
+// A callable is a function that is called and return an (optional) next scheduled time
+type callable func() time.Time
+
+type schedCall struct {
+	t time.Time
+	c callable
+}
+
+func (sc *schedCall) FromNow(now time.Time) time.Duration {
+	if sc.t.After(now) {
+		return time.Duration(sc.t.Sub(now).Nanoseconds())
+	}
+	return time.Duration(0)
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+// A min-heap of scheduled callables
+type schedCallsHeap []*schedCall
+
+func (ch *schedCallsHeap) Len() int           { return len(*ch) }
+func (ch *schedCallsHeap) Less(i, j int) bool { return (*ch)[i].t.Before((*ch)[j].t) }
+func (ch *schedCallsHeap) Swap(i, j int) {
+	(*ch)[i], (*ch)[j] = (*ch)[j], (*ch)[i]
+}
+
+func (ch *schedCallsHeap) Push(x interface{}) {
+	// Push and Pop use pointer receivers because they modify the slice's length,
+	// not just its contents.
+	entry := x.(*schedCall)
+	*ch = append(*ch, entry)
+}
+
+func (ch *schedCallsHeap) Pop() interface{} {
+	old := *ch
+	n := len(old)
+	item := old[n-1]
+	*ch = old[0 : n-1]
+	return item
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+// SchedQueue is queue of scheduled callables
+type SchedQueue struct {
+	clock      clock.Clock
+	callablesH schedCallsHeap
+	schedChan  chan *schedCall
+	closeChan  chan bool
+	counter    uint64 // number of calls invoked so far (used for stats). Note: it will wrap.
+}
+
+// NewSchedQueue creates a new scheduled queue
+func NewSchedQueue(clock clock.Clock) *SchedQueue {
+	cq := SchedQueue{
+		clock:     clock,
+		schedChan: make(chan *schedCall),
+		closeChan: make(chan bool),
+	}
+	heap.Init(&cq.callablesH)
+	return &cq
+}
+
+// Start starts the scheduled queue
+func (cq *SchedQueue) Start() {
+	go func() {
+		defer func() { Debug.Printf("[sched-q] Quitting (%d calls pending)", len(cq.callablesH)) }()
+
+		var now time.Time
+		timer := cq.clock.Timer(time.Duration(math.MaxInt64))
+
+		for {
+			// Wait until an insertion, until the next callable or until we are Stop()ed
+			select {
+			case ce := <-cq.schedChan:
+				heap.Push(&cq.callablesH, ce)
+				now = cq.clock.Now()
+				Debug.Printf("[sched-q] Call scheduled [len:%d]", len(cq.callablesH))
+			case <-cq.closeChan:
+				return
+			case now = <-timer.C:
+			}
+
+			timer.Stop()
+			durationUntilNext := cq.durationUntilNext(now)
+			for durationUntilNext == 0 {
+				sched := heap.Pop(&cq.callablesH).(*schedCall)
+
+				atomic.AddUint64(&cq.counter, 1)
+				schedNextTime := sched.c()
+
+				if !schedNextTime.IsZero() {
+					Debug.Printf("[sched-q] Re-scheduled at %s", schedNextTime)
+					sched.t = schedNextTime
+					heap.Push(&cq.callablesH, sched) // TODO: use a Fix() instead of Pop()&Push()
+				}
+
+				now = cq.clock.Now()
+				durationUntilNext = cq.durationUntilNext(now)
+			}
+
+			timer = cq.clock.Timer(durationUntilNext)
+		}
+	}()
+}
+
+// Stop stops the scheduled queue
+func (cq *SchedQueue) Stop() {
+	Debug.Printf("[sched-q] Stopping...")
+	cq.closeChan <- true
+}
+
+// Add schedules a call.
+// The callable should not modify the scheduled queue in any way.
+func (cq *SchedQueue) Add(c callable, t time.Time) {
+	Debug.Printf("[sched-q] Adding call at %s", t)
+	ce := schedCall{c: c, t: t}
+	cq.schedChan <- &ce
+}
+
+// Counter returns the number of calls invoked in the queued
+// Note: the result will wrap over time.
+func (cq *SchedQueue) Count() uint64 {
+	return atomic.LoadUint64(&cq.counter)
+}
+
+func (cq *SchedQueue) durationUntilNext(now time.Time) time.Duration {
+	if len(cq.callablesH) > 0 {
+		firstSched := cq.callablesH[0]
+		return firstSched.FromNow(now)
+	}
+	return time.Duration(math.MaxInt64)
+}

common/sched_queue_test.go

@@ -0,0 +1,87 @@
+package common
+
+import (
+	"testing"
+	"time"
+
+	"github.com/benbjohnson/clock"
+	wt "github.com/weaveworks/weave/testing"
+	"log"
+)
+
+// Ensure we can add new calls while forwarding the clock
+func TestSchedCallsBasic(t *testing.T) {
+	InitDefaultLogging(testing.Verbose())
+	Info.Println("TestSchedCallsBasic starting")
+
+	const testSecs = 1000
+	clk := clock.NewMock()
+	schedQueue := NewSchedQueue(clk)
+	schedQueue.Start()
+	defer schedQueue.Stop()
+
+	c := func() time.Time {
+		return clk.Now().Add(time.Second)
+	}
+
+	schedQueue.Add(c, clk.Now().Add(time.Second))
+	for i := 0; i < testSecs; i++ {
+		clk.Add(time.Second)
+	}
+
+	t.Logf("Now: %s - calls: %d", clk.Now(), schedQueue.Count())
+	wt.AssertEqualInt(t, (int)(schedQueue.Count()), testSecs, "Number of calls")
+}
+
+// Ensure we can create a 100 seconds gap in the middle of the time travel
+func TestSchedCallsGap(t *testing.T) {
+	InitDefaultLogging(testing.Verbose())
+	Info.Println("TestSchedCallsGap starting")
+
+	const testSecs = 1000
+	clk := clock.NewMock()
+	schedQueue := NewSchedQueue(clk)
+	schedQueue.Start()
+	defer schedQueue.Stop()
+
+	c2 := func() time.Time {
+		if schedQueue.Count() == testSecs/2 {
+			return clk.Now().Add(time.Duration(100) * time.Second)
+		}
+		return clk.Now().Add(time.Second)
+	}
+
+	schedQueue.Add(c2, clk.Now().Add(time.Second))
+	for i := 0; i < testSecs; i++ {
+		clk.Add(time.Second)
+	}
+
+	t.Logf("Now: %s - calls: %d", clk.Now(), schedQueue.Count())
+	wt.AssertEqualInt(t, (int)(schedQueue.Count()), testSecs-100+1, "Number of calls")
+}
+
+func TestSchedCallsStop(t *testing.T) {
+	InitDefaultLogging(testing.Verbose())
+	Info.Println("TestSchedCallsStop starting")
+
+	const testSecs = 1000
+	clk := clock.NewMock()
+	schedQueue := NewSchedQueue(clk)
+	schedQueue.Start()
+	defer schedQueue.Stop()
+
+	c2 := func() time.Time {
+		if schedQueue.Count() == testSecs/2 {
+			return time.Time{}
+		}
+		return clk.Now().Add(time.Second)
+	}
+
+	schedQueue.Add(c2, clk.Now().Add(time.Second))
+	for i := 0; i < testSecs; i++ {
+		clk.Add(time.Second)
+	}
+
+	t.Logf("Now: %s - calls: %d", clk.Now(), schedQueue.Count())
+	wt.AssertEqualInt(t, (int)(schedQueue.Count()), testSecs/2, "Number of calls")
+}

nameserver/mocks_test.go

@@ -432,6 +432,9 @@ func newMockedClock() *mockedClock {
 	return &mockedClock{clock.NewMock()}
 }
 
+// Note: moving the clock forward does not take into account the data that could
+//       be waiting in channels. So your code should not depend on the channels
+//       for creating new timers. Otherwise, time travelling will not be reliable...
 func (clk *mockedClock) Forward(secs int) {
 	Debug.Printf(">>>>>>> Moving clock forward %d seconds - Time traveling >>>>>>>", secs)
 	clk.Add(time.Duration(secs) * time.Second)