refactor tracing status writer done func

mafredri · mafredri · commit 791d763fa544 · 2025-03-05T22:36:23.000Z
diff --git a/coderd/tracing/status_writer.go b/coderd/tracing/status_writer.go
@@ -27,7 +27,7 @@ type StatusWriter struct {
 	http.ResponseWriter
 	Status       int
 	Hijacked     bool
-	Done         []func() // If non-nil, this function will be called when the handler is done.
+	doneFuncs    []func() // If non-nil, this function will be called when the handler is done.
 	responseBody []byte
 
 	wroteHeader      bool
@@ -38,12 +38,17 @@ func StatusWriterMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
 		sw := &StatusWriter{ResponseWriter: rw}
 		next.ServeHTTP(sw, r)
-		for _, done := range sw.Done {
+		for _, done := range sw.doneFuncs {
 			done()
 		}
 	})
 }
 
+func (w *StatusWriter) AddDoneFunc(f func()) {
+	// Prepend, as if deferred.
+	w.doneFuncs = append([]func(){f}, w.doneFuncs...)
+}
+
 func (w *StatusWriter) WriteHeader(status int) {
 	if buildinfo.IsDev() || flag.Lookup("test.v") != nil {
 		if w.wroteHeader {
diff --git a/coderd/tracing/status_writer_test.go b/coderd/tracing/status_writer_test.go
@@ -116,6 +116,27 @@ func TestStatusWriter(t *testing.T) {
 		require.Error(t, err)
 		require.Equal(t, "hijacked", err.Error())
 	})
+
+	t.Run("Middleware", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			sw   *tracing.StatusWriter
+			done = false
+			rr   = httptest.NewRecorder()
+		)
+		tracing.StatusWriterMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			sw = w.(*tracing.StatusWriter)
+			sw.AddDoneFunc(func() {
+				done = true
+			})
+			w.WriteHeader(http.StatusNoContent)
+		})).ServeHTTP(rr, httptest.NewRequest("GET", "/", nil))
+
+		require.Equal(t, http.StatusNoContent, rr.Code, "rr status code not set")
+		require.Equal(t, http.StatusNoContent, sw.Status, "sw status code not set")
+		require.True(t, done, "done func not called")
+	})
 }
 
 type hijacker struct {
diff --git a/coderd/workspaceapps/db.go b/coderd/workspaceapps/db.go
@@ -388,7 +388,7 @@ func (p *DBTokenProvider) auditInitAutocommitRequest(ctx context.Context, w http
 
 	// Set the commit function on the status writer to create an audit
 	// log, this ensures that the status and response body are available.
-	sw.Done = append(sw.Done, func() {
+	sw.AddDoneFunc(func() {
 		p.Logger.Info(ctx, "workspace app audit session", slog.F("status", sw.Status), slog.F("body", string(sw.ResponseBody())), slog.F("api_key", aReq.apiKey), slog.F("db_req", aReq.dbReq))
 
 		if sw.Status == http.StatusSeeOther {
