Go Worker Pool with Context Cancellation

package main

import (
	"context"
	"fmt"
	"sync"
	"time"
)

type Job struct {
	ID    int
	Value int
}

type Result struct {
	JobID int
	Out   int
	Err   error
}

func worker(ctx context.Context, jobs <-chan Job, results chan<- Result, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		select {
		case <-ctx.Done():
			return
		case job, ok := <-jobs:
			if !ok {
				return
			}
			if job.Value < 0 {
				results <- Result{JobID: job.ID, Err: fmt.Errorf("job %d: negative value %d", job.ID, job.Value)}
				continue
			}
			select {
			case <-ctx.Done():
				return
			case results <- Result{JobID: job.ID, Out: job.Value * job.Value}:
			}
		}
	}
}

func run(ctx context.Context, jobs []Job, workers int) ([]Result, error) {
	if workers < 1 {
		return nil, fmt.Errorf("workers must be at least 1, got %d", workers)
	}
	if len(jobs) == 0 {
		return nil, nil
	}

	jobCh := make(chan Job)
	resCh := make(chan Result, len(jobs))

	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go worker(ctx, jobCh, resCh, &wg)
	}

	go func() {
		defer close(jobCh)
		for _, job := range jobs {
			select {
			case <-ctx.Done():
				return
			case jobCh <- job:
			}
		}
	}()

	go func() {
		wg.Wait()
		close(resCh)
	}()

	results := make([]Result, 0, len(jobs))
	for result := range resCh {
		results = append(results, result)
	}
	if err := ctx.Err(); err != nil {
		return results, err
	}
	return results, nil
}

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	jobs := []Job{
		{ID: 1, Value: 4},
		{ID: 2, Value: 9},
		{ID: 3, Value: -1},
		{ID: 4, Value: 16},
		{ID: 5, Value: 25},
		{ID: 6, Value: 36},
	}

	results, err := run(ctx, jobs, 3)
	if err != nil {
		fmt.Println("stopped early:", err)
	}

	failed := 0
	for _, result := range results {
		if result.Err != nil {
			failed++
			fmt.Println("error:", result.Err)
			continue
		}
		fmt.Printf("job %d -> %d\n", result.JobID, result.Out)
	}
	fmt.Printf("%d results, %d failed\n", len(results), failed)
}