├── LICENSE ├── README.md ├── config.go ├── config_test.go ├── operation.go ├── operation_test.go ├── priority.go ├── priority_test.go ├── scheduler.go └── scheduler_test.go /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2015 boljen 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Scheduler (Go) 2 | 3 | Package scheduler implements a scheduler for rate limited operations 4 | using a prioritized queue. 5 | 6 | ## Use Case 7 | 8 | This package is built to schedule operations against rate limited API's. 9 | More specifically it's meant for applications which need to perform both 10 | real-time operations as well as a hefty amount of background scraping. 11 | 12 | ## Documentation 13 | 14 | See [godoc](https://godoc.org/github.com/boljen/go-scheduler) for more information. 15 | 16 | ## License 17 | 18 | Released under the MIT license. 19 | -------------------------------------------------------------------------------- /config.go: -------------------------------------------------------------------------------- 1 | package scheduler 2 | 3 | // Config configures the Ratelimitter. 4 | type Config struct { 5 | // OPS stands for operations per second and is the amount of operations 6 | // that the scheduler should allow during the course of one second. 7 | OPS float32 8 | 9 | // Workers is the amount of goroutine workers that process operations. 10 | // If this is 0 then no worker goroutines will be used and operations will 11 | // be executed synchronously from within the main tick loop. 12 | Workers int 13 | 14 | // MaxQueueSize is the maximum size of the operations queue. 15 | // This maximum is always enforced, even when priority-specific 16 | // queue's have a higher maximum queue size. 17 | MaxQueueSize int 18 | 19 | // ExecutionBufferSize is the capacity of the buffered channel which 20 | // forwards operations to the various workers. 21 | // This should be as low as possible to keep the scheduler in sync with 22 | // the remote rate limit window as much as possible. 23 | ExecutionBufferSize int 24 | 25 | // Fallback is an (optional) operation that will be executed every time that 26 | // no other operations are available. It will be executed from within the 27 | // same loop that processes ticks even if there are workers available. 28 | // This is by design and allows using this hook to refill the operations 29 | // queue whenever it's empty. 30 | Fallback Operation 31 | 32 | // PriorityAutoInit sets whether priorities are automatically initialized. 33 | // When this is false, the Scheduler will return an error every time an 34 | // operation uses an uninitialized priority. 35 | PriorityAutoInit bool 36 | 37 | // PriorityDefaultCapacity indicates the default capacity of a priority. 38 | // This is only relevant when PriorityAutoInit is true. 39 | PriorityDefaultCapacity int 40 | } 41 | 42 | func (c Config) rate() float32 { 43 | if c.OPS <= 0 { 44 | return 1 45 | } 46 | return c.OPS 47 | } 48 | 49 | func (c Config) maxops() uint32 { 50 | if c.MaxQueueSize <= 0 { 51 | return ^uint32(0) 52 | } 53 | return uint32(c.MaxQueueSize) 54 | } 55 | 56 | func (c Config) opbuf() int { 57 | if c.ExecutionBufferSize <= 0 { 58 | return 1 59 | } 60 | return c.ExecutionBufferSize 61 | } 62 | -------------------------------------------------------------------------------- /config_test.go: -------------------------------------------------------------------------------- 1 | package scheduler 2 | 3 | import "testing" 4 | 5 | func TestConfigMaxops(t *testing.T) { 6 | cfg := Config{} 7 | if cfg.maxops() != 4294967295 { 8 | t.Fatal("wrong default max operations") 9 | } 10 | cfg.MaxQueueSize = 5 11 | if cfg.maxops() != 5 { 12 | t.Fatal("wrong maxops") 13 | } 14 | 15 | if cfg.rate() != 1 { 16 | t.Fatal("wrong default rate") 17 | } 18 | 19 | cfg.OPS = 4 20 | if cfg.rate() != 4 { 21 | t.Fatal("wrong rate") 22 | } 23 | } 24 | -------------------------------------------------------------------------------- /operation.go: -------------------------------------------------------------------------------- 1 | package scheduler 2 | 3 | // Operation is an operation that can be executed by the scheduler. 4 | type Operation interface { 5 | Execute() 6 | } 7 | 8 | // Closure turns a closure into the Operation interface. 9 | // It should do so with virtually no overhead. 10 | func Closure(fx func()) Operation { 11 | return operationClosure(fx) 12 | } 13 | 14 | type operationClosure func() 15 | 16 | func (f operationClosure) Execute() { 17 | f() 18 | } 19 | -------------------------------------------------------------------------------- /operation_test.go: -------------------------------------------------------------------------------- 1 | package scheduler 2 | 3 | import "testing" 4 | 5 | type testOp struct{ T int } 6 | 7 | func (testOp) Execute() {} 8 | 9 | func TestOperationClosure(t *testing.T) { 10 | ok := false 11 | tf := func() { 12 | ok = true 13 | } 14 | cl := Closure(tf) 15 | cl.Execute() 16 | if !ok { 17 | t.Fatal("operation failed") 18 | } 19 | } 20 | -------------------------------------------------------------------------------- /priority.go: -------------------------------------------------------------------------------- 1 | package scheduler 2 | 3 | // Priority indicates a specific priority. 4 | // The higher the value, the higher the priority. 5 | type Priority int 6 | 7 | // (TODO): Refactor weight to "p" 8 | 9 | // priorityMetadata stores metadata of a priority inside the Scheduler. 10 | type priorityMetadata struct { 11 | priority Priority 12 | oplist map[int]Operation 13 | 14 | maxops uint32 // Maximum amount of operations 15 | curops uint32 // Current amount of operations 16 | 17 | first int 18 | last int 19 | 20 | Minimum uint32 21 | MinimumCallback func(Priority) 22 | } 23 | 24 | func getMaxops(maxops int) uint32 { 25 | maxops32 := uint32(maxops) 26 | if maxops32 == 0 { 27 | maxops32-- 28 | } 29 | return maxops32 30 | } 31 | 32 | func newPriorityMetadata(p Priority, maxops int) *priorityMetadata { 33 | return &priorityMetadata{ 34 | priority: p, 35 | oplist: make(map[int]Operation), 36 | curops: 0, 37 | maxops: getMaxops(maxops), 38 | } 39 | } 40 | 41 | // AddOperation adds a new operation to the priority. 42 | // It might return ErrPriorityCapacity when the priority-specific queue is full. 43 | func (p *priorityMetadata) AddOperation(o Operation) error { 44 | if p.curops == p.maxops { 45 | return ErrPriorityCapacity 46 | } 47 | p.curops++ 48 | p.oplist[p.last] = o 49 | p.last++ 50 | return nil 51 | } 52 | 53 | // GetOperation returns the next operation of this priority. 54 | // If no operation is available, the returned bool will be false. 55 | func (p *priorityMetadata) GetOperation() (Operation, bool) { 56 | if p.last == p.first { 57 | return nil, false 58 | } 59 | o := p.oplist[p.first] 60 | delete(p.oplist, p.first) 61 | p.first++ 62 | p.curops-- 63 | if p.curops == p.Minimum && p.MinimumCallback != nil { 64 | p.MinimumCallback(p.priority) 65 | } 66 | return o, true 67 | } 68 | -------------------------------------------------------------------------------- /priority_test.go: -------------------------------------------------------------------------------- 1 | package scheduler 2 | 3 | import "testing" 4 | 5 | func TestNewPriorityMetadata(t *testing.T) { 6 | p := newPriorityMetadata(12, 50) 7 | if p.priority != 12 { 8 | t.Fatal("wrong weight") 9 | } 10 | if p.maxops != 50 { 11 | t.Fatal("wrong maximum operations") 12 | } 13 | p = newPriorityMetadata(5, 0) 14 | if p.maxops < 100000 { 15 | t.Fatal("wrong maximum operations") 16 | } 17 | } 18 | 19 | func TestPriorityOperations(t *testing.T) { 20 | o1 := &testOp{} 21 | o2 := &testOp{} 22 | o3 := &testOp{} 23 | o4 := &testOp{} 24 | 25 | p := newPriorityMetadata(1, 2) 26 | if _, ok := p.GetOperation(); ok { 27 | t.Fatal("should not be ok") 28 | } 29 | if err := p.AddOperation(o1); err != nil { 30 | t.Fatal(err) 31 | } 32 | if err := p.AddOperation(o2); err != nil { 33 | t.Fatal(err) 34 | } 35 | if err := p.AddOperation(o3); err != ErrPriorityCapacity { 36 | t.Fatal(err) 37 | } 38 | if op, ok := p.GetOperation(); !ok || op != o1 { 39 | t.Fatal("should return operation 1") 40 | } 41 | if err := p.AddOperation(o4); err != nil { 42 | t.Fatal(err) 43 | } 44 | if op, ok := p.GetOperation(); !ok || op != o2 { 45 | t.Fatal("should return operation 2") 46 | } 47 | if op, ok := p.GetOperation(); !ok || op != o4 { 48 | t.Fatal("should return operation 4") 49 | } 50 | if _, ok := p.GetOperation(); ok { 51 | t.Fatal("should not be ok") 52 | } 53 | } 54 | -------------------------------------------------------------------------------- /scheduler.go: -------------------------------------------------------------------------------- 1 | // Package scheduler implements a scheduler for rate limited operations 2 | // using a prioritized queue. 3 | // 4 | // Use case 5 | // 6 | // This package is built to schedule operations against rate limited API's. 7 | // More specifically it's meant for applications which need to perform both 8 | // real-time operations as well as a hefty amount of background scraping. 9 | // 10 | // Scheduler 11 | // 12 | // The scheduler attempts to streamline the execution of operations by using 13 | // a continuous ticker at the allowed operation rate. At each tick, exactly one 14 | // operation will be executed. The order in which operations are executed is 15 | // first-in-first-out, but higher-priority operations will be executed first. 16 | // When no operations are available, a fallback operation will be called. 17 | // 18 | // By evenly spreading out the operations, there's a reasonable guarantee that 19 | // the highest priority operation will be executed within deterministic time 20 | // of one rate interval. 21 | // 22 | // Workers 23 | // 24 | // A scheduler can be configured to use one or multiple workers. Workers are 25 | // simply goroutines that continuously take operations from a buffered channel 26 | // and execute them. 27 | // 28 | // The advantage of using workers is that the main tick loop of the scheduler 29 | // won't be blocked by any single operation. It will only block when the 30 | // buffered channel transferring operations to workers has been filled entirely, 31 | // which will be the case when all of the workers together can't process the 32 | // flood of operations. 33 | // 34 | // The size of the buffered channel can be configured and does have an effect 35 | // on the accuracy of the scheduler. It should be kept as small as possible. 36 | // The moment an operation is sent to the buffered channel, it will seize to 37 | // exist inside the scheduler. 38 | // 39 | // When configured to use 0 workers, the scheduler will execute all operations 40 | // synchronously from within the loop that processes ticks from the internal 41 | // ticker. The disadvantage here is that it will cause the main loop to block, 42 | // the advantage is that it won't execute any expensive context switching. 43 | // 44 | // Bursts 45 | // 46 | // The Scheduler has no built-in support for bursts of operations. 47 | // This is because the way the rate is calculated highly depends on the 48 | // service that is being used and requires awareness of the strategy used. 49 | // 50 | // Implementing burst behavior can be done by lowering the allocated rate of 51 | // operations of the scheduler and using a separate system to allocate those 52 | // additional operations. It can also be done by 'saving up' operations through 53 | // the Fallback operation. 54 | package scheduler 55 | 56 | // (TODO): Provide hooks for "queue entries above/below x" for both the 57 | // global scheduler as well as the various priorities. This should be done 58 | // by creating a new Counter struct that provides this functionality and 59 | // which can be used by both the Priorities as well as the Scheduler itself. 60 | // (TODO): Optionally make the Scheduler stand-by until it receives an operation. 61 | // (TODO): Make the scheduler use an implementation of the "Tickable" interface. 62 | // (TODO): Create exhaustive unit tests. 63 | 64 | import ( 65 | "errors" 66 | "sync" 67 | "time" 68 | ) 69 | 70 | // These are possible errors that can be returned by the Scheduler instance. 71 | var ( 72 | ErrInvalidPriority = errors.New("Scheduler: Priority is not initizlaized") 73 | ErrMaxCapacity = errors.New("Scheduler: Maximum Queue Capacity Exceeded") 74 | ErrPriorityCapacity = errors.New("Priority: Maximum Priority-Specific Queue Capacity Exceeded") 75 | ) 76 | 77 | func worker(ch chan Operation) { 78 | for { 79 | op, more := <-ch 80 | if !more { 81 | break 82 | } 83 | op.Execute() 84 | } 85 | } 86 | 87 | // Scheduler schedules operations against a specific rate limit. 88 | type Scheduler struct { 89 | pause time.Time // The time until the scheduler must pause. 90 | usingWorkers bool // Whether separate goroutine workers are used. 91 | opqueue chan Operation // Queue of pending operations for the workers. 92 | fallback Operation // Fallback operation in case no operations are available. 93 | stop chan bool // Used to stop the ticker goroutine. 94 | ticker *time.Ticker // The internal ticker. 95 | 96 | pai bool // Priority Auto Initialization 97 | pdc int // Priority default capacity 98 | 99 | mu *sync.Mutex // Mutex 100 | pl map[Priority]*priorityMetadata // Mapped priority list. 101 | opl []*priorityMetadata // Ordered priority list. 102 | curops uint32 // total operations inside the scheduler queue. 103 | maxops uint32 // max is the maximum amount of operations that can be in the scheduler. 104 | } 105 | 106 | // New creates a newly initialized Scheduler instance. 107 | func New(c Config) *Scheduler { 108 | s := &Scheduler{ 109 | mu: new(sync.Mutex), 110 | pl: make(map[Priority]*priorityMetadata, 5), 111 | opl: make([]*priorityMetadata, 0, 5), 112 | pai: c.PriorityAutoInit, 113 | pdc: c.PriorityDefaultCapacity, 114 | maxops: c.maxops(), 115 | fallback: c.Fallback, 116 | stop: make(chan bool), 117 | } 118 | 119 | // When using workers we must initialize the workers and the operation queue. 120 | if c.Workers > 0 { 121 | s.opqueue = make(chan Operation, c.opbuf()) 122 | s.usingWorkers = true 123 | for i := 0; i < c.Workers; i++ { 124 | go worker(s.opqueue) 125 | } 126 | } 127 | 128 | // Start a new ticker based on the configured rate and start processing ticks. 129 | s.ticker = time.NewTicker(time.Duration(float32(time.Second) / c.rate())) 130 | go s.processTicks() 131 | 132 | return s 133 | } 134 | 135 | // processTicks processes ticks in a background goroutine. 136 | func (s *Scheduler) processTicks() { 137 | for { 138 | select { 139 | case t := <-s.ticker.C: 140 | if s.pause.Before(t) { 141 | s.execOp() 142 | } 143 | case <-s.stop: 144 | return 145 | } 146 | } 147 | } 148 | 149 | func (s *Scheduler) execOp() { 150 | o := s.getNextOp() 151 | if o == nil { 152 | if s.fallback == nil { 153 | return 154 | } 155 | s.fallback.Execute() 156 | return 157 | } 158 | 159 | if s.usingWorkers { 160 | s.opqueue <- o 161 | } else { 162 | o.Execute() 163 | } 164 | } 165 | 166 | // getOperation removes and returns the next pending operation. 167 | func (s *Scheduler) getNextOp() Operation { 168 | s.mu.Lock() 169 | defer s.mu.Unlock() 170 | for i := 0; i < len(s.opl); i++ { 171 | op, ok := s.opl[i].GetOperation() 172 | if ok { 173 | s.curops-- 174 | return op 175 | } 176 | } 177 | return nil 178 | } 179 | 180 | // InitPriority initializes a new priority and specifies the maximum 181 | // operation queue for the specific priority. If maxops equals 0, no 182 | // priority-specific limit will be applied. 183 | func (s *Scheduler) InitPriority(p Priority, maxops int) { 184 | s.mu.Lock() 185 | s.initPriority(p, maxops) 186 | s.mu.Unlock() 187 | } 188 | 189 | func (s *Scheduler) initPriority(p Priority, maxops int) { 190 | // If the priority already exists, simply overwrite the maxops. 191 | // Make sure to lock the mutex to avoid any race-conditions. 192 | if pr, ok := s.pl[p]; ok { 193 | pr.maxops = getMaxops(maxops) 194 | return 195 | } 196 | 197 | pm := newPriorityMetadata(p, maxops) 198 | s.pl[p] = pm 199 | 200 | // Reorder the ordered priority list slice from back to front. 201 | // This is most likely the most efficient algorithm. 202 | s.opl = append(s.opl, pm) 203 | for i := len(s.opl) - 1; i > 0; i-- { 204 | if s.opl[i].priority < s.opl[i-1].priority { 205 | s.opl[i] = s.opl[i-1] 206 | s.opl[i-1] = pm 207 | } 208 | } 209 | } 210 | 211 | // Add adds a new operation to the scheduler. 212 | // The priority must be initialized unless automated initialization is enabled. 213 | func (s *Scheduler) Add(p Priority, o Operation) error { 214 | s.mu.Lock() 215 | defer s.mu.Unlock() 216 | 217 | if s.curops >= s.maxops { 218 | return ErrMaxCapacity 219 | } 220 | 221 | pm, err := s.getPriorityMetadata(p) 222 | if err != nil { 223 | return err 224 | } 225 | 226 | if err := pm.AddOperation(o); err != nil { 227 | return err 228 | } 229 | 230 | s.curops++ 231 | return nil 232 | } 233 | 234 | // SetMinimumCallback sets a callback that will be executed each time 235 | // the amount of registered operations for a specific priority reaches 236 | // the specified minimum. Only one callback per priority can be set. 237 | // This will fail when the priority is not initialized and automated 238 | // initialization is disabled. 239 | func (s *Scheduler) SetMinimumCallback(p Priority, minimum int, cb func(Priority)) error { 240 | pm, err := s.getPriorityMetadata(p) 241 | if err != nil { 242 | return err 243 | } 244 | 245 | pm.Minimum = uint32(minimum) 246 | pm.MinimumCallback = cb 247 | if pm.Minimum >= pm.curops { 248 | pm.MinimumCallback(pm.priority) 249 | } 250 | return nil 251 | } 252 | 253 | func (s *Scheduler) getPriorityMetadata(p Priority) (*priorityMetadata, error) { 254 | pm, ok := s.pl[p] 255 | if !ok { 256 | if !s.pai { 257 | return nil, ErrInvalidPriority 258 | } 259 | s.initPriority(p, s.pdc) 260 | return s.pl[p], nil 261 | } 262 | return pm, nil 263 | } 264 | 265 | // Pause pauses the scheduler for the specified duration. 266 | // Use this when the rate limit has been exceeded and when you know 267 | // the moment where the next window will become active. 268 | func (s *Scheduler) Pause(d time.Duration) { 269 | s.pause = time.Now().Add(d) 270 | } 271 | 272 | // Stop stops the scheduler and all of it's background processes. 273 | // This might lead to skipping operations that are currently queued. 274 | // The operation is final. The scheduler shouldn't be used after 275 | // Stop has been called. 276 | func (s *Scheduler) Stop() { 277 | s.ticker.Stop() 278 | close(s.opqueue) 279 | s.stop <- true 280 | } 281 | -------------------------------------------------------------------------------- /scheduler_test.go: -------------------------------------------------------------------------------- 1 | package scheduler 2 | 3 | import ( 4 | "testing" 5 | "time" 6 | ) 7 | 8 | var defaultConfig = Config{} 9 | 10 | func TestWorker(t *testing.T) { 11 | ch := make(chan Operation) 12 | 13 | go func() { 14 | ch <- &testOp{} 15 | close(ch) 16 | }() 17 | 18 | worker(ch) 19 | } 20 | 21 | func TestNew(t *testing.T) { 22 | New(Config{}) 23 | rl := New(Config{ 24 | Workers: 10, 25 | }) 26 | time.Sleep(time.Second) 27 | 28 | rl.Stop() 29 | } 30 | 31 | func TestScheduler_Add(t *testing.T) { 32 | o := &testOp{} 33 | rl := New(Config{ 34 | Workers: 4, 35 | MaxQueueSize: 2, 36 | }) 37 | 38 | if err := rl.Add(1, o); err != ErrInvalidPriority { 39 | t.Fatal("wrong priority") 40 | } 41 | 42 | rl.InitPriority(1, 1) 43 | if err := rl.Add(1, o); err != nil { 44 | t.Fatal(err) 45 | } 46 | 47 | if err := rl.Add(1, o); err != ErrPriorityCapacity { 48 | t.Fatal("expected ErrPriorityCapacity") 49 | } 50 | 51 | rl.InitPriority(2, 0) 52 | if err := rl.Add(2, o); err != nil { 53 | t.Fatal(err) 54 | } 55 | 56 | if err := rl.Add(2, o); err != ErrMaxCapacity { 57 | t.Fatal("expected ErrMaxCapacity, got", err) 58 | } 59 | 60 | rl.Pause(time.Second) 61 | 62 | time.Sleep(time.Second) 63 | 64 | } 65 | 66 | func TestScheduler_getPriorityMetadata(t *testing.T) { 67 | rl := New(Config{}) 68 | rl.InitPriority(10, 100) 69 | if _, err := rl.getPriorityMetadata(10); err != nil { 70 | t.Fatal(err) 71 | } 72 | if _, err := rl.getPriorityMetadata(100); err == nil { 73 | t.Fatal("must return an error") 74 | } 75 | rl.pai = true 76 | rl.pdc = 1234 77 | if pm, err := rl.getPriorityMetadata(100); err != nil { 78 | t.Fatal(err) 79 | } else if pm.maxops != 1234 { 80 | t.Fatal("wrong oplist capacity") 81 | } 82 | } 83 | 84 | func TestScheduler_InitPriority(t *testing.T) { 85 | rl := New(Config{}) 86 | rl.InitPriority(10, 100) 87 | if rl.opl[0].priority != 10 { 88 | t.Fatal("wrong opl entry") 89 | } 90 | 91 | rl.InitPriority(5, 100) 92 | if rl.opl[0].priority != 5 || rl.opl[1].priority != 10 { 93 | t.Fatal("wrong opl entry") 94 | } 95 | 96 | rl.InitPriority(5, 50) 97 | if rl.opl[0].maxops != 50 { 98 | t.Fatal("wrong maxops entry") 99 | } 100 | } 101 | 102 | func TestSchedulerSetMinimumCallback(t *testing.T) { 103 | rl := New(Config{}) 104 | rl.InitPriority(10, 100) 105 | if err := rl.SetMinimumCallback(Priority(1), 5, nil); err != ErrInvalidPriority { 106 | t.Fatal("expected invalid priority error") 107 | } 108 | 109 | done := false 110 | if err := rl.SetMinimumCallback(10, 5, func(p Priority) { 111 | done = true 112 | }); err != nil { 113 | t.Fatal("unexpected error", err) 114 | } 115 | if !done { 116 | t.Fatal("should have launched the minimum callback ") 117 | } 118 | } 119 | --------------------------------------------------------------------------------