Automate the Boring Stuff with Go, Ch. 14: Concurrency

// goroutines run functions at the same time; a WaitGroup waits for them to finish
package main

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

func fetch(name string, took time.Duration) {
	time.Sleep(took) // stands in for a network call
	fmt.Printf("%s done after %s\n", name, took)
}

func main() {
	start := time.Now()

	// one after another: about 600 ms
	fetch("a", 300*time.Millisecond)
	fetch("b", 300*time.Millisecond)
	fmt.Println("sequential:", time.Since(start).Round(100*time.Millisecond))

	// at the same time: about 300 ms
	start = time.Now()
	var wg sync.WaitGroup
	for _, name := range []string{"c", "d", "e"} {
		wg.Go(func() { fetch(name, 300*time.Millisecond) })
	}
	wg.Wait() // blocks until every goroutine started with wg.Go has returned
	fmt.Println("concurrent:", time.Since(start).Round(100*time.Millisecond))
}

// channels carry values between goroutines; this is the worker-pool shape
package main

import (
	"fmt"
	"sync"
)

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

	// three workers read jobs until the channel is closed
	var wg sync.WaitGroup
	for w := 1; w <= 3; w++ {
		wg.Go(func() {
			for n := range jobs {
				results <- fmt.Sprintf("worker %d squared %d = %d", w, n, n*n)
			}
		})
	}

	// one goroutine feeds the jobs and closes the channel when there are no more
	go func() {
		for n := 1; n <= 6; n++ {
			jobs <- n
		}
		close(jobs)
	}()

	// close results once every worker is finished, so the range below ends
	go func() {
		wg.Wait()
		close(results)
	}()

	for r := range results {
		fmt.Println(r)
	}
}

// two goroutines touching one variable is a data race unless you say who goes first
package main

import (
	"fmt"
	"sync"
)

func main() {
	var wg sync.WaitGroup

	// wrong: the ++ is read-modify-write, and the goroutines interleave.
	// Run with `go run -race .` and the race detector will say so.
	unsafe := 0
	for range 1000 {
		wg.Go(func() { unsafe++ })
	}
	wg.Wait()
	fmt.Println("without a mutex:", unsafe, "(often less than 1000)")

	// right: a mutex makes the increment one step at a time
	var mu sync.Mutex
	safe := 0
	for range 1000 {
		wg.Go(func() {
			mu.Lock()
			safe++
			mu.Unlock()
		})
	}
	wg.Wait()
	fmt.Println("with a mutex:   ", safe)
}

// errgroup: run tasks concurrently, stop all of them at the first error, keep that error
package main

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

	"golang.org/x/sync/errgroup"
)

func work(ctx context.Context, name string, fail bool) error {
	for i := range 5 {
		select {
		case <-ctx.Done():
			fmt.Printf("%s: cancelled at step %d\n", name, i)
			return ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
		if fail && i == 1 {
			return errors.New(name + ": disk full")
		}
	}
	fmt.Printf("%s: finished\n", name)
	return nil
}

func main() {
	g, ctx := errgroup.WithContext(context.Background())
	g.SetLimit(2) // at most two tasks at once
	g.Go(func() error { return work(ctx, "alpha", false) })
	g.Go(func() error { return work(ctx, "beta", true) })
	g.Go(func() error { return work(ctx, "gamma", false) })

	if err := g.Wait(); err != nil {
		fmt.Println("first error:", err)
	}
}

// parallel-xkcd downloads the N most recent xkcd comics with a handful of goroutines at
// once. It is chapter 11's downloader with a worker pool, and it is the pattern for any
// "do this for a thousand URLs" job: a bounded number of workers, a context that stops
// everything on ctrl+c, and the first error reported instead of swallowed.
//
//	go run ./ch14/projects/parallel-xkcd -n 20 -workers 4 -dir xkcd
package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"path"
	"path/filepath"
	"time"

	"golang.org/x/sync/errgroup"
)

var client = &http.Client{Timeout: 30 * time.Second}

// Comic is the subset of xkcd's JSON we use.
type Comic struct {
	Num   int    `json:"num"`
	Title string `json:"title"`
	Img   string `json:"img"`
}

// Fetcher knows where the API lives, so tests can point it at a local server.
type Fetcher struct {
	Base string
}

func (f Fetcher) comic(ctx context.Context, url string) (Comic, error) {
	var c Comic
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return c, err
	}
	req.Header.Set("User-Agent", "automate-the-boring-stuff-with-go/2026")
	resp, err := client.Do(req)
	if err != nil {
		return c, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return c, fmt.Errorf("%s: %s", url, resp.Status)
	}
	return c, json.NewDecoder(resp.Body).Decode(&c)
}

// Latest returns the newest comic number.
func (f Fetcher) Latest(ctx context.Context) (int, error) {
	c, err := f.comic(ctx, f.Base+"/info.0.json")
	return c.Num, err
}

