139 lines
3.6 KiB
Go
139 lines
3.6 KiB
Go
package counting
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"runtime/debug"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// QueueConfig sizes the counting queue.
|
|
type QueueConfig struct {
|
|
Size int // queued jobs before Enqueue refuses (default 10 000)
|
|
Workers int // concurrent workers (default 8)
|
|
Timeout time.Duration // budget per job: plan, Redis, log (default 4s)
|
|
}
|
|
|
|
// Observer receives queue events for metrics and logging; nil funcs are ignored.
|
|
type Observer struct {
|
|
Depth func(n int) // queue depth after every enqueue/dequeue
|
|
Wait func(d time.Duration) // time a job waited in the queue
|
|
Panic func(job Job, p any, stack []byte) // a handler panicked (recovered)
|
|
}
|
|
|
|
// Queue decouples counting from the HTTP request: /in and /out validate,
|
|
// enqueue a Job and answer 202; workers run the handler with their own
|
|
// bounded context, so a client that hangs up or times out can no longer
|
|
// cancel a count (review #1).
|
|
type Queue struct {
|
|
cfg QueueConfig
|
|
jobs chan Job
|
|
handler func(ctx context.Context, job Job)
|
|
obs Observer
|
|
|
|
pending atomic.Int64 // enqueued and not yet finished
|
|
mu sync.RWMutex // Enqueue vs Drain (sending on a closed channel panics)
|
|
closed bool
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
// NewQueue applies defaults and starts the workers.
|
|
func NewQueue(cfg QueueConfig, handler func(ctx context.Context, job Job), obs Observer) *Queue {
|
|
if cfg.Size <= 0 {
|
|
cfg.Size = 10000
|
|
}
|
|
if cfg.Workers <= 0 {
|
|
cfg.Workers = 8
|
|
}
|
|
if cfg.Timeout <= 0 {
|
|
cfg.Timeout = 4 * time.Second
|
|
}
|
|
q := &Queue{cfg: cfg, jobs: make(chan Job, cfg.Size), handler: handler, obs: obs}
|
|
for i := 0; i < cfg.Workers; i++ {
|
|
q.wg.Add(1)
|
|
go q.worker()
|
|
}
|
|
return q
|
|
}
|
|
|
|
// Enqueue hands a job to the workers without blocking. It returns false when
|
|
// the queue is full or draining; the caller answers 503 and logs the drop.
|
|
func (q *Queue) Enqueue(job Job) bool {
|
|
if job.Received.IsZero() {
|
|
job.Received = time.Now()
|
|
}
|
|
q.mu.RLock()
|
|
defer q.mu.RUnlock()
|
|
if q.closed {
|
|
return false
|
|
}
|
|
select {
|
|
case q.jobs <- job:
|
|
q.pending.Add(1)
|
|
q.observeDepth()
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
// Depth is the number of jobs waiting for a worker.
|
|
func (q *Queue) Depth() int { return len(q.jobs) }
|
|
|
|
// Pending is the number of jobs enqueued and not yet finished (waiting or
|
|
// in flight); tests and the shutdown log use it.
|
|
func (q *Queue) Pending() int { return int(q.pending.Load()) }
|
|
|
|
// Drain stops accepting, lets the workers finish every queued job and waits
|
|
// for them, bounded by ctx. On an expired ctx it reports how many jobs were
|
|
// still pending.
|
|
func (q *Queue) Drain(ctx context.Context) error {
|
|
q.mu.Lock()
|
|
if !q.closed {
|
|
q.closed = true
|
|
close(q.jobs)
|
|
}
|
|
q.mu.Unlock()
|
|
done := make(chan struct{})
|
|
go func() { q.wg.Wait(); close(done) }()
|
|
select {
|
|
case <-done:
|
|
return nil
|
|
case <-ctx.Done():
|
|
return fmt.Errorf("counting queue: %d jobs still pending: %w", q.Pending(), ctx.Err())
|
|
}
|
|
}
|
|
|
|
func (q *Queue) worker() {
|
|
defer q.wg.Done()
|
|
for job := range q.jobs {
|
|
if q.obs.Wait != nil {
|
|
q.obs.Wait(time.Since(job.Received))
|
|
}
|
|
q.run(job)
|
|
q.pending.Add(-1)
|
|
q.observeDepth()
|
|
}
|
|
}
|
|
|
|
// run executes one job under the queue's own timeout, never a request
|
|
// context, and keeps the worker alive when the handler panics.
|
|
func (q *Queue) run(job Job) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), q.cfg.Timeout)
|
|
defer cancel()
|
|
defer func() {
|
|
if p := recover(); p != nil && q.obs.Panic != nil {
|
|
q.obs.Panic(job, p, debug.Stack())
|
|
}
|
|
}()
|
|
q.handler(ctx, job)
|
|
}
|
|
|
|
func (q *Queue) observeDepth() {
|
|
if q.obs.Depth != nil {
|
|
q.obs.Depth(len(q.jobs))
|
|
}
|
|
}
|