stop scheduler from advancing through empty buckets without accept from ratelimiter

Author: Jim Minter

pkg/image/controller/scheduled_image_controller_test.go

@@ -57,8 +57,9 @@ func TestScheduledImport(t *testing.T) {
 	}
 
 	// encountering a not found error for image streams should drop the stream
-	sched.scheduler.RunOnce() // we need to run it twice since we have 2 buckets
-	sched.scheduler.RunOnce()
+	for i := 0; i < 3; i++ { // loop all the buckets (2 + the additional internal one)
+		sched.scheduler.RunOnce()
+	}
 	if sched.scheduler.Len() != 0 {
 		t.Fatalf("should have removed item in scheduler: %#v", sched.scheduler)
 	}
@@ -72,8 +73,9 @@ func TestScheduledImport(t *testing.T) {
 	isInformer.Informer().GetIndexer().Add(stream)
 
 	// run a background import
-	sched.scheduler.RunOnce() // we need to run it twice since we have 2 buckets
-	sched.scheduler.RunOnce()
+	for i := 0; i < 3; i++ { // loop all the buckets (2 + the additional internal one)
+		sched.scheduler.RunOnce()
+	}
 	if sched.scheduler.Len() != 1 {
 		t.Fatalf("should have left item in scheduler: %#v", sched.scheduler)
 	}
@@ -85,8 +87,9 @@ func TestScheduledImport(t *testing.T) {
 	sched.enabled = false
 	fake.ClearActions()
 
-	sched.scheduler.RunOnce() // we need to run it twice since we have 2 buckets
-	sched.scheduler.RunOnce()
+	for i := 0; i < 3; i++ { // loop all the buckets (2 + the additional internal one)
+		sched.scheduler.RunOnce()
+	}
 	if sched.scheduler.Len() != 0 {
 		t.Fatalf("should have removed item from scheduler: %#v", sched.scheduler)
 	}

pkg/image/controller/scheduler.go

@@ -69,16 +69,12 @@ func (s *Scheduler) next() (interface{}, interface{}, bool) {
 	defer s.mu.Unlock()
 
 	last := s.buckets[s.position]
-	if len(last) == 0 {
-		s.position = s.at(1)
-		last = s.buckets[s.position]
-	}
-
 	for k, v := range last {
 		delete(last, k)
 		s.buckets[s.at(-1)][k] = v
 		return k, v, false
 	}
+	s.position = s.at(1)
 	return nil, nil, true
 }
 

pkg/image/controller/scheduler_test.go

@@ -1,9 +1,14 @@
 package controller
 
 import (
+	"container/heap"
+	"fmt"
 	"reflect"
+	"sync"
 	"testing"
+	"time"
 
+	"github.com/juju/ratelimit"
 	flowcontrol "k8s.io/client-go/util/flowcontrol"
 )
 
@@ -39,14 +44,11 @@ func TestScheduler(t *testing.T) {
 		t.Fatal("expected to find key in a bucket")
 	}
 
-	for i := 0; i < 10; i++ {
-		s.Delay("first")
-		if _, ok := s.buckets[(s.position-1+len(s.buckets))%len(s.buckets)]["first"]; !ok {
-			t.Fatal("key was not in the last bucket")
-		}
-	}
+	s.Delay("first")
 
-	s.RunOnce()
+	for i := 0; i < 2; i++ { // Delay shouldn't have put the item in the current bucket
+		s.RunOnce()
+	}
 	if len(keys) != 0 {
 		t.Fatal(keys)
 	}
@@ -126,3 +128,102 @@ func TestSchedulerRemove(t *testing.T) {
 		t.Fatal(s)
 	}
 }
+
+type int64Heap []int64
+
+var _ heap.Interface = &int64Heap{}
+
+func (h int64Heap) Len() int            { return len(h) }
+func (h int64Heap) Less(i, j int) bool  { return h[i] < h[j] }
+func (h int64Heap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
+func (h *int64Heap) Push(x interface{}) { *h = append(*h, x.(int64)) }
+func (h *int64Heap) Pop() interface{} {
+	x := (*h)[len(*h)-1]
+	*h = (*h)[:len(*h)-1]
+	return x
+}
+
+type wallClock struct{}
+
+var _ ratelimit.Clock = &wallClock{}
+
+func (c wallClock) Now() time.Time        { return time.Now() }
+func (c wallClock) Sleep(d time.Duration) { time.Sleep(d) }
+
+// fakeClock implements ratelimit.Clock.  Its time starts at the UNIX epoch.
+// When all known threads are in Sleep(), its time advances just enough to wake
+// the first thread (or threads) due to wake.
+type fakeClock struct {
+	c       sync.Cond
+	threads int
+	now     int64
+	wake    int64Heap
+}
+
+var _ ratelimit.Clock = &fakeClock{}
+
+func newFakeClock(threads int) ratelimit.Clock {
+	return &fakeClock{threads: threads, c: sync.Cond{L: &sync.Mutex{}}}
+}
+
+func (c *fakeClock) Now() time.Time {
+	c.c.L.Lock()
+	defer c.c.L.Unlock()
+	return time.Unix(0, c.now)
+}
+
+func (c *fakeClock) Sleep(d time.Duration) {
+	c.c.L.Lock()
+	defer c.c.L.Unlock()
+	wake := c.now + int64(d)
+	heap.Push(&c.wake, wake)
+	if len(c.wake) == c.threads {
+		// everyone is asleep, advance the clock.
+		c.now = heap.Pop(&c.wake).(int64)
+		for len(c.wake) > 0 && c.wake[0] == c.now {
+			// pop any additional threads waiting on the same time.
+			heap.Pop(&c.wake)
+		}
+		c.c.Broadcast()
+	}
+	for c.now < wake {
+		c.c.Wait()
+	}
+}
+
+func TestSchedulerSanity(t *testing.T) {
+	const (
+		buckets = 4
+		items   = 10
+	)
+
+	clock := newFakeClock(2) // 2 threads: us and the scheduler.
+
+	// 1 token per second => one bucket's worth of items should get scheduled
+	// per second.
+	limiter := flowcontrol.NewTokenBucketRateLimiterWithClock(1, 1, clock)
+
+	m := map[int]int{}
+	s := NewScheduler(buckets, limiter, func(key, value interface{}) {
+		fmt.Printf("%v: %v\n", clock.Now().UTC(), key)
+		m[key.(int)]++
+	})
+	for i := 0; i < items; i++ {
+		s.Add(i, nil)
+	}
+
+	go s.RunUntil(make(chan struct{}))
+
+	// run the clock just long enough to expect to have scheduled each item
+	// exactly twice.
+	clock.Sleep((2*buckets-1)*time.Second + 1)
+
+	expected := map[int]int{}
+	for i := 0; i < items; i++ {
+		expected[i] = 2
+	}
+
+	if !reflect.DeepEqual(m, expected) {
+		t.Errorf("m did not match expected: %#v\n", m)
+	}
+}