62 lines
1.2 KiB
Go
62 lines
1.2 KiB
Go
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)
|
|
}
|
|
}
|