diff --git a/internal/cmd/discoverworkload/discoverworkload.go b/internal/cmd/discoverworkload/discoverworkload.go index 8eb75ee..2f02a89 100644 --- a/internal/cmd/discoverworkload/discoverworkload.go +++ b/internal/cmd/discoverworkload/discoverworkload.go @@ -7,6 +7,8 @@ import ( "io" "log/slog" "os" + "os/signal" + "syscall" "time" "github.com/spf13/cobra" @@ -75,6 +77,8 @@ func NewCommand(ctx context.Context) *cobra.Command { logger.Info("starting to watch for workloads", "duration", cfg.Timeout) + go gracefulShutdown(cancel) + opts := discover.NewManifestJSONProcessorFnOptions{ CompactOutput: cfg.CompactOutput, } @@ -143,3 +147,10 @@ func newLogger(level string, out io.Writer) (*slog.Logger, error) { return logger, nil } + +func gracefulShutdown(cancel context.CancelFunc) { + quit := make(chan os.Signal, 1) + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + <-quit + cancel() +} diff --git a/internal/discover/processor_test.go b/internal/discover/processor_test.go index d2122d7..bdd0532 100644 --- a/internal/discover/processor_test.go +++ b/internal/discover/processor_test.go @@ -394,6 +394,11 @@ func TestManifestJSONProcessor(t *testing.T) { compact bool expected []byte }{ + "no pods": { + ctx: context.TODO(), + input: []corev1.Pod{}, + expected: []byte{}, + }, "initContainer only": { ctx: context.TODO(), input: []corev1.Pod{ @@ -461,7 +466,7 @@ func TestManifestJSONProcessor(t *testing.T) { actual, err := io.ReadAll(buffer) if testFnErr != nil { - t.Fatalf("processor function threw an error unexpectedly: %q", err) + t.Fatalf("processor function threw an error unexpectedly: %q", testFnErr) } if err != nil { @@ -475,6 +480,60 @@ func TestManifestJSONProcessor(t *testing.T) { } } +func TestManifestJSONProcessorContextCancelled(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.TODO()) + input := corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "pod-1"}, + Spec: corev1.PodSpec{ + Containers: []corev1.Container{ + { + Name: "cname", + Image: "example.com/namespace/image:0.0.1", + }, + }, + }, + } + expected := []byte("{\"DiscoveredImages\":[{\"Image\":\"example.com/namespace/image:0.0.1\",\"Containers\":[{\"Name\":\"cname\",\"Type\":\"Container\",\"Pod\":{\"Name\":\"pod-1\",\"Namespace\":\"\"}}]}]}\n") + + testLogger := NewSlogDiscardLogger() + buffer := bytes.NewBuffer([]byte{}) + opts := NewManifestJSONProcessorFnOptions{ + CompactOutput: true, + } + fn := NewManifestJSONProcessorFn(buffer, opts) + + ch := make(chan *corev1.Pod) + + var wg sync.WaitGroup + wg.Add(1) + var testFnErr error + go func() { + defer wg.Done() + testFnErr = fn(ctx, ch, testLogger) + }() + + // If the context is cancelled prematurely, the function should + // still output JSON for what has already been processed. + ch <- &input + cancel() + + wg.Wait() + actual, err := io.ReadAll(buffer) + + if testFnErr != nil { + t.Fatalf("processor function threw an error unexpectedly: %q", testFnErr) + } + + if err != nil { + t.Fatalf("unable to read output written to buffer: %q", err) + } + + if !bytes.Equal(actual, expected) { + t.Fatalf("ManifestJSONPrinter processing function returned the wrong output. actual: %q expected %q", actual, expected) + } +} + func TestContainerProcessing(t *testing.T) { t.Parallel() testcases := map[string]struct {