feat/throttled-dispatcher

throttled_dispatcher.go

@@ -0,0 +1,70 @@
+package artifex
+
+import "time"
+
+// ThrottleDispatcher wraps a Dispatcher with rate-limiting behavior,
+// allowing at most maxJobs to be dispatched per interval.
+type ThrottleDispatcher struct {
+	*Dispatcher
+	maxJobs  int
+	interval time.Duration
+	tokens   chan struct{}
+	quit     chan bool
+}
+
+// NewThrottledDispatcher creates a new dispatcher that limits job execution
+// to at most maxJobs per interval (e.g., 5 jobs per time.Minute).
+func NewThrottledDispatcher(maxWorkers int, maxQueue int, maxJobs int, interval time.Duration) *ThrottleDispatcher {
+	tokens := make(chan struct{}, maxJobs)
+	// Pre-fill the token bucket
+	for i := 0; i < maxJobs; i++ {
+		tokens <- struct{}{}
+	}
+	return &ThrottleDispatcher{
+		Dispatcher: NewDispatcher(maxWorkers, maxQueue),
+		maxJobs:    maxJobs,
+		interval:   interval,
+		tokens:     tokens,
+		quit:       make(chan bool),
+	}
+}
+
+// Start begins the throttled dispatcher and starts the token replenishment ticker.
+func (td *ThrottleDispatcher) Start() {
+	td.Dispatcher.Start()
+	go td.replenish()
+}
+
+// Stop halts the throttled dispatcher and stops token replenishment.
+func (td *ThrottleDispatcher) Stop() {
+	td.quit <- true
+	td.Dispatcher.Stop()
+}
+
+// replenish refills the token bucket at the configured rate.
+// Tokens are added in batches of maxJobs every interval.
+func (td *ThrottleDispatcher) replenish() {
+	ticker := time.NewTicker(td.interval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			for i := 0; i < td.maxJobs; i++ {
+				select {
+				case td.tokens <- struct{}{}:
+				default:
+					// Bucket is full; discard extra tokens
+				}
+			}
+		case <-td.quit:
+			return
+		}
+	}
+}
+
+// Dispatch waits for a throttle token before queuing the job,
+// ensuring no more than maxJobs are dispatched per interval.
+func (td *ThrottleDispatcher) Dispatch(run func()) error {
+	<-td.tokens // Block until a token is available
+	return td.Dispatcher.Dispatch(run)
+}

throttled_dispatcher_test.go

@@ -0,0 +1,39 @@
+package artifex
+
+import (
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestThrottledDispatcher_LimitsRate(t *testing.T) {
+	// Allow 3 jobs per 500ms window
+	td := NewThrottledDispatcher(5, 100, 3, 500*time.Millisecond)
+	td.Start()
+	defer td.Stop()
+
+	var count int32
+
+	// Dispatch 6 jobs; first 3 should run immediately, next 3 after ~500ms
+	for i := 0; i < 6; i++ {
+		go func() {
+			td.Dispatch(func() {
+				atomic.AddInt32(&count, 1)
+			})
+		}()
+	}
+
+	// After a short wait, only first batch should have been dispatched
+	time.Sleep(100 * time.Millisecond)
+	firstBatch := atomic.LoadInt32(&count)
+	if firstBatch > 3 {
+		t.Errorf("expected at most 3 jobs in first window, got %d", firstBatch)
+	}
+
+	// Wait for second window to open
+	time.Sleep(600 * time.Millisecond)
+	total := atomic.LoadInt32(&count)
+	if total != 6 {
+		t.Errorf("expected 6 total jobs, got %d", total)
+	}
+}