learning/go-patterns/patterns/cancellation/cancellation.go

// Package cancellation demonstrates context cancellation in a channel loop.
package cancellation

import "context"

// Process consumes jobs until the channel closes or ctx is canceled. Results
// completed before cancellation are returned together with ctx.Err(). The
// callback runs synchronously, so there is no worker goroutine to leak.
func Process(ctx context.Context, jobs <-chan int, fn func(int) int) ([]int, error) {
	var results []int
	for {
		// Checking before select makes an already-canceled context deterministic
		// even when jobs is also ready.
		if err := ctx.Err(); err != nil {
			return results, err
		}

		select {
		case <-ctx.Done():
			return results, ctx.Err()
		case job, ok := <-jobs:
			if !ok {
				return results, nil
			}
			// The callback may cancel ctx. Do not start another job after that.
			if err := ctx.Err(); err != nil {
				return results, err
			}
			results = append(results, fn(job))
		}
	}
}