From bbc20f2c259fba4ad6c2536b878bbb51b4326d1a Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Fri, 26 May 2023 03:24:04 +0800
Subject: [PATCH] add

---
 web/Web.go | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/web/Web.go b/web/Web.go
index ff168a2..65c4c71 100644
--- a/web/Web.go
+++ b/web/Web.go
@@ -10,6 +10,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/dustin/go-humanize"
 	psync "github.com/qydysky/part/sync"
 	sys "github.com/qydysky/part/sys"
 )
@@ -229,10 +230,50 @@ type Cache struct {
 	gcL atomic.Int64
 }
 
+type cacheRes struct {
+	headerf      func() http.Header
+	writef       func([]byte) (int, error)
+	writeHeaderf func(statusCode int)
+}
+
+func (t cacheRes) Header() http.Header {
+	return t.headerf()
+}
+func (t cacheRes) Write(b []byte) (int, error) {
+	return t.writef(b)
+}
+func (t cacheRes) WriteHeader(statusCode int) {
+	t.writeHeaderf(statusCode)
+}
+
 func (t *Cache) IsCache(key string) (res *[]byte, isCache bool) {
 	return t.g.Load(key)
 }
 
+func (t *Cache) Cache(key string, aliveDur time.Duration, w http.ResponseWriter) http.ResponseWriter {
+	var (
+		res    cacheRes
+		called atomic.Bool
+	)
+	res.headerf = w.Header
+	res.writeHeaderf = w.WriteHeader
+	res.writef = func(b []byte) (int, error) {
+		if called.CompareAndSwap(false, true) {
+			if len(b) < humanize.MByte {
+				t.g.Store(key, &b, aliveDur)
+			}
+		} else {
+			panic("Cache Write called")
+		}
+		return w.Write(b)
+	}
+	if s := int64(t.g.Len()); s > 10 && t.gcL.Load() <= s {
+		t.gcL.Store(s * 2)
+		t.g.GC()
+	}
+	return res
+}
+
 func (t *Cache) Store(key string, aliveDur time.Duration, data *[]byte) {
 	t.g.Store(key, data, aliveDur)
 	if s := int64(t.g.Len()); s > 10 && t.gcL.Load() <= s {
-- 
2.39.2