Priority Worker Pool in Go

package main

import (
	"container/heap"
	"context"
	"fmt"
	"strings"
	"sync"
	"time"
)

type Job struct {
	ID       int
	Priority int
	Payload  string
}

type Result struct {
	JobID  int
	Output string
}

type jobQueue []Job

func (q jobQueue) Len() int { return len(q) }

func (q jobQueue) Less(i, j int) bool {
	if q[i].Priority != q[j].Priority {
		return q[i].Priority > q[j].Priority
	}
	return q[i].ID < q[j].ID
}

func (q jobQueue) Swap(i, j int) { q[i], q[j] = q[j], q[i] }

func (q *jobQueue) Push(value any) { *q = append(*q, value.(Job)) }

func (q *jobQueue) Pop() any {
	old := *q
	last := len(old) - 1
	item := old[last]
	old[last] = Job{}
	*q = old[:last]
	return item
}

func runPool(ctx context.Context, jobs []Job, workers int) ([]Result, error) {
	if workers < 1 {
		return nil, fmt.Errorf("workers must be at least 1")
	}

	queue := make(jobQueue, len(jobs))
	copy(queue, jobs)
	heap.Init(&queue)

	jobCh := make(chan Job)
	resultCh := make(chan Result)
	var wg sync.WaitGroup

	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for job := range jobCh {
				output := strings.ToUpper(job.Payload)
				select {
				case resultCh <- Result{JobID: job.ID, Output: output}:
				case <-ctx.Done():
					return
				}
			}
		}()
	}

	go func() {
		defer close(jobCh)
		for queue.Len() > 0 {
			select {
			case <-ctx.Done():
				return
			case jobCh <- heap.Pop(&queue).(Job):
			}
		}
	}()

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

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

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

	jobs := []Job{
		{ID: 1, Priority: 2, Payload: "billing"},
		{ID: 2, Priority: 5, Payload: "auth"},
		{ID: 3, Priority: 5, Payload: "search"},
		{ID: 4, Priority: 1, Payload: "cleanup"},
	}

	results, err := runPool(ctx, jobs, 2)
	if err != nil {
		fmt.Println("pool stopped:", err)
	}
	for _, result := range results {
		fmt.Printf("job %d -> %s\n", result.JobID, result.Output)
	}
}