From 52e116fffc20f5cf159522840b6a82d61ee58c99 Mon Sep 17 00:00:00 2001
From: James Rouzier <rouzier@gmail.com>
Date: Wed, 22 Nov 2023 13:49:24 +0000
Subject: [PATCH] Create a global pool to share contexts between all routes

---
 context.go | 24 ++++++++++++++++++------
 mux.go     | 18 ++++--------------
 2 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/context.go b/context.go
index 21b1384e..4d3b305b 100644
--- a/context.go
+++ b/context.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/http"
 	"strings"
+	"sync"
 )
 
 // URLParam returns the url parameter from a http.Request object.
@@ -39,17 +40,29 @@ var (
 	RouteCtxKey = &contextKey{"RouteContext"}
 )
 
+var contextPool = sync.Pool{
+	New: func() interface{} {
+		return NewRouteContext()
+	},
+}
+
+func GetRouteContext(mx *Mux) *Context {
+	rctx := contextPool.Get().(*Context)
+	rctx.Routes = mx
+	return rctx
+}
+
+func PutRouteContext(rctx *Context) {
+	rctx.Reset()
+	contextPool.Put(rctx)
+}
+
 // Context is the default routing context set on the root node of a
 // request context to track route patterns, URL parameters and
 // an optional routing path.
 type Context struct {
 	Routes Routes
 
-	// parentCtx is the parent of this one, for using Context as a
-	// context.Context directly. This is an optimization that saves
-	// 1 allocation.
-	parentCtx context.Context
-
 	// Routing path/method override used during the route search.
 	// See Mux#routeHTTP method.
 	RoutePath   string
@@ -92,7 +105,6 @@ func (x *Context) Reset() {
 	x.routeParams.Keys = x.routeParams.Keys[:0]
 	x.routeParams.Values = x.routeParams.Values[:0]
 	x.methodNotAllowed = false
-	x.parentCtx = nil
 }
 
 // URLParam returns the corresponding URL parameter value from the request
diff --git a/mux.go b/mux.go
index 735ab232..60a34951 100644
--- a/mux.go
+++ b/mux.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net/http"
 	"strings"
-	"sync"
 )
 
 var _ Router = &Mux{}
@@ -33,9 +32,6 @@ type Mux struct {
 	// to a parent mux
 	parent *Mux
 
-	// Routing context pool
-	pool *sync.Pool
-
 	// Custom route not found handler
 	notFoundHandler http.HandlerFunc
 
@@ -50,10 +46,7 @@ type Mux struct {
 // NewMux returns a newly initialized Mux object that implements the Router
 // interface.
 func NewMux() *Mux {
-	mux := &Mux{tree: &node{}, pool: &sync.Pool{}}
-	mux.pool.New = func() interface{} {
-		return NewRouteContext()
-	}
+	mux := &Mux{tree: &node{}}
 	return mux
 }
 
@@ -78,17 +71,14 @@ func (mx *Mux) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// mx.handler that is comprised of mx.middlewares + mx.routeHTTP.
 	// Once the request is finished, reset the routing context and put it back
 	// into the pool for reuse from another request.
-	rctx = mx.pool.Get().(*Context)
-	rctx.Reset()
-	rctx.Routes = mx
-	rctx.parentCtx = r.Context()
+	rctx = GetRouteContext(mx)
+	defer PutRouteContext(rctx)
 
 	// NOTE: r.WithContext() causes 2 allocations and context.WithValue() causes 1 allocation
 	r = r.WithContext(context.WithValue(r.Context(), RouteCtxKey, rctx))
 
 	// Serve the request and once its done, put the request context back in the sync pool
 	mx.handler.ServeHTTP(w, r)
-	mx.pool.Put(rctx)
 }
 
 // Use appends a middleware handler to the Mux middleware stack.
@@ -243,7 +233,7 @@ func (mx *Mux) With(middlewares ...func(http.Handler) http.Handler) Router {
 	mws = append(mws, middlewares...)
 
 	im := &Mux{
-		pool: mx.pool, inline: true, parent: mx, tree: mx.tree, middlewares: mws,
+		inline: true, parent: mx, tree: mx.tree, middlewares: mws,
 		notFoundHandler: mx.notFoundHandler, methodNotAllowedHandler: mx.methodNotAllowedHandler,
 	}