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

@@ -39,6 +39,11 @@ const (
 func main() {
 	flag.Parse()
 
+	// A terminating pod should not leave the worker blocked on a peer that is
+	// never going to arrive. The stop func is discarded rather than deferred
+	// because every exit path here is os.Exit, which would skip it anyway.
+	ctx, _ := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+
 	if *enableProfile {
 		go func() {
 			server := &http.Server{
@@ -74,7 +79,7 @@ func main() {
 	}
 
 	if *imageListPtr == "" {
-		nonCompliantImages, err := util.ReadImagesPipe(context.Background(), util.ScanErasePath)
+		nonCompliantImages, err := util.ReadImagesPipe(ctx, util.ScanErasePath)
 		if err != nil {
 			log.Error(err, "error reading non-compliant images")
 			os.Exit(generalErr)
@@ -106,33 +111,30 @@ func main() {
 		log.Info("no images to exclude")
 	}
 
-	removed, err := removeImages(client, imagelist)
+	removed, err := removeImages(ctx, client, imagelist)
 	if err != nil {
 		log.Error(err, "failed to remove images")
 		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")
 			}

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"
 )
@@ -36,7 +38,7 @@ func TestImagesHandoffRoundTrip(t *testing.T) {
 	}
 
 	errCh := make(chan error, 1)
-	go func() { errCh <- WriteImagesPipe(path, want) }()
+	go func() { errCh <- WriteImagesPipe(context.Background(), path, want) }()
 
 	got, err := ReadImagesPipe(context.Background(), path)
 	if err != nil {
@@ -66,7 +68,7 @@ func TestCompletionHandoffRoundTrip(t *testing.T) {
 	defer func() { _ = pipe.Close() }()
 
 	errCh := make(chan error, 1)
-	go func() { errCh <- WriteCompletionPipe(path) }()
+	go func() { errCh <- WriteCompletionPipe(context.Background(), path) }()
 
 	data, err := pipe.Await()
 	if err != nil {
@@ -86,7 +88,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 +97,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")
 

pkg/utils/handoff_unix.go

@@ -68,8 +68,9 @@ func (p *CompletionPipe) Close() error {
 	return nil
 }
 
-// WriteImagesPipe publishes the endpoint and blocks until the reader connects.
-func WriteImagesPipe(path string, images []unversioned.Image) error {
+// WriteImagesPipe publishes the endpoint and blocks until the reader connects,
+// or until ctx is done.
+func WriteImagesPipe(ctx context.Context, path string, images []unversioned.Image) error {
 	data, err := json.Marshal(images)
 	if err != nil {
 		return err
@@ -79,20 +80,85 @@ func WriteImagesPipe(path string, images []unversioned.Image) error {
 		return err
 	}
 
-	//nolint:gosec // G304: Opening pipe file is intended functionality
-	file, err := os.OpenFile(path, os.O_WRONLY, 0)
+	file, err := openForWrite(ctx, path)
 	if err != nil {
 		return err
 	}
 
-	_, err = file.Write(data)
+	return writeAndClose(ctx, file, data)
+}
+
+// writeAndClose writes the payload and closes, which is what frames the message
+// for the reader. The open is not the only place this can block: once the pipe
+// buffer fills, a reader that stops draining blocks the write too, so the
+// watcher closes the file to unblock it.
+func writeAndClose(ctx context.Context, file *os.File, payload []byte) error {
+	done := make(chan struct{})
+	defer close(done)
+
+	go func() {
+		select {
+		case <-ctx.Done():
+			_ = file.Close()
+		case <-done:
+		}
+	}()
+
+	_, err := file.Write(payload)
+
+	// The watcher may already have closed the file, which is what surfaced as the
+	// write error, so the context is checked before the error is trusted.
+	if ctxErr := ctx.Err(); ctxErr != nil {
+		_ = file.Close()
+		return ctxErr
+	}
+
 	if closeErr := file.Close(); closeErr != nil && err == nil {
 		err = closeErr
 	}
 
 	return err
 }
 
+// openForWrite opens a FIFO for writing, which blocks in the kernel until a
+// reader arrives. The open itself is left untouched -- the rendezvous, and the
+// behavior every existing deployment depends on, is exactly as before. Only
+// the waiting is made interruptible, by doing it on a goroutine that hands the
+// file over if anyone is still listening and closes it if not.
+func openForWrite(ctx context.Context, path string) (*os.File, error) {
+	type opened struct {
+		file *os.File
+		err  error
+	}
+
+	// Unbuffered, and paired with abandoned rather than a default case. A
+	// buffered channel would accept the file after the caller had already
+	// returned, orphaning the descriptor; a default case would close a file the
+	// caller was about to ask for, if the open won the race to this select.
+	ch := make(chan opened)
+	abandoned := make(chan struct{})
+
+	go func() {
+		//nolint:gosec // G304: Opening pipe file is intended functionality
+		file, err := os.OpenFile(path, os.O_WRONLY, 0)
+		select {
+		case ch <- opened{file: file, err: err}:
+		case <-abandoned:
+			if file != nil {
+				_ = file.Close()
+			}
+		}
+	}()
+
+	select {
+	case <-ctx.Done():
+		close(abandoned)
+		return nil, ctx.Err()
+	case o := <-ch:
+		return o.file, o.err
+	}
+}
+
 // ReadImagesPipe waits for the endpoint to appear, then reads until the writer
 // finishes. It returns ctx.Err() if the context is canceled while waiting.
 func ReadImagesPipe(ctx context.Context, path string) ([]unversioned.Image, error) {
@@ -143,17 +209,18 @@ func ReadImagesPipe(ctx context.Context, path string) ([]unversioned.Image, erro
 // WriteCompletionPipe signals a peer that this stage is done. The returned error
 // satisfies os.IsNotExist when the peer never published the endpoint, which is
 // how an absent scanner is detected.
-func WriteCompletionPipe(path string) error {
-	//nolint:gosec // G304: Opening pipe file is intended functionality
-	file, err := os.OpenFile(path, os.O_WRONLY, 0)
-	if err != nil {
+func WriteCompletionPipe(ctx context.Context, path string) error {
+	// Checked before the open so that an absent peer is reported as such even
+	// when ctx is already done; otherwise a terminating remover could mistake a
+	// disabled scanner for a cancellation, and vice versa.
+	if _, err := os.Stat(path); err != nil {
 		return err
 	}
 
-	_, err = file.WriteString(EraseCompleteMessage)
-	if closeErr := file.Close(); closeErr != nil && err == nil {
-		err = closeErr
+	file, err := openForWrite(ctx, path)
+	if err != nil {
+		return err
 	}
 
-	return err
+	return writeAndClose(ctx, file, []byte(EraseCompleteMessage))
 }