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)
+	}
+}