feat: add WaitUntilEmpty to LogSender

spikecurtis · spikecurtis · commit f39d152ebea2 · 2024-02-15T16:46:32.000+04:00
diff --git a/codersdk/agentsdk/logs.go b/codersdk/agentsdk/logs.go
@@ -483,6 +483,34 @@ func (l *LogSender) GetScriptLogger(logSourceID uuid.UUID) ScriptLogger {
 	return ScriptLogger{srcID: logSourceID, sender: l}
 }
 
+// WaitUntilEmpty waits until the LogSender's queues are empty or the given context expires.
+func (l *LogSender) WaitUntilEmpty(ctx context.Context) error {
+	ctxDone := false
+	nevermind := make(chan struct{})
+	defer close(nevermind)
+	go func() {
+		select {
+		case <-ctx.Done():
+			l.L.Lock()
+			defer l.L.Unlock()
+			ctxDone = true
+			l.Broadcast()
+			return
+		case <-nevermind:
+			return
+		}
+	}()
+	l.L.Lock()
+	defer l.L.Unlock()
+	for len(l.queues) != 0 && !ctxDone {
+		l.Wait()
+	}
+	if len(l.queues) == 0 {
+		return nil
+	}
+	return ctx.Err()
+}
+
 type ScriptLogger struct {
 	sender *LogSender
 	srcID  uuid.UUID
diff --git a/codersdk/agentsdk/logs_internal_test.go b/codersdk/agentsdk/logs_internal_test.go
@@ -56,6 +56,12 @@ func TestLogSender_Mainline(t *testing.T) {
 		loopErr <- err
 	}()
 
+	empty := make(chan error, 1)
+	go func() {
+		err := uut.WaitUntilEmpty(ctx)
+		empty <- err
+	}()
+
 	// since neither source has even been flushed, it should immediately Flush
 	// both, although the order is not controlled
 	var logReqs []*proto.BatchCreateLogsRequest
@@ -104,8 +110,11 @@ func TestLogSender_Mainline(t *testing.T) {
 	require.Equal(t, proto.Log_DEBUG, req.Logs[0].GetLevel())
 	require.Equal(t, t1, req.Logs[0].GetCreatedAt().AsTime())
 
+	err := testutil.RequireRecvCtx(ctx, t, empty)
+	require.NoError(t, err)
+
 	cancel()
-	err := testutil.RequireRecvCtx(testCtx, t, loopErr)
+	err = testutil.RequireRecvCtx(testCtx, t, loopErr)
 	require.NoError(t, err)
 
 	// we can still enqueue more logs after SendLoop returns
@@ -363,6 +372,33 @@ func TestLogSender_SendError(t *testing.T) {
 	uut.L.Unlock()
 }
 
+func TestLogSender_WaitUntilEmpty_ContextExpired(t *testing.T) {
+	t.Parallel()
+	testCtx := testutil.Context(t, testutil.WaitShort)
+	ctx, cancel := context.WithCancel(testCtx)
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
+	uut := NewLogSender(logger)
+
+	t0 := dbtime.Now()
+
+	ls1 := uuid.UUID{0x11}
+	uut.Enqueue(ls1, Log{
+		CreatedAt: t0,
+		Output:    "test log 0, src 1",
+		Level:     codersdk.LogLevelInfo,
+	})
+
+	empty := make(chan error, 1)
+	go func() {
+		err := uut.WaitUntilEmpty(ctx)
+		empty <- err
+	}()
+
+	cancel()
+	err := testutil.RequireRecvCtx(testCtx, t, empty)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
 type fakeLogDest struct {
 	reqs  chan *proto.BatchCreateLogsRequest
 	resps chan *proto.BatchCreateLogsResponse
