diff --git a/context.go b/context.go new file mode 100644 index 0000000..7fa97c6 --- /dev/null +++ b/context.go @@ -0,0 +1,30 @@ +package main + +import ( + "context" +) + +// RouteCtxKey is a private key for the routing context. +type RouteCtxKey struct{} + +// Context represents the routing context. +type Context struct { + //... existing fields + refCount int +} + +// RouteContext returns the routing context from the given context. +func RouteContext(ctx context.Context) *Context { + if ctx == nil { + return &Context{} + } + if rc, ok := ctx.Value(RouteCtxKey{}).(*Context); ok { + return rc.CloneContext() + } + return AcquireContext() +} + +// WithRouteContext returns a new context with the provided routing context. +func WithRouteContext(parent context.Context, routeCtx *Context) context.Context { + return context.WithValue(parent, RouteCtxKey{}, routeCtx.CloneContext()) +} diff --git a/context_pool.go b/context_pool.go new file mode 100644 index 0000000..9a9a4e7 --- /dev/null +++ b/context_pool.go @@ -0,0 +1,38 @@ +package main + +import ( + "sync" +) + +var contextPool = sync.Pool{ + New: func() interface{} { + return &Context{refCount: 0} + }, +} + +// Context represents the routing context. +type Context struct { + //... existing fields + refCount int +} + +// AcquireContext acquires a context from the pool. +func AcquireContext() *Context { + ctx := contextPool.Get().(*Context) + ctx.refCount = 1 + return ctx +} + +// ReleaseContext releases the context back to the pool if it is not referenced. +func (c *Context) ReleaseContext() { + c.refCount-- + if c.refCount == 0 { + contextPool.Put(c) + } +} + +// CloneContext clones the context and increments the reference count. +func (c *Context) CloneContext() *Context { + c.refCount++ + return c +} diff --git a/middleware.go b/middleware.go new file mode 100644 index 0000000..283e072 --- /dev/null +++ b/middleware.go @@ -0,0 +1,23 @@ +package main + +import ( + "net/http" + "github.com/go-chi/chi/v5" +) + +// PreserveRouteContextMiddleware is a middleware that preserves the routing context when the request is cloned. +func PreserveRouteContextMiddleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Get the current routing context + routeCtx := chi.RouteContext(r.Context()) + + // Create a new context with the preserved routing context + newCtx := chi.WithRouteContext(r.Context(), routeCtx) + + // Clone the request with the new context + newR := r.WithContext(newCtx) + + // Call the next handler with the cloned request + next.ServeHTTP(w, newR) + }) +} diff --git a/router.go b/router.go new file mode 100644 index 0000000..6c411ac --- /dev/null +++ b/router.go @@ -0,0 +1,27 @@ +package main + +import ( + "net/http" + "github.com/go-chi/chi/v5" +) + +// NewRouter creates a new router with the necessary middlewares. +func NewRouter() *chi.Mux { + r := chi.NewRouter() + + // Apply the PreserveRouteContextMiddleware globally + r.Use(PreserveRouteContextMiddleware) + + // Add your routes here + r.Get("/", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("Welcome!")) + }) + + r.Route("/sub", func(r chi.Router) { + r.Get("/", func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("Sub route!")) + }) + }) + + return r +}