internal/covsnap/covsnap.go

Ref: Size: 1.8 KiB History
// Package covsnap lets a long-running binary flush integration-coverage
// counters on demand. It is a no-op unless built with `go build -cover` AND
// GOCOVERDIR is set, so it costs nothing in a normal build.
package covsnap

import (
	"context"
	"log/slog"
	"os"
	"os/signal"
	"runtime/coverage"
	"sync"
	"syscall"
)

// write is a seam for tests; production writes real coverage data.
var write = func(dir string) error {
	if err := coverage.WriteMetaDir(dir); err != nil {
		return err
	}
	return coverage.WriteCountersDir(dir)
}

// Install starts a goroutine that, on each SIGUSR1, snapshots coverage into
// $GOCOVERDIR. Unset GOCOVERDIR -> returns without registering anything.
//
// The returned stop tears the handler down synchronously: it unregisters the
// signal and blocks until the goroutine has exited, so no handler outlives the
// caller. Production wires this to the process lifetime and lets ctx end it, so
// it ignores the return; a test that installs and uninstalls in the same
// process must call stop, or its handler leaks into the next test and a stray
// SIGUSR1 races the package-level write seam. stop is idempotent.
func Install(ctx context.Context) (stop func()) {
	dir := os.Getenv("GOCOVERDIR")
	if dir == "" {
		return func() {}
	}
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGUSR1)
	quit := make(chan struct{})
	done := make(chan struct{})
	go func() {
		defer close(done)
		defer signal.Stop(ch)
		for {
			select {
			case <-ctx.Done():
				return
			case <-quit:
				return
			case <-ch:
				if err := write(dir); err != nil {
					slog.Warn("covsnap: write failed", "dir", dir, "err", err)
					continue
				}
				slog.Info("covsnap: coverage written", "dir", dir)
			}
		}
	}()
	var once sync.Once
	return func() {
		once.Do(func() {
			signal.Stop(ch)
			close(quit)
		})
		<-done
	}
}