learning/go-patterns/stages/07-cancel-join/main.go

// Stage 07 keeps the concurrency limit and adds cancellation plus joining.
package main

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

// Long-running tasks must observe ctx or pass it to a cancellable API.
type Task func(context.Context) (string, error)

// RunFailFast returns ordered results after every worker has stopped.
// Unstarted tasks retain empty results. The first worker to record an error
// cancels the group; simultaneous failures have no fixed winner.
func RunFailFast(parent context.Context, tasks []Task, limit int) ([]string, error) {
	if limit < 1 {
		return nil, errors.New("limit must be positive")
	}
	for i, task := range tasks {
		if task == nil {
			return nil, fmt.Errorf("task %d is nil", i)
		}
	}
	values := make([]string, len(tasks))
	ctx, cancel := context.WithCancel(parent)
	defer cancel()
	jobs := make(chan int)
	var workers sync.WaitGroup
	var first sync.Once
	var firstErr error
	for range min(limit, len(tasks)) {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for index := range jobs {
				if ctx.Err() != nil {
					return
				}
				value, err := tasks[index](ctx)
				values[index] = value
				if err != nil {
					first.Do(func() { firstErr = err; cancel() })
					return
				}
			}
		}()
	}
dispatch:
	for index := range tasks {
		select {
		case jobs <- index:
		case <-ctx.Done():
			break dispatch
		}
	}
	close(jobs)
	workers.Wait()
	if firstErr != nil {
		return values, firstErr
	}
	return values, parent.Err()
}

func main() {
	started := make(chan struct{})
	stopped := make(chan struct{})
	wantErr := errors.New("shipping service unavailable")
	_, err := RunFailFast(context.Background(), []Task{
		func(ctx context.Context) (string, error) {
			close(started)
			defer close(stopped)
			<-ctx.Done()
			return "", ctx.Err()
		},
		func(context.Context) (string, error) { <-started; return "", wantErr },
	}, 2)
	<-stopped
	fmt.Println("error:", err)
	fmt.Println("workers stopped")
}