mirror of
https://github.com/navidrome/navidrome.git
synced 2026-08-31 07:30:32 +00:00
fix(plugins): harden TaskQueue host service with validation and safety improvements
Add input validation (queue name length, payload size limits), extract status string constants to eliminate raw SQL literals, make CreateQueue idempotent via upsert for crash recovery, fix RetentionMs default check for negative values, cap exponential backoff at 1 hour to prevent overflow, and replace manual mutex-based delay enforcement with rate.Limiter from golang.org/x/time/rate for correct concurrent worker serialization.
This commit is contained in:
Deluan
Deluan Quintão
parent
a7545689ee
commit
dd12d34fe4
@ -18,6 +18,7 @@ import (
|
||||
"github.com/navidrome/navidrome/model/id"
|
||||
"github.com/navidrome/navidrome/plugins/capabilities"
|
||||
"github.com/navidrome/navidrome/plugins/host"
|
||||
"golang.org/x/time/rate"
|
||||
)
|
||||
|
||||
const (
|
||||
@ -26,11 +27,22 @@ const (
|
||||
defaultRetentionMs int64 = 3_600_000 // 1 hour
|
||||
minRetentionMs int64 = 60_000 // 1 minute
|
||||
maxRetentionMs int64 = 604_800_000 // 1 week
|
||||
maxQueueNameLength = 128
|
||||
maxPayloadSize = 1 * 1024 * 1024 // 1MB
|
||||
maxBackoffMs = 3_600_000 // 1 hour
|
||||
cleanupInterval = 5 * time.Minute
|
||||
pollInterval = 5 * time.Second
|
||||
shutdownTimeout = 10 * time.Second
|
||||
)
|
||||
|
||||
const (
|
||||
taskStatusPending = "pending"
|
||||
taskStatusRunning = "running"
|
||||
taskStatusCompleted = "completed"
|
||||
taskStatusFailed = "failed"
|
||||
taskStatusCancelled = "cancelled"
|
||||
)
|
||||
|
||||
// CapabilityTaskWorker indicates the plugin can receive task execution callbacks.
|
||||
// Detected when the plugin exports the task worker callback function.
|
||||
const CapabilityTaskWorker Capability = "TaskWorker"
|
||||
@ -43,10 +55,9 @@ func init() {
|
||||
|
||||
// queueState holds in-memory state for a single task queue.
|
||||
type queueState struct {
|
||||
config host.QueueConfig
|
||||
signal chan struct{}
|
||||
lastDispatchAt time.Time
|
||||
mu sync.Mutex
|
||||
config host.QueueConfig
|
||||
signal chan struct{}
|
||||
limiter *rate.Limiter // rate limiter for delay enforcement between dispatches
|
||||
}
|
||||
|
||||
// taskQueueServiceImpl implements host.TaskQueueService with SQLite persistence
|
||||
@ -141,6 +152,14 @@ func createTaskQueueSchema(db *sql.DB) error {
|
||||
|
||||
// CreateQueue creates a named task queue with the given configuration.
|
||||
func (s *taskQueueServiceImpl) CreateQueue(ctx context.Context, name string, config host.QueueConfig) error {
|
||||
// Validate queue name
|
||||
if len(name) == 0 {
|
||||
return fmt.Errorf("queue name cannot be empty")
|
||||
}
|
||||
if len(name) > maxQueueNameLength {
|
||||
return fmt.Errorf("queue name exceeds maximum length of %d bytes", maxQueueNameLength)
|
||||
}
|
||||
|
||||
// Apply defaults
|
||||
if config.Concurrency <= 0 {
|
||||
config.Concurrency = defaultConcurrency
|
||||
@ -148,7 +167,7 @@ func (s *taskQueueServiceImpl) CreateQueue(ctx context.Context, name string, con
|
||||
if config.BackoffMs <= 0 {
|
||||
config.BackoffMs = defaultBackoffMs
|
||||
}
|
||||
if config.RetentionMs == 0 {
|
||||
if config.RetentionMs <= 0 {
|
||||
config.RetentionMs = defaultRetentionMs
|
||||
}
|
||||
|
||||
@ -187,10 +206,16 @@ func (s *taskQueueServiceImpl) CreateQueue(ctx context.Context, name string, con
|
||||
return fmt.Errorf("queue %q already exists", name)
|
||||
}
|
||||
|
||||
// Insert into queues table
|
||||
// Upsert into queues table (idempotent across restarts)
|
||||
_, err := s.db.ExecContext(ctx, `
|
||||
INSERT INTO queues (name, concurrency, max_retries, backoff_ms, delay_ms, retention_ms)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(name) DO UPDATE SET
|
||||
concurrency = excluded.concurrency,
|
||||
max_retries = excluded.max_retries,
|
||||
backoff_ms = excluded.backoff_ms,
|
||||
delay_ms = excluded.delay_ms,
|
||||
retention_ms = excluded.retention_ms
|
||||
`, name, config.Concurrency, config.MaxRetries, config.BackoffMs, config.DelayMs, config.RetentionMs)
|
||||
if err != nil {
|
||||
return fmt.Errorf("creating queue: %w", err)
|
||||
@ -199,8 +224,8 @@ func (s *taskQueueServiceImpl) CreateQueue(ctx context.Context, name string, con
|
||||
// Reset stale running tasks from previous crash
|
||||
now := time.Now().UnixMilli()
|
||||
_, err = s.db.ExecContext(ctx, `
|
||||
UPDATE tasks SET status = 'pending', updated_at = ? WHERE queue_name = ? AND status = 'running'
|
||||
`, now, name)
|
||||
UPDATE tasks SET status = ?, updated_at = ? WHERE queue_name = ? AND status = ?
|
||||
`, taskStatusPending, now, name, taskStatusRunning)
|
||||
if err != nil {
|
||||
return fmt.Errorf("resetting stale tasks: %w", err)
|
||||
}
|
||||
@ -210,6 +235,11 @@ func (s *taskQueueServiceImpl) CreateQueue(ctx context.Context, name string, con
|
||||
config: config,
|
||||
signal: make(chan struct{}, 1),
|
||||
}
|
||||
if config.DelayMs > 0 {
|
||||
// Rate limit dispatches to enforce delay between tasks.
|
||||
// Burst of 1 allows one immediate dispatch, then enforces the delay interval.
|
||||
qs.limiter = rate.NewLimiter(rate.Every(time.Duration(config.DelayMs)*time.Millisecond), 1)
|
||||
}
|
||||
s.queues[name] = qs
|
||||
|
||||
// Start worker goroutines
|
||||
@ -234,13 +264,17 @@ func (s *taskQueueServiceImpl) Enqueue(ctx context.Context, queueName string, pa
|
||||
return "", fmt.Errorf("queue %q does not exist", queueName)
|
||||
}
|
||||
|
||||
if len(payload) > maxPayloadSize {
|
||||
return "", fmt.Errorf("payload size %d exceeds maximum of %d bytes", len(payload), maxPayloadSize)
|
||||
}
|
||||
|
||||
taskID := id.NewRandom()
|
||||
now := time.Now().UnixMilli()
|
||||
|
||||
_, err := s.db.ExecContext(ctx, `
|
||||
INSERT INTO tasks (id, queue_name, payload, status, attempt, max_retries, next_run_at, created_at, updated_at)
|
||||
VALUES (?, ?, ?, 'pending', 0, ?, ?, ?, ?)
|
||||
`, taskID, queueName, payload, qs.config.MaxRetries, now, now, now)
|
||||
VALUES (?, ?, ?, ?, 0, ?, ?, ?, ?)
|
||||
`, taskID, queueName, payload, taskStatusPending, qs.config.MaxRetries, now, now, now)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("enqueuing task: %w", err)
|
||||
}
|
||||
@ -272,8 +306,8 @@ func (s *taskQueueServiceImpl) GetTaskStatus(ctx context.Context, taskID string)
|
||||
func (s *taskQueueServiceImpl) CancelTask(ctx context.Context, taskID string) error {
|
||||
now := time.Now().UnixMilli()
|
||||
result, err := s.db.ExecContext(ctx, `
|
||||
UPDATE tasks SET status = 'cancelled', updated_at = ? WHERE id = ? AND status = 'pending'
|
||||
`, now, taskID)
|
||||
UPDATE tasks SET status = ?, updated_at = ? WHERE id = ? AND status = ?
|
||||
`, taskStatusCancelled, now, taskID, taskStatusPending)
|
||||
if err != nil {
|
||||
return fmt.Errorf("cancelling task: %w", err)
|
||||
}
|
||||
@ -344,14 +378,14 @@ func (s *taskQueueServiceImpl) processTask(queueName string, qs *queueState) boo
|
||||
var attempt int32
|
||||
var maxRetries int32
|
||||
err := s.db.QueryRowContext(s.ctx, `
|
||||
UPDATE tasks SET status = 'running', attempt = attempt + 1, updated_at = ?
|
||||
UPDATE tasks SET status = ?, attempt = attempt + 1, updated_at = ?
|
||||
WHERE id = (
|
||||
SELECT id FROM tasks
|
||||
WHERE queue_name = ? AND status = 'pending' AND next_run_at <= ?
|
||||
WHERE queue_name = ? AND status = ? AND next_run_at <= ?
|
||||
ORDER BY next_run_at, created_at LIMIT 1
|
||||
)
|
||||
RETURNING id, payload, attempt, max_retries
|
||||
`, now, queueName, now).Scan(&taskID, &payload, &attempt, &maxRetries)
|
||||
`, taskStatusRunning, now, queueName, taskStatusPending, now).Scan(&taskID, &payload, &attempt, &maxRetries)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return false
|
||||
}
|
||||
@ -360,27 +394,14 @@ func (s *taskQueueServiceImpl) processTask(queueName string, qs *queueState) boo
|
||||
return false
|
||||
}
|
||||
|
||||
// Enforce delay between task dispatches
|
||||
if qs.config.DelayMs > 0 {
|
||||
qs.mu.Lock()
|
||||
elapsed := time.Since(qs.lastDispatchAt)
|
||||
delay := time.Duration(qs.config.DelayMs) * time.Millisecond
|
||||
if elapsed < delay {
|
||||
waitTime := delay - elapsed
|
||||
qs.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-s.ctx.Done():
|
||||
// Put the task back to pending on shutdown
|
||||
s.revertTaskToPending(taskID)
|
||||
return false
|
||||
case <-time.After(waitTime):
|
||||
}
|
||||
|
||||
qs.mu.Lock()
|
||||
// Enforce delay between task dispatches using a rate limiter.
|
||||
// This is done after dequeue so that empty polls don't consume rate tokens.
|
||||
if qs.limiter != nil {
|
||||
if err := qs.limiter.Wait(s.ctx); err != nil {
|
||||
// Context cancelled during wait — revert task to pending for recovery
|
||||
s.revertTaskToPending(taskID)
|
||||
return false
|
||||
}
|
||||
qs.lastDispatchAt = time.Now()
|
||||
qs.mu.Unlock()
|
||||
}
|
||||
|
||||
// Invoke callback
|
||||
@ -396,7 +417,7 @@ func (s *taskQueueServiceImpl) processTask(queueName string, qs *queueState) boo
|
||||
now = time.Now().UnixMilli()
|
||||
if callbackErr == nil {
|
||||
// Success: mark as completed
|
||||
_, err = s.db.ExecContext(s.ctx, `UPDATE tasks SET status = 'completed', updated_at = ? WHERE id = ?`, now, taskID)
|
||||
_, err = s.db.ExecContext(s.ctx, `UPDATE tasks SET status = ?, updated_at = ? WHERE id = ?`, taskStatusCompleted, now, taskID)
|
||||
if err != nil {
|
||||
log.Error(s.ctx, "Failed to mark task as completed", "plugin", s.pluginName, "taskID", taskID, err)
|
||||
}
|
||||
@ -409,10 +430,13 @@ func (s *taskQueueServiceImpl) processTask(queueName string, qs *queueState) boo
|
||||
if attempt <= maxRetries {
|
||||
// Retry with exponential backoff: backoffMs * 2^(attempt-1)
|
||||
backoff := qs.config.BackoffMs * int64(math.Pow(2, float64(attempt-1)))
|
||||
if backoff < 0 || backoff > maxBackoffMs {
|
||||
backoff = maxBackoffMs
|
||||
}
|
||||
nextRunAt := now + backoff
|
||||
_, err = s.db.ExecContext(s.ctx, `
|
||||
UPDATE tasks SET status = 'pending', next_run_at = ?, updated_at = ? WHERE id = ?
|
||||
`, nextRunAt, now, taskID)
|
||||
UPDATE tasks SET status = ?, next_run_at = ?, updated_at = ? WHERE id = ?
|
||||
`, taskStatusPending, nextRunAt, now, taskID)
|
||||
if err != nil {
|
||||
log.Error(s.ctx, "Failed to reschedule task for retry", "plugin", s.pluginName, "taskID", taskID, err)
|
||||
}
|
||||
@ -428,7 +452,7 @@ func (s *taskQueueServiceImpl) processTask(queueName string, qs *queueState) boo
|
||||
})
|
||||
} else {
|
||||
// Exhausted retries: mark as failed
|
||||
_, err = s.db.ExecContext(s.ctx, `UPDATE tasks SET status = 'failed', updated_at = ? WHERE id = ?`, now, taskID)
|
||||
_, err = s.db.ExecContext(s.ctx, `UPDATE tasks SET status = ?, updated_at = ? WHERE id = ?`, taskStatusFailed, now, taskID)
|
||||
if err != nil {
|
||||
log.Error(s.ctx, "Failed to mark task as failed", "plugin", s.pluginName, "taskID", taskID, err)
|
||||
}
|
||||
@ -443,7 +467,7 @@ func (s *taskQueueServiceImpl) processTask(queueName string, qs *queueState) boo
|
||||
// counter (used during shutdown to ensure the interrupted attempt doesn't count).
|
||||
func (s *taskQueueServiceImpl) revertTaskToPending(taskID string) {
|
||||
now := time.Now().UnixMilli()
|
||||
_, err := s.db.Exec(`UPDATE tasks SET status = 'pending', attempt = MAX(attempt - 1, 0), updated_at = ? WHERE id = ? AND status = 'running'`, now, taskID)
|
||||
_, err := s.db.Exec(`UPDATE tasks SET status = ?, attempt = MAX(attempt - 1, 0), updated_at = ? WHERE id = ? AND status = ?`, taskStatusPending, now, taskID, taskStatusRunning)
|
||||
if err != nil {
|
||||
log.Error("Failed to revert task to pending", "plugin", s.pluginName, "taskID", taskID, err)
|
||||
}
|
||||
@ -508,8 +532,8 @@ func (s *taskQueueServiceImpl) runCleanup() {
|
||||
now := time.Now().UnixMilli()
|
||||
for name, qs := range queues {
|
||||
result, err := s.db.Exec(`
|
||||
DELETE FROM tasks WHERE queue_name = ? AND status IN ('completed', 'failed', 'cancelled') AND updated_at + ? < ?
|
||||
`, name, qs.config.RetentionMs, now)
|
||||
DELETE FROM tasks WHERE queue_name = ? AND status IN (?, ?, ?) AND updated_at + ? < ?
|
||||
`, name, taskStatusCompleted, taskStatusFailed, taskStatusCancelled, qs.config.RetentionMs, now)
|
||||
if err != nil {
|
||||
log.Error(s.ctx, "Failed to cleanup tasks", "plugin", s.pluginName, "queue", name, err)
|
||||
continue
|
||||
@ -541,7 +565,7 @@ func (s *taskQueueServiceImpl) Close() error {
|
||||
// Mark running tasks as pending for recovery on next startup
|
||||
if s.db != nil {
|
||||
now := time.Now().UnixMilli()
|
||||
_, err := s.db.Exec(`UPDATE tasks SET status = 'pending', updated_at = ? WHERE status = 'running'`, now)
|
||||
_, err := s.db.Exec(`UPDATE tasks SET status = ?, updated_at = ? WHERE status = ?`, taskStatusPending, now, taskStatusRunning)
|
||||
if err != nil {
|
||||
log.Error("Failed to reset running tasks on shutdown", "plugin", s.pluginName, err)
|
||||
}
|
||||
|
||||
@ -12,6 +12,9 @@ import (
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
@ -88,6 +91,30 @@ var _ = Describe("TaskQueueService", func() {
|
||||
})
|
||||
})
|
||||
|
||||
Describe("CreateQueue name validation", func() {
|
||||
It("rejects empty queue name", func() {
|
||||
err := service.CreateQueue(ctx, "", host.QueueConfig{})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("queue name cannot be empty"))
|
||||
})
|
||||
|
||||
It("rejects over-length queue name", func() {
|
||||
longName := strings.Repeat("a", maxQueueNameLength+1)
|
||||
err := service.CreateQueue(ctx, longName, host.QueueConfig{})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("exceeds maximum length"))
|
||||
})
|
||||
|
||||
It("accepts queue name at maximum length", func() {
|
||||
service.invokeCallbackFn = func(_ context.Context, _, _ string, _ []byte, _ int32) error {
|
||||
return nil
|
||||
}
|
||||
exactName := strings.Repeat("a", maxQueueNameLength)
|
||||
err := service.CreateQueue(ctx, exactName, host.QueueConfig{})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
Describe("CreateQueue defaults", func() {
|
||||
It("applies defaults for zero-value config", func() {
|
||||
err := service.CreateQueue(ctx, "defaults-queue", host.QueueConfig{})
|
||||
@ -102,6 +129,23 @@ var _ = Describe("TaskQueueService", func() {
|
||||
})
|
||||
})
|
||||
|
||||
Describe("CreateQueue defaults with negative values", func() {
|
||||
It("applies default RetentionMs for negative value", func() {
|
||||
service.invokeCallbackFn = func(_ context.Context, _, _ string, _ []byte, _ int32) error {
|
||||
return nil
|
||||
}
|
||||
err := service.CreateQueue(ctx, "neg-retention", host.QueueConfig{
|
||||
RetentionMs: -500,
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
service.mu.Lock()
|
||||
qs := service.queues["neg-retention"]
|
||||
service.mu.Unlock()
|
||||
Expect(qs.config.RetentionMs).To(Equal(defaultRetentionMs))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("CreateQueue clamping", func() {
|
||||
It("clamps concurrency exceeding maxConcurrency", func() {
|
||||
// maxConcurrency is 5; request 10
|
||||
@ -162,6 +206,20 @@ var _ = Describe("TaskQueueService", func() {
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("does not exist"))
|
||||
})
|
||||
|
||||
It("rejects payload exceeding maximum size", func() {
|
||||
bigPayload := make([]byte, maxPayloadSize+1)
|
||||
_, err := service.Enqueue(ctx, "enqueue-test", bigPayload)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("exceeds maximum"))
|
||||
})
|
||||
|
||||
It("accepts payload at maximum size", func() {
|
||||
exactPayload := make([]byte, maxPayloadSize)
|
||||
taskID, err := service.Enqueue(ctx, "enqueue-test", exactPayload)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(taskID).ToNot(BeEmpty())
|
||||
})
|
||||
})
|
||||
|
||||
Describe("GetTaskStatus", func() {
|
||||
@ -204,38 +262,30 @@ var _ = Describe("TaskQueueService", func() {
|
||||
})
|
||||
|
||||
It("cancels a pending task", func() {
|
||||
// Create queue with 0 concurrency via a trick: use delayMs to slow down processing
|
||||
// Actually, just stop workers by closing and recreating without workers
|
||||
service.Close()
|
||||
service = nil
|
||||
|
||||
// Recreate without starting workers - we'll create the queue after overriding invokeCallbackFn
|
||||
managerCtx2, cancel2 := context.WithCancel(ctx)
|
||||
DeferCleanup(cancel2)
|
||||
manager2 := &Manager{
|
||||
plugins: make(map[string]*plugin),
|
||||
ctx: managerCtx2,
|
||||
}
|
||||
var err error
|
||||
service, err = newTaskQueueService("test_plugin_cancel", manager2, 5)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Block the callback so task stays pending while we try to cancel
|
||||
// Block the callback so the first task occupies the worker
|
||||
started := make(chan struct{})
|
||||
service.invokeCallbackFn = func(ctx context.Context, _, _ string, _ []byte, _ int32) error {
|
||||
time.Sleep(10 * time.Second)
|
||||
return nil
|
||||
close(started)
|
||||
<-ctx.Done()
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
err = service.CreateQueue(ctx, "cancel-test", host.QueueConfig{
|
||||
err := service.CreateQueue(ctx, "cancel-test", host.QueueConfig{
|
||||
Concurrency: 1,
|
||||
DelayMs: 5000, // Large delay so worker doesn't grab it immediately
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Enqueue a blocker task to occupy the single worker
|
||||
_, err = service.Enqueue(ctx, "cancel-test", []byte("blocker"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Wait for the blocker task to start running
|
||||
Eventually(started).WithTimeout(5 * time.Second).Should(BeClosed())
|
||||
|
||||
// Enqueue a second task — it stays pending since the worker is busy
|
||||
taskID, err := service.Enqueue(ctx, "cancel-test", []byte("cancel-me"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Cancel quickly before worker picks it up
|
||||
err = service.CancelTask(ctx, taskID)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
@ -367,6 +417,89 @@ var _ = Describe("TaskQueueService", func() {
|
||||
})
|
||||
})
|
||||
|
||||
Describe("Backoff overflow cap", func() {
|
||||
It("caps backoff at maxRetentionMs to prevent overflow", func() {
|
||||
var callCount atomic.Int32
|
||||
|
||||
service.invokeCallbackFn = func(_ context.Context, _, _ string, _ []byte, _ int32) error {
|
||||
callCount.Add(1)
|
||||
return fmt.Errorf("always fail")
|
||||
}
|
||||
|
||||
err := service.CreateQueue(ctx, "backoff-overflow", host.QueueConfig{
|
||||
MaxRetries: 3,
|
||||
BackoffMs: 1_000_000_000, // Very large backoff to trigger overflow on exponentiation
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
taskID, err := service.Enqueue(ctx, "backoff-overflow", []byte("overflow-test"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Wait for first attempt to fail
|
||||
Eventually(func() int32 {
|
||||
return callCount.Load()
|
||||
}).WithTimeout(5 * time.Second).WithPolling(50 * time.Millisecond).Should(BeNumerically(">=", int32(1)))
|
||||
|
||||
// Check next_run_at is positive and reasonable (capped at maxRetentionMs from now)
|
||||
var nextRunAt int64
|
||||
err = service.db.QueryRow(`SELECT next_run_at FROM tasks WHERE id = ?`, taskID).Scan(&nextRunAt)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
now := time.Now().UnixMilli()
|
||||
Expect(nextRunAt).To(BeNumerically(">", int64(0)), "next_run_at should be positive")
|
||||
Expect(nextRunAt).To(BeNumerically("<=", now+maxBackoffMs+1000), "next_run_at should be at most maxBackoffMs from now")
|
||||
})
|
||||
})
|
||||
|
||||
Describe("Delay enforcement with concurrent workers", func() {
|
||||
It("enforces delay between dispatches even with multiple workers", func() {
|
||||
var mu sync.Mutex
|
||||
var dispatchTimes []time.Time
|
||||
|
||||
service.invokeCallbackFn = func(_ context.Context, _, _ string, _ []byte, _ int32) error {
|
||||
mu.Lock()
|
||||
dispatchTimes = append(dispatchTimes, time.Now())
|
||||
mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
err := service.CreateQueue(ctx, "delay-concurrent", host.QueueConfig{
|
||||
Concurrency: 3,
|
||||
DelayMs: 200,
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Enqueue 5 tasks
|
||||
for i := 0; i < 5; i++ {
|
||||
_, err := service.Enqueue(ctx, "delay-concurrent", []byte(fmt.Sprintf("task-%d", i)))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
}
|
||||
|
||||
// Wait for all tasks to complete
|
||||
Eventually(func() int {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
return len(dispatchTimes)
|
||||
}).WithTimeout(10 * time.Second).WithPolling(50 * time.Millisecond).Should(Equal(5))
|
||||
|
||||
// Sort dispatch times and verify gaps
|
||||
mu.Lock()
|
||||
sort.Slice(dispatchTimes, func(i, j int) bool {
|
||||
return dispatchTimes[i].Before(dispatchTimes[j])
|
||||
})
|
||||
times := make([]time.Time, len(dispatchTimes))
|
||||
copy(times, dispatchTimes)
|
||||
mu.Unlock()
|
||||
|
||||
// Consecutive dispatches should have at least ~160ms gap (80% of 200ms)
|
||||
for i := 1; i < len(times); i++ {
|
||||
gap := times[i].Sub(times[i-1])
|
||||
Expect(gap).To(BeNumerically(">=", 160*time.Millisecond),
|
||||
fmt.Sprintf("gap between dispatch %d and %d was %v, expected >= 160ms", i-1, i, gap))
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
Describe("Shutdown recovery", func() {
|
||||
It("resets stale running tasks on CreateQueue", func() {
|
||||
// Create a first service and queue, enqueue a task
|
||||
@ -405,15 +538,7 @@ var _ = Describe("TaskQueueService", func() {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Re-create the queue - this should reset stale running tasks
|
||||
// First we need to re-insert the queue row since it was from the old service
|
||||
// Actually the queue row is already there from the first service, but
|
||||
// CreateQueue will fail because the row exists. We need to handle this differently.
|
||||
// The queue metadata exists in DB, but not in the new service's memory map.
|
||||
// The schema has the queue row already. Let's delete it and re-create.
|
||||
_, err = service.db.Exec(`DELETE FROM queues WHERE name = 'recovery-queue'`)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Re-create the queue - the upsert handles the existing row from the old service
|
||||
err = service.CreateQueue(ctx, "recovery-queue", host.QueueConfig{})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
@ -752,9 +877,8 @@ var _ = Describe("TaskQueueService Integration", Ordered, func() {
|
||||
ctx := GinkgoT().Context()
|
||||
|
||||
// Create queue with concurrency=1 and a large delay between dispatches.
|
||||
// After the first task is dispatched (no delay for the first), the
|
||||
// second task will be dequeued but the worker will block waiting for
|
||||
// the 60s delay. Tasks 3+ remain in 'pending' status and can be cancelled.
|
||||
// The first task completes immediately (burst token), the second is dequeued
|
||||
// but blocks on the rate limiter. Tasks 3+ remain in 'pending' and can be cancelled.
|
||||
_, err := callTestTaskQueue(ctx, testTaskQueueInput{
|
||||
Operation: "create_queue",
|
||||
QueueName: "test-cancel",
|
||||
@ -765,8 +889,8 @@ var _ = Describe("TaskQueueService Integration", Ordered, func() {
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Enqueue several tasks - the first will be processed immediately,
|
||||
// the second will block in the delay wait (status=running),
|
||||
// Enqueue several tasks - the first will complete immediately,
|
||||
// the second will be dequeued but block on the rate limiter (status=running),
|
||||
// the rest will stay pending.
|
||||
var taskIDs []string
|
||||
for i := 0; i < 5; i++ {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user