sync from project

middleware/cleanup/cleanup.go

@@ -0,0 +1,61 @@
+package cleanup
+
+import (
+	"context"
+	"log/slog"
+	"net/http"
+	"reflect"
+
+	"git.jeffthecoder.xyz/public/lazyhandler/middleware"
+)
+
+type ctxKey int
+
+const (
+	cleanupCtxKey ctxKey = iota
+)
+
+type CleanupContext struct {
+	funcs []Cleanup
+}
+
+type Cleanup interface {
+	Name() string
+	Cleanup()
+}
+
+type CleanupFunc func()
+
+func (fn CleanupFunc) Name() string {
+	return reflect.TypeOf(fn).String()
+}
+
+func (fn CleanupFunc) Cleanup() {
+	defer func() {
+		if v := recover(); v != nil {
+			slog.With("v", v).Warn("cleanup panicked")
+		}
+	}()
+
+	fn()
+}
+
+func Collect() middleware.Middleware {
+	return middleware.WrapFunc(func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := &CleanupContext{}
+			next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), cleanupCtxKey, ctx)))
+			for _, cleanupFunc := range ctx.funcs {
+				defer cleanupFunc.Cleanup()
+			}
+		})
+	})
+}
+
+// Register adds a Cleanup function to the CleanupContext in the provided context.
+// If the CleanupContext is not found in the context, the Cleanup function is not registered.
+func Register(ctx context.Context, c Cleanup) {
+	if ctx, ok := ctx.Value(cleanupCtxKey).(*CleanupContext); ok {
+		ctx.funcs = append(ctx.funcs, c)
+	}
+}

middleware/cleanup/cleanup_test.go

@@ -0,0 +1,54 @@
+package cleanup
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+)
+
+func TestCleanupMiddleware(t *testing.T) {
+	var executionOrder []string
+
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		Register(r.Context(), CleanupFunc(func() {
+			executionOrder = append(executionOrder, "first")
+		}))
+		Register(r.Context(), CleanupFunc(func() {
+			executionOrder = append(executionOrder, "second")
+		}))
+		w.WriteHeader(http.StatusOK)
+	})
+
+	req := httptest.NewRequest("GET", "/", nil)
+	rr := httptest.NewRecorder()
+
+	middleware := Collect().WrapHandler(handler)
+	middleware.ServeHTTP(rr, req)
+
+	if rr.Code != http.StatusOK {
+		t.Errorf("handler returned wrong status code: got %v want %v",
+			rr.Code, http.StatusOK)
+	}
+
+	expectedOrder := "second,first"
+	actualOrder := strings.Join(executionOrder, ",")
+	if actualOrder != expectedOrder {
+		t.Errorf("cleanup functions executed in wrong order: got %s want %s",
+			actualOrder, expectedOrder)
+	}
+}
+
+func TestCleanupMissingContext(t *testing.T) {
+	// This test ensures that Register does not panic when the context is missing.
+	// The function should fail silently.
+	defer func() {
+		if r := recover(); r != nil {
+			t.Errorf("The code panicked when it should not have")
+		}
+	}()
+
+	req := httptest.NewRequest("GET", "/", nil)
+	// No middleware, so no context
+	Register(req.Context(), CleanupFunc(func() {}))
+}