// Download fetches comic n's metadata and image into dir. It reports whether it wrote a
// file (false when it was already there).
func (f Fetcher) Download(ctx context.Context, n int, dir string) (bool, error) {
	c, err := f.comic(ctx, fmt.Sprintf("%s/%d/info.0.json", f.Base, n))
	if err != nil {
		return false, fmt.Errorf("comic %d: %w", n, err)
	}
	name := filepath.Join(dir, fmt.Sprintf("%d-%s", c.Num, path.Base(c.Img)))
	if _, err := os.Stat(name); err == nil {
		return false, nil
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.Img, nil)
	if err != nil {
		return false, err
	}
	resp, err := client.Do(req)
	if err != nil {
		return false, fmt.Errorf("image %d: %w", n, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return false, fmt.Errorf("image %d: %s", n, resp.Status)
	}
	tmp := name + ".part"
	out, err := os.Create(tmp)
	if err != nil {
		return false, err
	}
	if _, err := io.Copy(out, resp.Body); err != nil {
		out.Close()
		return false, err
	}
	if err := out.Close(); err != nil {
		return false, err
	}
	return true, os.Rename(tmp, name)
}

// DownloadRecent fetches the n newest comics with at most workers in flight, and returns
// how many files were written. The first error cancels the rest.
func (f Fetcher) DownloadRecent(ctx context.Context, n, workers int, dir string) (int, error) {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return 0, err
	}
	latest, err := f.Latest(ctx)
	if err != nil {
		return 0, err
	}
	g, ctx := errgroup.WithContext(ctx)
	g.SetLimit(workers)
	written := make(chan int, n) // each worker sends 1 when it writes a file
	for num := latest; num > latest-n && num > 0; num-- {
		g.Go(func() error {
			wrote, err := f.Download(ctx, num, dir)
			if err != nil {
				return err
			}
			if wrote {
				written <- 1
				fmt.Printf("#%d\n", num)
			}
			return nil
		})
	}
	err = g.Wait()
	close(written)
	count := 0
	for range written {
		count++
	}
	return count, err
}

func run(args []string) error {
	flags := flag.NewFlagSet("parallel-xkcd", flag.ContinueOnError)
	n := flags.Int("n", 20, "how many recent comics")
	workers := flags.Int("workers", 4, "how many downloads at once")
	dir := flags.String("dir", "xkcd", "folder to save into")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if *n < 1 || *workers < 1 {
		return errors.New("-n and -workers must be at least 1")
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	start := time.Now()
	count, err := Fetcher{Base: "https://xkcd.com"}.DownloadRecent(ctx, *n, *workers, *dir)
	fmt.Printf("%d new file(s) in %s after %s\n", count, *dir, time.Since(start).Round(time.Millisecond))
	return err
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// find-duplicates lists files under a folder that have identical contents, by hashing
// every file with SHA-256. Hashing is the slow part, so a pool of workers does it while
// one goroutine walks the tree and one collects the answers.
//
//	go run ./ch14/projects/find-duplicates ~/Pictures
package main

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"runtime"
	"slices"
	"sync"
)

// hashed is one file's result.
type hashed struct {
	path string
	sum  string
	err  error
}

// HashFile returns the hex SHA-256 of a file's contents, streaming so size does not matter.
func HashFile(path string) (string, error) {
	f, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}

// FindDuplicates returns groups of paths with identical contents (groups of size >= 2),
// each group sorted, groups sorted by their first path. minSize skips tiny files.
func FindDuplicates(root string, minSize int64, workers int) ([][]string, error) {
	paths := make(chan string)
	results := make(chan hashed)

	// the walker: finds files and sends their paths
	var walkErr error
	go func() {
		defer close(paths)
		walkErr = filepath.WalkDir(root, func(p string, d fs.DirEntry, err error) error {
			if err != nil {
				return nil // unreadable: skip, do not stop
			}
			if d.IsDir() {
				return nil
			}
			info, err := d.Info()
			if err != nil || info.Size() < minSize || !info.Mode().IsRegular() {
				return nil
			}
			paths <- p
			return nil
		})
	}()

	// the workers: hash whatever the walker sends
	var wg sync.WaitGroup
	for range workers {
		wg.Go(func() {
			for p := range paths {
				sum, err := HashFile(p)
				results <- hashed{p, sum, err}
			}
		})
	}
	go func() {
		wg.Wait()
		close(results)
	}()

	// the collector: this goroutine, grouping by hash
	bySum := map[string][]string{}
	for r := range results {
		if r.err != nil {
			continue // a file that vanished or cannot be read is not a duplicate of anything
		}
		bySum[r.sum] = append(bySum[r.sum], r.path)
	}
	if walkErr != nil {
		return nil, walkErr
	}

	var groups [][]string
	for _, g := range bySum {
		if len(g) >= 2 {
			slices.Sort(g)
			groups = append(groups, g)
		}
	}
	slices.SortFunc(groups, func(a, b []string) int { return slices.Compare(a, b) })
	return groups, nil
}

func run(args []string) error {
	flags := flag.NewFlagSet("find-duplicates", flag.ContinueOnError)
	minSize := flags.Int64("min", 1, "ignore files smaller than this many bytes")
	workers := flags.Int("workers", runtime.NumCPU(), "how many files to hash at once")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 1 {
		return errors.New("usage: find-duplicates [-min bytes] [-workers N] <folder>")
	}
	groups, err := FindDuplicates(flags.Arg(0), *minSize, *workers)
	if err != nil {
		return err
	}
	for i, g := range groups {
		fmt.Printf("group %d (%d copies):\n", i+1, len(g))
		for _, p := range g {
			fmt.Println("  ", p)
		}
	}
	fmt.Fprintf(os.Stderr, "%d group(s) of duplicates\n", len(groups))
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}