feat: make the worker handoff write path cancellable

pkg/collector/collector.go

@@ -1,11 +1,14 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/eraser-dev/eraser/pkg/cri"
@@ -85,11 +88,31 @@ func main() {
 		os.Exit(1)
 	}
 
-	if err := util.WriteImagesPipe(path, finalImages); err != nil {
+	// Registering the handler suppresses the default SIGTERM exit, so it covers
+	// exactly the one call that observes ctx. Everything above builds its own
+	// timeouts from Background, and Await below has no context at all; holding
+	// the handler across either would swallow the signal.
+	ctx, stopSignals := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+
+	if err := util.WriteImagesPipe(ctx, path, finalImages); err != nil {
+		stopSignals()
 		log.Error(err, "failed to send images", "pipeFile", path)
 		os.Exit(1)
 	}
 
+	// Read before stopping, because stopSignals cancels this context itself:
+	// checked afterwards it would always report Canceled, and the collector
+	// would exit on every successful run instead of waiting for the erase.
+	sigErr := ctx.Err()
+	stopSignals()
+
+	// A signal that landed while the handler was registered was consumed rather
+	// than killing the process, so it has to be acted on here.
+	if sigErr != nil {
+		log.Error(sigErr, "terminating before waiting for completion")
+		os.Exit(1)
+	}
+
 	data, err := completion.Await()
 	if err != nil {
 		log.Error(err, "failed to read pipe", "pipeFile", util.EraseCompleteCollectPath)

pkg/remover/helpers.go

@@ -8,10 +8,13 @@ import (
 	util "github.com/eraser-dev/eraser/pkg/utils"
 )
 
-func removeImages(c cri.Remover, targetImages []string) (int, error) {
+func removeImages(ctx context.Context, c cri.Remover, targetImages []string) (int, error) {
 	removed := 0
 
-	backgroundContext, cancel := context.WithTimeout(context.Background(), timeout)
+	// Derived from the caller's context, not Background: signal notification is
+	// registered for the whole process, so nothing would observe a SIGTERM during
+	// the deletion loop otherwise.
+	backgroundContext, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
 	images, err := c.ListImages(backgroundContext)

pkg/remover/remover.go

@@ -106,33 +106,45 @@ func main() {
 		log.Info("no images to exclude")
 	}
 
-	removed, err := removeImages(client, imagelist)
+	// Registering the handler suppresses the default SIGTERM exit, so it starts
+	// only here: once the peer publishes its endpoint the read above blocks in a
+	// call no context can interrupt, and covering it would swallow the signal
+	// until SIGKILL. The stop func is discarded rather than deferred because
+	// every exit path below is os.Exit, which would skip it anyway.
+	ctx, _ := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+
+	removed, err := removeImages(ctx, client, imagelist)
 	if err != nil {
 		log.Error(err, "failed to remove images")
 		os.Exit(generalErr)
 	}
 
+	// A signal that landed during removal was consumed rather than killing the
+	// process, and with --imagelist there is no completion write below to report
+	// it, so an interrupted run would otherwise exit 0 having removed nothing.
+	if err := ctx.Err(); err != nil {
+		log.Error(err, "terminating before removal finished", "removed", removed)
+		os.Exit(generalErr)
+	}
+
 	if os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT") != "" {
 		// record metrics
-		ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
-
 		exporter, reader, provider := metrics.ConfigureMetrics(ctx, log, os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT"))
 		otel.SetMeterProvider(provider)
 
 		if err := metrics.RecordMetricsRemover(ctx, otel.GetMeterProvider(), int64(removed)); err != nil {
 			log.Error(err, "error recording metrics")
 		}
 		metrics.ExportMetrics(log, exporter, reader)
-		cancel()
 	}
 
 	if *imageListPtr == "" {
-		if err := util.WriteCompletionPipe(util.EraseCompleteCollectPath); err != nil {
+		if err := util.WriteCompletionPipe(ctx, util.EraseCompleteCollectPath); err != nil {
 			log.Error(err, "unable to signal completion", "pipeFile", util.EraseCompleteCollectPath)
 			os.Exit(generalErr)
 		}
 
-		err := util.WriteCompletionPipe(util.EraseCompleteScanPath)
+		err := util.WriteCompletionPipe(ctx, util.EraseCompleteScanPath)
 		// if the scanner is disabled
 		if os.IsNotExist(err) {
 			return

pkg/remover/remover_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"testing"
 
 	v1 "k8s.io/cri-api/pkg/apis/runtime/v1"
@@ -50,7 +51,7 @@ func TestRemoveImages(t *testing.T) {
 				}
 			}
 
-			_, err := removeImages(client, tc.remove)
+			_, err := removeImages(context.Background(), client, tc.remove)
 			if tc.shouldErr && err == nil {
 				t.Fatal("expected error, got none")
 			}
@@ -85,3 +86,26 @@ func TestRemoveImages(t *testing.T) {
 		})
 	}
 }
+
+// removeImages builds its deadline from the caller rather than Background, so
+// that a SIGTERM registered process-wide actually reaches the runtime. Nothing
+// above notices if that regresses -- every case passes a context that is never
+// done -- so this pins it: a caller who has already gone must not get images
+// deleted on their behalf.
+func TestRemoveImagesPassesTheCallersContextToTheRuntime(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	client := &testClient{t: t, images: []*v1.Image{{Id: "sha256:aaaa"}}}
+
+	removed, err := removeImages(ctx, client, []string{"sha256:aaaa"})
+	if err != nil {
+		t.Fatalf("removeImages: %v", err)
+	}
+	if removed != 0 {
+		t.Errorf("removed = %d, want 0", removed)
+	}
+	if len(client.images) != 1 {
+		t.Error("the image was deleted for a caller that was already gone")
+	}
+}

pkg/remover/test_client_test.go

@@ -87,8 +87,12 @@ func (c *testClient) removeImageFromSlice(index int) {
 	c.images = s
 }
 
-func (c *testClient) DeleteImage(_ context.Context, image string) (err error) {
+func (c *testClient) DeleteImage(ctx context.Context, image string) (err error) {
 	c.logf("DeleteImage: %s", image)
+	// a real CRI client fails the call rather than deleting anyway
+	if err := ctx.Err(); err != nil {
+		return err
+	}
 	if image == "" {
 		return errImageEmpty
 	}

pkg/scanners/template/scanner_template.go

@@ -88,7 +88,7 @@ func (cfg *config) SendImages(nonCompliantImages, failedImages []unversioned.Ima
 		nonCompliantImages = append(nonCompliantImages, failedImages...)
 	}
 
-	if err := util.WriteScanErasePipe(nonCompliantImages); err != nil {
+	if err := util.WriteImagesPipe(cfg.ctx, util.ScanErasePath, nonCompliantImages); err != nil {
 		cfg.log.Error(err, "unable to write non-compliant images to scan erase pipe")
 		return err
 	}

pkg/utils/handoff_test.go

@@ -2,9 +2,11 @@ package utils
 
 import (
 	"context"
+	"errors"
 	"os"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/eraser-dev/eraser/api/unversioned"
 )
@@ -27,18 +29,33 @@ func shortTempDir(t *testing.T) string {
 	return dir
 }
 
+// testContext bounds the round trips. Both halves of a handoff block until the
+// peer shows up, so on context.Background a rendezvous that never completes
+// hangs until the package-wide test timeout -- ten minutes of nothing, with no
+// indication of which test is stuck. A deadline turns that into a failure in
+// the test that caused it.
+func testContext(t *testing.T) context.Context {
+	t.Helper()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	t.Cleanup(cancel)
+
+	return ctx
+}
+
 func TestImagesHandoffRoundTrip(t *testing.T) {
 	path := filepath.Join(shortTempDir(t), "images")
+	ctx := testContext(t)
 
 	want := []unversioned.Image{
 		{ImageID: "sha256:aaaa", Names: []string{"repo/one:v1"}},
 		{ImageID: "sha256:bbbb", Names: []string{"repo/two:v2"}},
 	}
 
 	errCh := make(chan error, 1)
-	go func() { errCh <- WriteImagesPipe(path, want) }()
+	go func() { errCh <- WriteImagesPipe(ctx, path, want) }()
 
-	got, err := ReadImagesPipe(context.Background(), path)
+	got, err := ReadImagesPipe(ctx, path)
 	if err != nil {
 		t.Fatalf("ReadImagesPipe: %v", err)
 	}
@@ -58,6 +75,7 @@ func TestImagesHandoffRoundTrip(t *testing.T) {
 
 func TestCompletionHandoffRoundTrip(t *testing.T) {
 	path := filepath.Join(shortTempDir(t), "complete")
+	ctx := testContext(t)
 
 	pipe, err := CreateCompletionPipe(path)
 	if err != nil {
@@ -66,18 +84,36 @@ func TestCompletionHandoffRoundTrip(t *testing.T) {
 	defer func() { _ = pipe.Close() }()
 
 	errCh := make(chan error, 1)
-	go func() { errCh <- WriteCompletionPipe(path) }()
+	go func() { errCh <- WriteCompletionPipe(ctx, path) }()
 
-	data, err := pipe.Await()
-	if err != nil {
-		t.Fatalf("Await: %v", err)
+	type awaited struct {
+		data []byte
+		err  error
+	}
+
+	// Await takes no context, so it cannot be given the deadline directly; the
+	// select is what enforces it.
+	awaitCh := make(chan awaited, 1)
+	go func() {
+		data, err := pipe.Await()
+		awaitCh <- awaited{data: data, err: err}
+	}()
+
+	var got awaited
+	select {
+	case got = <-awaitCh:
+	case <-ctx.Done():
+		t.Fatalf("Await did not return within the deadline: %v", ctx.Err())
+	}
+	if got.err != nil {
+		t.Fatalf("Await: %v", got.err)
 	}
 	if err := <-errCh; err != nil {
 		t.Fatalf("WriteCompletionPipe: %v", err)
 	}
 
-	if string(data) != EraseCompleteMessage {
-		t.Errorf("payload = %q, want %q", string(data), EraseCompleteMessage)
+	if string(got.data) != EraseCompleteMessage {
+		t.Errorf("payload = %q, want %q", string(got.data), EraseCompleteMessage)
 	}
 }
 
@@ -86,7 +122,7 @@ func TestCompletionHandoffRoundTrip(t *testing.T) {
 func TestWriteCompletionPipeAbsentPeerIsNotExist(t *testing.T) {
 	path := filepath.Join(shortTempDir(t), "no-such-peer")
 
-	err := WriteCompletionPipe(path)
+	err := WriteCompletionPipe(context.Background(), path)
 	if err == nil {
 		t.Fatal("expected an error writing to an endpoint nobody published")
 	}
@@ -95,6 +131,48 @@ func TestWriteCompletionPipeAbsentPeerIsNotExist(t *testing.T) {
 	}
 }
 
+// The pre-open Stat exists so this precedence holds: an endpoint nobody
+// published reports IsNotExist even when the caller is already shutting down.
+// Left to a select, the two would race and "the scanner is disabled" would
+// become indistinguishable from "we are terminating".
+func TestWriteCompletionPipeAbsentPeerBeatsACanceledContext(t *testing.T) {
+	path := filepath.Join(shortTempDir(t), "no-such-peer")
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err := WriteCompletionPipe(ctx, path)
+	if err == nil {
+		t.Fatal("expected an error writing to an endpoint nobody published")
+	}
+	if !os.IsNotExist(err) {
+		t.Errorf("os.IsNotExist(%v) = false, want true", err)
+	}
+}
+
+// The whole point of taking a context: a worker whose peer never arrives has to
+// be able to give up, on either platform.
+func TestWriteImagesPipeHonoursACanceledContext(t *testing.T) {
+	path := filepath.Join(shortTempDir(t), "never-read")
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	errCh := make(chan error, 1)
+	go func() { errCh <- WriteImagesPipe(ctx, path, []unversioned.Image{{ImageID: "sha256:aaaa"}}) }()
+
+	// nothing ever reads the endpoint, so the write is still waiting
+	cancel()
+
+	select {
+	case err := <-errCh:
+		if !errors.Is(err, context.Canceled) {
+			t.Errorf("WriteImagesPipe = %v, want context.Canceled", err)
+		}
+	case <-time.After(30 * time.Second):
+		t.Fatal("WriteImagesPipe ignored the canceled context")
+	}
+}
+
 func TestCompletionPipeCloseIsIdempotentlySafe(t *testing.T) {
 	path := filepath.Join(shortTempDir(t), "closed")