two/pkg/worker/queue.go

package worker

import (
	"log"
	"sync"
)

// Task is a function to be executed asynchronously by a worker.
type Task func()

// Queue is a FIFO channel-backed task queue consumed by worker goroutines.
type Queue struct {
	tasks   chan Task
	wg      sync.WaitGroup
	mu      sync.RWMutex
	stopped bool
}

// New creates a Queue with the given channel buffer size.
func New(bufferSize int) *Queue {
	return &Queue{tasks: make(chan Task, bufferSize)}
}

// Submit enqueues a task. Blocks if the queue is full. Tasks submitted after
// Stop are rejected and logged rather than enqueued.
func (q *Queue) Submit(t Task) {
	q.mu.RLock()
	defer q.mu.RUnlock()

	if q.stopped {
		log.Print("worker: queue stopped, task rejected")
		return
	}
	q.tasks <- t
}

// Start launches n worker goroutines that consume and execute tasks.
func (q *Queue) Start(n int) {
	log.Printf("worker: starting %d workers", n)
	for i := range n {
		q.wg.Add(1)
		go func(id int) {
			defer q.wg.Done()
			for task := range q.tasks {
				task()
			}
		}(i)
	}
}

// Stop rejects new tasks, then waits for the queued and in-flight ones to
// finish. The write lock is what makes closing the channel safe: it is only
// taken once every in-flight Submit has released its read lock.
func (q *Queue) Stop() {
	q.mu.Lock()
	if q.stopped {
		q.mu.Unlock()
		return
	}
	q.stopped = true
	close(q.tasks)
	q.mu.Unlock()

	q.wg.Wait()
}