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,8 +88,60 @@ func main() {
 		os.Exit(1)
 	}
 
-	if err := util.WriteImagesPipe(path, finalImages); err != nil {
-		log.Error(err, "failed to send images", "pipeFile", path)
+	// Registering a 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.
+	//
+	// signal.Notify rather than NotifyContext, because deregistering has to stay
+	// distinguishable from receiving a signal, and NotifyContext's stop func
+	// cancels the very context a signal would.
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	signaled := make(chan struct{})
+	watching := make(chan struct{})
+	go func() {
+		defer close(watching)
+		select {
+		case <-sigCh:
+			close(signaled)
+			cancel()
+		case <-ctx.Done():
+		}
+	}()
+
+	writeErr := util.WriteImagesPipe(ctx, path, finalImages)
+
+	// Deregister, then join the watcher. Afterwards a signal has either closed
+	// signaled or is still sitting in the buffer, and there is no third
+	// outcome -- previously one that landed between the check and the stop was
+	// consumed by the handler and lost, leaving the collector in an Await it
+	// could not be interrupted out of.
+	signal.Stop(sigCh)
+	cancel()
+	<-watching
+
+	terminating := false
+	select {
+	case <-signaled:
+		terminating = true
+	default:
+		select {
+		case <-sigCh:
+			terminating = true
+		default:
+		}
+	}
+
+	if writeErr != nil {
+		log.Error(writeErr, "failed to send images", "pipeFile", path)
+		os.Exit(1)
+	}
+
+	if terminating {
+		log.Error(context.Canceled, "terminating before waiting for completion")
 		os.Exit(1)
 	}
 

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
 	}