Context, Pipelines and Worker Pools

context.Context is how a Go program carries cancellation, deadlines and request-scoped values across API boundaries. Every standard-library call that can block takes one, and yours should too - as the first parameter, by convention, named ctx.

package main

import (
	"context"
	"errors"
	"fmt"
	"time"
)

func main() {
	// A context with a deadline. The cancel function must always be called -
	// defer it on the next line, every time, or you leak the timer.
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()

	if err := slowWork(ctx, 200*time.Millisecond); err != nil {
		fmt.Println("first:", err, "| deadline exceeded:", errors.Is(err, context.DeadlineExceeded))
	}

	// Enough time this time.
	ctx2, cancel2 := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel2()
	fmt.Println("second:", slowWork(ctx2, 20*time.Millisecond))

	// Cancellation by hand, which is what a user closing a connection looks like.
	ctx3, cancel3 := context.WithCancel(context.Background())
	go func() {
		time.Sleep(20 * time.Millisecond)
		cancel3()
	}()
	fmt.Println("third:", slowWork(ctx3, time.Second))
}

func slowWork(ctx context.Context, duration time.Duration) error {
	select {
	case <-time.After(duration):
		return nil
	case <-ctx.Done():
		return ctx.Err()          // DeadlineExceeded or Canceled
	}
}

Pass the context down; never store one in a struct, and never pass nil - context.Background() is the root, context.TODO() the placeholder while you work out where the real one comes from.

A worker pool

package main

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

type result struct {
	job   int
	value int
}

func main() {
	jobs := make(chan int)
	results := make(chan result)

	const workers = 4
	var wg sync.WaitGroup

	for w := 1; w <= workers; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for job := range jobs {              // ends when jobs is closed
				time.Sleep(10 * time.Millisecond)
				results <- result{job: job, value: job * job}
			}
		}()
	}

	// The sender closes jobs when there is no more work.
	go func() {
		defer close(jobs)
		for i := 1; i <= 12; i++ {
			jobs <- i
		}
	}()

	// A separate goroutine closes results once every worker has finished, which
	// is what lets the range below terminate.
	go func() {
		wg.Wait()
		close(results)
	}()

	total := 0
	for r := range results {
		total += r.value
	}
	fmt.Println("sum of squares 1..12 =", total)
}

Three channel rules are doing the work: the producer closes jobs, each worker ranges until it is closed, and a closer goroutine waits for all of them before closing results. Get one wrong and the program either deadlocks or panics on a send to a closed channel.

A pipeline

Each stage is a function that takes a channel and returns one, so stages compose and each runs concurrently with the others.

package main

import "fmt"

func main() {
	for value := range square(filterEven(generate(1, 10))) {
		fmt.Print(value, " ")
	}
	fmt.Println()
}

func generate(from, to int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for i := from; i <= to; i++ {
			out <- i
		}
	}()
	return out
}

func filterEven(in <-chan int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for n := range in {
			if n%2 == 0 {
				out <- n
			}
		}
	}()
	return out
}

func square(in <-chan int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for n := range in {
			out <- n * n
		}
	}()
	return out
}

<-chan int in the return type is a receive-only channel: the caller can read but not send or close. Directional channel types are free documentation the compiler enforces.

errgroup, for when any failure should cancel the rest

The standard library's sync.WaitGroup has no error handling. The common pattern - and what golang.org/x/sync/errgroup packages up - is a group that collects the first error and cancels a shared context.

package main

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

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

	var (
		wg    sync.WaitGroup
		once  sync.Once
		first error
	)

	fail := func(err error) {
		once.Do(func() {           // only the FIRST error is kept
			first = err
			cancel()               // and everyone else is told to stop
		})
	}

	for i := 1; i <= 4; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := task(ctx, i); err != nil {
				fail(err)
			}
		}()
	}
	wg.Wait()

	fmt.Println("first error:", first)
	fmt.Println("ctx cancelled:", ctx.Err())
}

func task(ctx context.Context, id int) error {
	select {
	case <-time.After(time.Duration(id*20) * time.Millisecond):
		if id == 2 {
			return fmt.Errorf("task %d failed", id)
		}
		fmt.Printf("  task %d done\n", id)
		return nil
	case <-ctx.Done():
		fmt.Printf("  task %d cancelled\n", id)
		return ctx.Err()
	}
}

Goroutine leaks

A goroutine blocked forever on a channel is never collected, and its stack and everything it references stay in memory. This is Go's characteristic leak.

package main

import (
	"context"
	"fmt"
	"runtime"
	"time"
)

func main() {
	before := runtime.NumGoroutine()

	// Leaks: nobody ever receives, so this goroutine blocks forever.
	leak := make(chan int)
	go func() { leak <- 1 }()

	time.Sleep(20 * time.Millisecond)
	fmt.Printf("goroutines: %d -> %d (one leaked)\n", before, runtime.NumGoroutine())

	// The fix: give every goroutine a way to be told to stop.
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		select {
		case <-ctx.Done():
			return
		case leak <- 1:
		}
	}()
	cancel()
	time.Sleep(20 * time.Millisecond)
	fmt.Println("goroutines after cancelling the second:", runtime.NumGoroutine())
}

Rule: whoever starts a goroutine is responsible for how it ends. A context, a done channel, or a closed input - but always something.

What to take away

• ctx first, always passed down, never stored in a struct; defer cancel() every time.
• ctx.Done() in a select is how a goroutine notices cancellation.
• Worker pool: producer closes the jobs channel, a closer goroutine waits then closes results.
• Pipelines compose stages that each own and close their output channel.
• <-chan T and chan<- T document direction and are checked by the compiler.
• Whoever starts a goroutine must arrange for it to finish, or it leaks.