concurrency-producer-consumer-pipeline

Answer

Here's a complete multi-stage pipeline in Go using channels and context cancellation:

package main

import (
    "context"
    "fmt"
    "os"
    "os/signal"
    "syscall"
)

// generator emits integers 1..n on the output channel.
// It stops early if ctx is cancelled.
func generator(ctx context.Context, n int) <-chan int {
    out := make(chan int)
    go func() {
        defer close(out)
        for i := 1; i <= n; i++ {
            select {
            case out <- i:
            case <-ctx.Done():
                return
            }
        }
    }()
    return out
}

// square reads ints from in, squares each one, and sends the result.
// Stops early if ctx is cancelled, draining the input to prevent
// upstream goroutine leaks.
func square(ctx context.Context, in <-chan int) <-chan int {
    out := make(chan int)
    go func() {
        defer close(out)
        for v := range in {
            result := v * v
            select {
            case out <- result:
            case <-ctx.Done():
                for range in { // drain upstream
                }
                return
            }
        }
    }()
    return out
}

// filter passes through only values >= threshold.
// Stops early if ctx is cancelled.
func filter(ctx context.Context, threshold int, in <-chan int) <-chan int {
    out := make(chan int)
    go func() {
        defer close(out)
        for v := range in {
            if v < threshold {
                continue
            }
            select {
            case out <- v:
            case <-ctx.Done():
                for range in { // drain upstream
                }
                return
            }
        }
    }()
    return out
}

// print reads ints from in and prints them to stdout.
// Returns once in is closed or ctx is cancelled.
func printFunc(ctx context.Context, in <-chan int) {
    for {
        select {
        case v, ok := <-in:
            if !ok {
                return
            }
            fmt.Println(v)
        case <-ctx.Done():
            return
        }
    }
}

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    // Trap SIGINT/SIGTERM for graceful shutdown.
    sigCh := make(chan os.Signal, 1)
    signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
    go func() {
        <-sigCh
        fmt.Fprintln(os.Stderr, "\nshutting down...")
        cancel()
    }()

    // Pipeline: generator -> square -> filter -> print.
    c1 := generator(ctx, 20)
    c2 := square(ctx, c1)
    c3 := filter(ctx, 50, c2) // only keep squares >= 50
    printFunc(ctx, c3)

    fmt.Println("done")
}

Key design decisions:

1. Each stage is a goroutine that reads from an input <-chan int, processes values, and writes to an output <-chan int (owned by the stage).
2. defer close(out) ensures the next stage knows when no more data is coming.
3. select with ctx.Done() in every send allows instant cancellation at any blocking point.
4. Input draining (for range in {}) in downstream stages upon cancellation prevents upstream goroutines from blocking forever trying to send to a channel nobody reads.
5. printFunc uses a select with ctx.Done() for the receiver side, so it can exit even if no value is ready.