fix(log): sendLogsV2: cancel LogSender.SendLoop when doneFunc is called

cmd/envbuilder/main.go

@@ -59,6 +59,7 @@ func envbuilderCmd() serpent.Command {
 					preExec = append(preExec, func() {
 						o.Logger(log.LevelInfo, "Closing logs")
 						closeLogs()
+						o.Logger(log.LevelInfo, "Closed logs")
 					})
 					// This adds the envbuilder subsystem.
 					// If telemetry is enabled in a Coder deployment,

integration/integration_test.go

@@ -110,8 +110,9 @@ func TestLogs(t *testing.T) {
 			"Dockerfile": fmt.Sprintf(`FROM %s`, testImageUbuntu),
 		},
 	})
-	_, err := runEnvbuilder(t, runOpts{env: []string{
+	ctr, err := runEnvbuilder(t, runOpts{env: []string{
 		envbuilderEnv("GIT_URL", srv.URL),
+		envbuilderEnv("INIT_SCRIPT", "date > /.date.txt"),
 		"CODER_AGENT_URL=" + logSrv.URL,
 		"CODER_AGENT_TOKEN=" + token,
 	}})
@@ -123,6 +124,8 @@ func TestLogs(t *testing.T) {
 		t.Fatal("timed out waiting for logs")
 	case <-logsDone:
 	}
+	output := execContainer(t, ctr, "cat /date.txt")
+	require.NotEmpty(t, strings.TrimSpace(output))
 }
 
 func TestInitScriptInitCommand(t *testing.T) {

log/coder.go

@@ -105,10 +105,15 @@ func sendLogsV1(ctx context.Context, client *agentsdk.Client, l slog.Logger) (Fu
 				Level:     codersdk.LogLevel(lvl),
 			}
 			if err := sendLogs(ctx, log); err != nil {
-				l.Warn(ctx, "failed to send logs to Coder", slog.Error(err))
+				if !errors.Is(err, context.Canceled) {
+					l.Warn(ctx, "failed to send logs to Coder", slog.Error(err))
+				}
 			}
 		}, func() {
-			if err := flushLogs(ctx); err != nil {
+			// Wait for up to 10 seconds for logs to finish sending.
+			sendCtx, sendCancel := context.WithTimeout(context.Background(), logSendGracePeriod)
+			defer sendCancel()
+			if err := flushLogs(sendCtx); err != nil {
 				l.Warn(ctx, "failed to flush logs", slog.Error(err))
 			}
 		}
@@ -118,9 +123,11 @@ func sendLogsV1(ctx context.Context, client *agentsdk.Client, l slog.Logger) (Fu
 func sendLogsV2(ctx context.Context, dest agentsdk.LogDest, ls coderLogSender, l slog.Logger) (Func, func()) {
 	done := make(chan struct{})
 	uid := uuid.New()
+	sendLoopCtx, cancelSendLoop := context.WithCancel(ctx)
+	defer cancelSendLoop()
 	go func() {
 		defer close(done)
-		if err := ls.SendLoop(ctx, dest); err != nil {
+		if err := ls.SendLoop(sendLoopCtx, dest); err != nil {
 			if !errors.Is(err, context.Canceled) {
 				l.Warn(ctx, "failed to send logs to Coder", slog.Error(err))
 			}
@@ -144,14 +151,20 @@ func sendLogsV2(ctx context.Context, dest agentsdk.LogDest, ls coderLogSender, l
 	}()
 
 	logFunc := func(l Level, msg string, args ...any) {
-		ls.Enqueue(uid, agentsdk.Log{
-			CreatedAt: time.Now(),
-			Output:    fmt.Sprintf(msg, args...),
-			Level:     codersdk.LogLevel(l),
-		})
+		select {
+		case <-sendLoopCtx.Done():
+			return
+		default:
+			ls.Enqueue(uid, agentsdk.Log{
+				CreatedAt: time.Now(),
+				Output:    fmt.Sprintf(msg, args...),
+				Level:     codersdk.LogLevel(l),
+			})
+		}
 	}
 
 	doneFunc := func() {
+		cancelSendLoop()
 		<-done
 	}