improve cache handler, embracing #2210 too

HISTORY.md

@@ -23,6 +23,14 @@ Developers are not forced to upgrade if they don't really need it. Upgrade whene
 
 Changes apply to `main` branch.
 
+- The `cache` sub-package has an update, after 4 years:
+
+    - Add support for custom storage on `cache` package, through the `Handler#Store` method.
+    - Add support for custom expiration duration on `cache` package, trough the `Handler#MaxAge` method.
+    - Improve the overral performance of the `cache` package.
+    - The `cache.Handler` input and output arguments remain as it is.
+    - The `cache.Cache` input argument changed from `time.Duration` to `func(iris.Context) time.Duration`.
+
 # Mon, 25 Sep 2023 | v12.2.7
 
 Minor bug fixes and support of multiple `block` and `define` directives in multiple layouts and templates in the `Blocks` view engine.

_examples/response-writer/cache/simple/main.go

@@ -29,6 +29,10 @@ func main() {
 	app := iris.New()
 	app.Logger().SetLevel("debug")
 	app.Get("/", cache.Handler(10*time.Second), writeMarkdown)
+	// To customize the cache handler:
+	// cache.Cache(nil).MaxAge(func(ctx iris.Context) time.Duration {
+	// 	return time.Duration(ctx.MaxAge()) * time.Second
+	// }).AddRule(...).Store(...)
 	// saves its content on the first request and serves it instead of re-calculating the content.
 	// After 10 seconds it will be cleared and reset.
 

cache/cache.go

@@ -46,8 +46,25 @@ func WithKey(key string) context.Handler {
 	}
 }
 
+// DefaultMaxAge is a function which returns the
+// `context#MaxAge` as time.Duration.
+// It's the default expiration function for the cache handler.
+var DefaultMaxAge = func(ctx *context.Context) time.Duration {
+	return time.Duration(ctx.MaxAge()) * time.Second
+}
+
+// MaxAge is a shortcut to set a simple duration as a MaxAgeFunc.
+//
+// Usage:
+// app.Get("/", cache.Cache(cache.MaxAge(1*time.Minute), mainHandler)
+func MaxAge(dur time.Duration) client.MaxAgeFunc {
+	return func(*context.Context) time.Duration {
+		return dur
+	}
+}
+
 // Cache accepts the cache expiration duration.
-// If the "expiration" input argument is invalid, <=2 seconds,
+// If the "maxAgeFunc" input argument is nil,
 // then expiration is taken by the "cache-control's maxage" header.
 // Returns a Handler structure which you can use to customize cache further.
 //
