learning/go-patterns/stages/06-bounded/main.go

// Stage 06: run a batch with a fixed number of workers.
package main

import (
	"context"
	"errors"
	"fmt"
	"sync"
)

// Task is one batch item. This stage focuses on limiting active work; the
// context-aware cancellation protocol is introduced in stage 07.
type Task func() (string, error)

type indexedTask struct {
	index int
	task  Task
}

type indexedResult struct {
	index int
	value string
	err   error
}

// RunBounded runs every task with at most limit tasks active at once. It
// preserves input order and waits for the whole batch before returning. If
// several tasks fail, the error belonging to the lowest input index is
// returned so the result is deterministic.
func RunBounded(tasks []Task, limit int) ([]string, error) {
	if limit < 1 {
		return nil, errors.New("limit must be positive")
	}
	for index, task := range tasks {
		if task == nil {
			return nil, fmt.Errorf("task %d is nil", index)
		}
	}
	values := make([]string, len(tasks))
	if len(tasks) == 0 {
		return values, nil
	}

	workerCount := limit
	if workerCount > len(tasks) {
		workerCount = len(tasks)
	}
	jobs := make(chan indexedTask)
	results := make(chan indexedResult, len(tasks))

	var workers sync.WaitGroup
	workers.Add(workerCount)
	for worker := 0; worker < workerCount; worker++ {
		go func() {
			defer workers.Done()
			for job := range jobs {
				value, err := job.task()
				results <- indexedResult{index: job.index, value: value, err: err}
			}
		}()
	}

	// Sending from this goroutine applies back-pressure: once all workers are
	// busy, dispatch waits until one worker takes another job.
	for index, task := range tasks {
		jobs <- indexedTask{index: index, task: task}
	}
	close(jobs)
	workers.Wait()
	close(results)

	errs := make([]error, len(tasks))
	for result := range results {
		values[result.index] = result.value
		errs[result.index] = result.err
	}
	for _, err := range errs {
		if err != nil {
			return values, err
		}
	}
	return values, nil
}

func main() {
	service := NewService(Catalog{"book": 2100, "pen": 300, "shipping": 500})
	prices, err := service.QuoteMany(context.Background(), []string{"book", "pen"}, 2)
	if err != nil {
		fmt.Println("quotes failed:", err)
		return
	}
	for _, price := range prices {
		fmt.Printf("%s: %d cents\n", price.Item, price.Cents)
	}
}