diff --git a/common/sched_queue.go b/common/sched_queue.go new file mode 100644 index 00000000..df6a74ad --- /dev/null +++ b/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) +} diff --git a/common/sched_queue_test.go b/common/sched_queue_test.go new file mode 100644 index 00000000..81375a79 --- /dev/null +++ b/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") +} diff --git a/nameserver/mocks_test.go b/nameserver/mocks_test.go index 1bec091d..d2cb9810 100644 --- a/nameserver/mocks_test.go +++ b/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)