fix: prevent LogStreamer.Flush deadlock when context is cancelled

src/pkg/logs.go

@@ -21,17 +21,18 @@ type LogStreamer struct {
 	processors []LogProcessor
 	logger     zerolog.Logger
 	quit       chan bool
+	done       chan struct{}
 	logBuffer  *ring.Ring
 }
 
 func NewLogStreamer(logger zerolog.Logger, processors ...LogProcessor) LogStreamer {
-	quit := make(chan bool)
 	return LogStreamer{
 		Stdout:     &SafeBuffer{},
 		Stderr:     &SafeBuffer{},
 		processors: processors,
 		logger:     logger,
-		quit:       quit,
+		quit:       make(chan bool),
+		done:       make(chan struct{}),
 		logBuffer:  ring.New(20),
 	}
 }
@@ -52,17 +53,18 @@ func (s *LogStreamer) streams() []logStream {
 	}
 }
 
-func (s *LogStreamer) processLine(stream logStream) {
+func (s *LogStreamer) processLine(stream logStream) bool {
 	line, _ := stream.buf.ReadString('\n')
 	if line == "" {
-		return
+		return false
 	}
 	line = strings.TrimSuffix(line, "\n")
 	for _, processor := range s.processors {
 		line = stream.fn(processor, line)
 	}
 	s.logBuffer.Value = line
 	s.logBuffer = s.logBuffer.Next()
+	return true
 }
 
 func (s *LogStreamer) GetLogBuffer() []string {
@@ -76,6 +78,7 @@ func (s *LogStreamer) GetLogBuffer() []string {
 }
 
 func (s *LogStreamer) Run(ctx context.Context) {
+	defer close(s.done)
 	s.logger.Trace().Msg("Starting log streamer ...")
 	ticker := time.NewTicker(50 * time.Millisecond)
 	defer ticker.Stop()
@@ -102,22 +105,32 @@ func (s *LogStreamer) Flush(outcome JobOutcome) {
 	ticker := time.NewTicker(200 * time.Millisecond)
 	defer ticker.Stop()
 	timeout := time.After(30 * time.Second)
+wait:
 	for strings.Contains(s.Stderr.String(), "\n") || strings.Contains(s.Stdout.String(), "\n") {
 		select {
 		case <-ticker.C:
 			// Continue waiting
+		case <-s.done:
+			// 'Run' already exited (context cancelled); drain the rest below.
+			break wait
 		case <-timeout:
 			s.logger.Warn().Msg("Flush timeout reached, proceeding with remaining data")
-			goto done
+			break wait
 		}
 	}
-done:
+	// Stop 'Run' if it is still going — it may have already exited via ctx.Done,
+	// in which case there is no receiver for quit and sending would block forever.
+	select {
+	case s.quit <- true:
+	case <-s.done:
+	}
+	<-s.done
 	s.logger.Trace().Msg("Finished log streamer flush ...")
-	s.quit <- true
-	time.Sleep(200 * time.Millisecond) // Allow 'Run' goroutine to quit
-	// Drain any partial line that never received a terminating newline.
+	// Drain anything 'Run' did not process, including a partial line that never
+	// received a terminating newline.
 	for _, stream := range s.streams() {
-		s.processLine(stream)
+		for s.processLine(stream) {
+		}
 	}
 	s.logger.Trace().Msg("Flushing log processors ...")
 	for i := len(s.processors) - 1; i >= 0; i-- {

src/pkg/logs_test.go

@@ -46,6 +46,36 @@ func TestLogStreamerPartialLineStdout(t *testing.T) {
 	autopilot.Equals(t, []string{"partial", "trailing-no-newline"}, cap.lines)
 }
 
+func TestLogStreamerFlushAfterContextCancel(t *testing.T) {
+	cap := &captureProcessor{}
+	s := NewLogStreamer(zerolog.Nop(), cap)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go s.Run(ctx)
+
+	_, _ = s.Stdout.Write([]byte("line-one\nline-two\n"))
+	time.Sleep(100 * time.Millisecond)
+
+	// Cancel the context so 'Run' exits on its own, then write more lines that
+	// only Flush can drain.
+	cancel()
+	<-s.done
+	_, _ = s.Stdout.Write([]byte("late-one\nlate-two"))
+
+	flushed := make(chan struct{})
+	go func() {
+		s.Flush(JobOutcome{})
+		close(flushed)
+	}()
+	select {
+	case <-flushed:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Flush deadlocked after context cancellation")
+	}
+
+	autopilot.Equals(t, []string{"line-one", "line-two", "late-one", "late-two"}, cap.lines)
+}
+
 func TestLogStreamerPartialLineStderr(t *testing.T) {
 	cap := &captureProcessor{}
 	s := NewLogStreamer(zerolog.Nop(), cap)