@@ -57,15 +74,23 @@ func WithKey(key string) context.Handler {
 // may be more suited to your needs.
 //
 // You can add validators with this function.
-func Cache(expiration time.Duration) *client.Handler {
-	return client.NewHandler(expiration)
+func Cache(maxAgeFunc client.MaxAgeFunc) *client.Handler {
+	if maxAgeFunc == nil {
+		maxAgeFunc = DefaultMaxAge
+	}
+
+	return client.NewHandler(maxAgeFunc)
 }
 
 // Handler like `Cache` but returns an Iris Handler to be used as a middleware.
 // For more options use the `Cache`.
 //
 // Examples can be found at: https://github.com/kataras/iris/tree/main/_examples/response-writer/cache
 func Handler(expiration time.Duration) context.Handler {
-	h := Cache(expiration).ServeHTTP
+	maxAgeFunc := func(*context.Context) time.Duration {
+		return expiration
+	}
+
+	h := Cache(maxAgeFunc).ServeHTTP
 	return h
 }

cache/cache_test.go

@@ -162,18 +162,18 @@ func TestCacheValidator(t *testing.T) {
 		ctx.Write([]byte(expectedBodyStr))
 	}
 
-	validCache := cache.Cache(cacheDuration)
-	app.Get("/", validCache.ServeHTTP, h)
+	validCache := cache.Handler(cacheDuration)
+	app.Get("/", validCache, h)
 
-	managedCache := cache.Cache(cacheDuration)
+	managedCache := cache.Cache(cache.MaxAge(cacheDuration))
 	managedCache.AddRule(rule.Validator([]rule.PreValidator{
 		func(ctx *context.Context) bool {
 			// should always invalid for cache, don't bother to go to try to get or set cache
 			return ctx.Request().URL.Path != "/invalid"
 		},
 	}, nil))
 
-	managedCache2 := cache.Cache(cacheDuration)
+	managedCache2 := cache.Cache(cache.MaxAge(cacheDuration))
 	managedCache2.AddRule(rule.Validator(nil,
 		[]rule.PostValidator{
 			func(ctx *context.Context) bool {
@@ -183,7 +183,7 @@ func TestCacheValidator(t *testing.T) {
 		},
 	))
 
-	app.Get("/valid", validCache.ServeHTTP, h)
+	app.Get("/valid", validCache, h)
 
 	app.Get("/invalid", managedCache.ServeHTTP, h)
 	app.Get("/invalid2", managedCache2.ServeHTTP, func(ctx *context.Context) {

cache/client/handler.go

@@ -1,8 +1,8 @@
 package client
 
 import (
+	"net/http"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/kataras/iris/v12/cache/client/rule"
@@ -23,19 +23,23 @@ type Handler struct {
 	// See more at ruleset.go
 	rule rule.Rule
 	// when expires.
-	expiration time.Duration
+	maxAgeFunc MaxAgeFunc
 	// entries the memory cache stored responses.
-	entries map[string]*entry.Entry
-	mu      sync.RWMutex
+	entryPool  *entry.Pool
+	entryStore entry.Store
 }
 
+type MaxAgeFunc func(*context.Context) time.Duration
+
 // NewHandler returns a new cached handler for the "bodyHandler"
 // which expires every "expiration".
-func NewHandler(expiration time.Duration) *Handler {
+func NewHandler(maxAgeFunc MaxAgeFunc) *Handler {
 	return &Handler{
 		rule:       DefaultRuleSet,
-		expiration: expiration,
-		entries:    make(map[string]*entry.Entry),
+		maxAgeFunc: maxAgeFunc,
+
+		entryPool:  entry.NewPool(),
+		entryStore: entry.NewMemStore(),
 	}
 }
 
@@ -64,14 +68,20 @@ func (h *Handler) AddRule(r rule.Rule) *Handler {
 	return h
 }
 
-var emptyHandler = func(ctx *context.Context) {
-	ctx.StopWithText(500, "cache: empty body handler")
+// Store sets a custom store for this handler.
+func (h *Handler) Store(store entry.Store) *Handler {
+	h.entryStore = store
+	return h
 }
 
-func parseLifeChanger(ctx *context.Context) entry.LifeChanger {
-	return func() time.Duration {
-		return time.Duration(ctx.MaxAge()) * time.Second
-	}
+// MaxAge customizes the expiration duration for this handler.
+func (h *Handler) MaxAge(fn MaxAgeFunc) *Handler {
+	h.maxAgeFunc = fn
+	return h
+}
+
+var emptyHandler = func(ctx *context.Context) {
+	ctx.StopWithText(500, "cache: empty body handler")
 }
 
 const entryKeyContextKey = "iris.cache.server.entry.key"
@@ -133,33 +143,10 @@ func (h *Handler) ServeHTTP(ctx *context.Context) {
 		return
 	}
 
-	var (
-		response *entry.Response
-		valid    = false
-		// unique per subdomains and paths with different url query.
-		key = getOrSetKey(ctx)
-	)
-
-	h.mu.RLock()
-	e, found := h.entries[key]
-	h.mu.RUnlock()
-
-	if found {
-		// the entry is here, .Response will give us
-		// if it's expired or no
-		response, valid = e.Response()
-	} else {
-		// create the entry now.
-		// fmt.Printf("create new cache entry\n")
-		// fmt.Printf("key: %s\n", key)
-
-		e = entry.NewEntry(h.expiration)
-		h.mu.Lock()
-		h.entries[key] = e
-		h.mu.Unlock()
-	}
+	key := getOrSetKey(ctx) // unique per subdomains and paths with different url query.
 
-	if !valid {
+	e := h.entryStore.Get(key)
+	if e == nil {
 		// if it's expired, then execute the original handler
 		// with our custom response recorder response writer
 		// because the net/http doesn't give us
@@ -182,30 +169,61 @@ func (h *Handler) ServeHTTP(ctx *context.Context) {
 			return
 		}
 
-		// check for an expiration time if the
-		// given expiration was not valid then check for GetMaxAge &
-		// update the response & release the recorder
-		e.Reset(
-			recorder.StatusCode(),
-			recorder.Header(),
-			body,
-			parseLifeChanger(ctx),
-		)
-
 		// fmt.Printf("reset cache entry\n")
 		// fmt.Printf("key: %s\n", key)
 		// fmt.Printf("content type: %s\n", recorder.Header().Get(cfg.ContentTypeHeader))
 		// fmt.Printf("body len: %d\n", len(body))
+
+		r := entry.NewResponse(recorder.StatusCode(), recorder.Header(), body)
+		e = h.entryPool.Acquire(h.maxAgeFunc(ctx), r, func() {
+			h.entryStore.Delete(key)
+		})
+
+		h.entryStore.Set(key, e)
 		return
 	}
 
 	// if it's valid then just write the cached results
-	entry.CopyHeaders(ctx.ResponseWriter().Header(), response.Headers())
+	r := e.Response()
+	// if !ok {
+	// 	// it shouldn't be happen because if it's not valid (= expired)
+	// 	// then it shouldn't be found on the store, we return as it is, the body was written.
+	// 	return
+	// }
+
+	copyHeaders(ctx.ResponseWriter().Header(), r.Headers())
 	ctx.SetLastModified(e.LastModified)
-	ctx.StatusCode(response.StatusCode())
-	ctx.Write(response.Body())
+	ctx.StatusCode(r.StatusCode())
+	ctx.Write(r.Body())
 
 	// fmt.Printf("key: %s\n", key)
 	// fmt.Printf("write content type: %s\n", response.Headers()["ContentType"])
 	// fmt.Printf("write body len: %d\n", len(response.Body()))
 }
+
+func copyHeaders(dst, src http.Header) {
+	// Clone returns a copy of h or nil if h is nil.
+	if src == nil {
+		return
+	}
+
+	// Find total number of values.
+	nv := 0
+	for _, vv := range src {
+		nv += len(vv)
+	}
+
+	sv := make([]string, nv) // shared backing array for headers' values
+	for k, vv := range src {
+		if vv == nil {
+			// Preserve nil values. ReverseProxy distinguishes
+			// between nil and zero-length header values.
+			dst[k] = nil
+			continue
+		}
+
+		n := copy(sv, vv)
+		dst[k] = sv[:n:n]
+		sv = sv[n:]
+	}
+}