Fan-Out and Fan-In Pipeline in Go

package main

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

type indexed struct {
	index int
	value int
}

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

	jobs := make(chan indexed)
	outputs := make(chan indexed)
	var wg sync.WaitGroup

	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for job := range jobs {
				select {
				case outputs <- indexed{index: job.index, value: job.value * job.value}:
				case <-ctx.Done():
					return
				}
			}
		}()
	}

	go func() {
		defer close(jobs)
		for index, value := range inputs {
			select {
			case jobs <- indexed{index: index, value: value}:
			case <-ctx.Done():
				return
			}
		}
	}()

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

	results := make([]int, len(inputs))
	seen := make([]bool, len(inputs))
	for output := range outputs {
		results[output.index] = output.value
		seen[output.index] = true
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	for index, ok := range seen {
		if !ok {
			return nil, fmt.Errorf("missing result for input %d", index)
		}
	}
	return results, nil
}

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

	inputs := []int{3, 1, 4, 1, 5, 9, 2, 6}
	squares, err := process(ctx, inputs, 3)
	if err != nil {
		fmt.Println("error:", err)
		return
	}
	fmt.Println(squares)
}