go-worker-pool-rate-limit

Answer

The implementation is a concurrent worker pool in Go with: - Configurable rate limiting via a time.Ticker — jobs are throttled to at most N per second before being queued - Graceful shutdown via context cancellation — workers exit cleanly when the context is cancelled, and Stop() waits for all in-flight work to finish - No channel panics — uses a done channel to signal shutdown so Submit returns false instead of panicking on a closed channel

Package: workerpool/workerpool.go

package workerpool

import (
    "context"
    "sync"
    "sync/atomic"
    "time"
)

// Job represents a unit of work to be processed by the pool.
type Job struct {
    ID      int
    Payload interface{}
}

// Result represents the outcome of processing a job.
type Result struct {
    Job   Job
    Err   error
    Value interface{}
}

// ProcessFunc is a function that processes a job and returns a result.
type ProcessFunc func(Job) (interface{}, error)

// WorkerPool manages a pool of goroutine workers that process jobs
// with a configurable rate limit.
type WorkerPool struct {
    numWorkers  int
    rateLimit   int
    jobChan     chan Job
    resultsChan chan Result
    cancel      context.CancelFunc
    process     ProcessFunc

    started   int64
    completed int64
    dropped   int64

    wg      sync.WaitGroup
    limiter *time.Ticker

    stopOnce sync.Once
    done     chan struct{}
}

func New(numWorkers, rateLimit, bufferSize int) *WorkerPool {
    if numWorkers <= 0 {
        numWorkers = 1
    }
    if rateLimit < 0 {
        rateLimit = 0
    }
    if bufferSize <= 0 {
        bufferSize = numWorkers * 2
    }

    p := &WorkerPool{
        numWorkers:  numWorkers,
        rateLimit:   rateLimit,
        jobChan:     make(chan Job, bufferSize),
        resultsChan: make(chan Result, bufferSize),
        done:        make(chan struct{}),
    }

    if rateLimit > 0 {
        interval := time.Second / time.Duration(rateLimit)
        p.limiter = time.NewTicker(interval)
    }

    return p
}

func (p *WorkerPool) Start(ctx context.Context) {
    ctx, p.cancel = context.WithCancel(ctx)

    for i := 0; i < p.numWorkers; i++ {
        p.wg.Add(1)
        go p.worker(ctx)
    }

    if p.limiter != nil {
        go func() {
            <-ctx.Done()
            p.limiter.Stop()
        }()
    }
}

func (p *WorkerPool) StartWithProcessor(ctx context.Context, process ProcessFunc) {
    p.process = process
    p.Start(ctx)
}

func (p *WorkerPool) Submit(ctx context.Context, job Job) bool {
    select {
    case <-ctx.Done():
        atomic.AddInt64(&p.dropped, 1)
        return false
    case <-p.done:
        atomic.AddInt64(&p.dropped, 1)
        return false
    default:
    }

    if p.limiter != nil {
        select {
        case <-ctx.Done():
            atomic.AddInt64(&p.dropped, 1)
            return false
        case <-p.done:
            atomic.AddInt64(&p.dropped, 1)
            return false
        case <-p.limiter.C:
        }
    }

    select {
    case <-ctx.Done():
        atomic.AddInt64(&p.dropped, 1)
        return false
    case <-p.done:
        atomic.AddInt64(&p.dropped, 1)
        return false
    case p.jobChan <- job:
        atomic.AddInt64(&p.started, 1)
        return true
    }
}

func (p *WorkerPool) Results() <-chan Result {
    return p.resultsChan
}

func (p *WorkerPool) Stop() {
    p.stopOnce.Do(func() {
        if p.cancel != nil {
            p.cancel()
        }
        p.wg.Wait()
        close(p.resultsChan)
        close(p.done)
    })
}

func (p *WorkerPool) Stats() (started, completed, dropped int64) {
    return atomic.LoadInt64(&p.started),
        atomic.LoadInt64(&p.completed),
        atomic.LoadInt64(&p.dropped)
}

func (p *WorkerPool) worker(ctx context.Context) {
    defer p.wg.Done()

    for {
        select {
        case <-ctx.Done():
            return
        case job, ok := <-p.jobChan:
            if !ok {
                return
            }
            if p.process != nil {
                val, err := p.process(job)
                p.resultsChan <- Result{Job: job, Value: val, Err: err}
            } else {
                p.resultsChan <- Result{Job: job, Value: job.Payload}
            }
            atomic.AddInt64(&p.completed, 1)
        }
    }
}

Key design decisions