Ensure the log streamer respects forced shutdown of the agent

agent/log_streamer.go

@@ -66,6 +66,9 @@ type LogStreamer struct {
 
 	// Have we stopped?
 	stopped bool
+
+	// have we been asked to exit immediately?
+	exitImmediately bool
 }
 
 // NewLogStreamer creates a new instance of the log streamer.
@@ -161,7 +164,7 @@ func (ls *LogStreamer) Process(ctx context.Context, output []byte) error {
 }
 
 // Stop stops the streamer.
-func (ls *LogStreamer) Stop() {
+func (ls *LogStreamer) Stop(graceful bool) {
 	ls.processMutex.Lock()
 	if ls.stopped {
 		ls.processMutex.Unlock()
@@ -171,8 +174,13 @@ func (ls *LogStreamer) Stop() {
 	close(ls.queue)
 	ls.processMutex.Unlock()
 
-	ls.logger.Debug("[LogStreamer] Waiting for workers to shut down")
-	ls.workerWG.Wait()
+	if graceful {
+		ls.workerWG.Wait()
+		ls.logger.Info("[LogStreamer] Waiting for workers to shut down and outstanding chunks to be uploaded")
+	} else {
+		ls.exitImmediately = true
+		ls.logger.Warn("[LogStreamer] NOT waiting for outstanding chunks to be uploaded")
+	}
 }
 
 // The actual log streamer worker
@@ -189,6 +197,15 @@ func (ls *LogStreamer) worker(ctx context.Context, id int) {
 	for {
 		setStat("⌚️ Waiting for a chunk")
 
+		if len(ls.queue) > 0 {
+			ls.logger.Debug("[LogStreamer/Worker#%d] Queue length: %d", id, len(ls.queue))
+		}
+
+		if ls.exitImmediately {
+			ls.logger.Warn("[LogStreamer/Worker#%d] Worker is shutting down immediately, Outstanding Queue length: %d", id, len(ls.queue))
+			return
+		}
+
 		// Get the next chunk (pointer) from the queue. This will block
 		// until something is returned.
 		var chunk *api.Chunk
@@ -201,6 +218,9 @@ func (ls *LogStreamer) worker(ctx context.Context, id int) {
 			return
 		}
 
+		// used to simulate a slow upload
+		// time.Sleep(30 * time.Second)
+
 		setStat("📨 Uploading chunk")
 
 		// Upload the chunk

agent/log_streamer_test.go

@@ -46,7 +46,7 @@ func TestLogStreamer(t *testing.T) {
 		t.Errorf("LogStreamer.Process(ctx, %q) = %v", input, err)
 	}
 
-	ls.Stop()
+	ls.Stop(true)
 
 	want := []*api.Chunk{
 		{
@@ -94,3 +94,47 @@ func TestLogStreamer(t *testing.T) {
 		t.Errorf("after Stop: LogStreamer.Process(ctx, %q) err = %v, want %v", input, err, errStreamerStopped)
 	}
 }
+
+func TestLogStreamerExitImmediately(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	logger := logger.NewConsoleLogger(
+		logger.NewTextPrinter(os.Stderr),
+		func(c int) { t.Errorf("exit(%d)", c) },
+	)
+
+	var mu sync.Mutex
+	var got []*api.Chunk
+	callback := func(ctx context.Context, chunk *api.Chunk) error {
+		mu.Lock()
+		got = append(got, chunk)
+		mu.Unlock()
+		return nil
+	}
+
+	ls := NewLogStreamer(logger, callback, LogStreamerConfig{
+		Concurrency:       3,
+		MaxChunkSizeBytes: 10,
+		MaxSizeBytes:      30,
+	})
+
+	if err := ls.Start(ctx); err != nil {
+		t.Fatalf("LogStreamer.Start(ctx) = %v", err)
+	}
+
+	input := "0123456789abcdefghijklmnopqrstuvwxyz!@#$%^&*()" // 46 bytes
+	if err := ls.Process(ctx, []byte(input)); err != nil {
+		t.Errorf("LogStreamer.Process(ctx, %q) = %v", input, err)
+	}
+
+	ls.Stop(false)
+
+	if !ls.exitImmediately {
+		t.Errorf("LogStreamer.Stop(false) did not set exitImmediately")
+	}
+
+	if len(got) > 0 {
+		t.Errorf("LogStreamer.Stop(false) did not exit immediately")
+	}
+}

agent/run_job.go

@@ -355,7 +355,7 @@ func (r *JobRunner) cleanup(ctx context.Context, wg *sync.WaitGroup, exit core.P
 	r.logStreamer.Process(ctx, r.output.ReadAndTruncate())
 
 	// Stop the log streamer. This will block until all the chunks have been uploaded
-	r.logStreamer.Stop()
+	r.logStreamer.Stop(true)
 
 	// Stop the header time streamer. This will block until all the chunks have been uploaded
 	r.headerTimesStreamer.Stop()
@@ -471,6 +471,7 @@ func (r *JobRunner) CancelAndStop() error {
 	r.cancelLock.Lock()
 	r.stopped = true
 	r.cancelLock.Unlock()
+	r.logStreamer.Stop(false)
 	return r.Cancel()
 }