feat: implement RistrettoStorage for improved caching with HTTP middleware

Author: MagicFun1241

main.go

@@ -19,6 +19,7 @@ import (
 	"media-proxy/metrics"
 	fiberprometheus "media-proxy/middlewares/prometheus"
 	"media-proxy/routes"
+	"media-proxy/storage"
 
 	"github.com/dgraph-io/ristretto/v2"
 )
@@ -50,6 +51,13 @@ func main() {
 		BufferItems: 64,      // number of keys per Get buffer.
 	}
 
+	// HTTP cache configuration for middleware
+	httpCacheConfig := &ristretto.Config[string, []byte]{
+		NumCounters: 1e6,     // number of keys to track frequency of (1M).
+		MaxCost:     1 << 28, // maximum cost of cache (256MB).
+		BufferItems: 64,      // number of keys per Get buffer.
+	}
+
 	if config.HTTPCacheTTL == 0 {
 		config.HTTPCacheTTL = 1800 // 30 minutes
 	}
@@ -75,6 +83,12 @@ func main() {
 		logger.Fatal(err.Error())
 	}
 
+	// Create HTTP cache for middleware
+	httpCacheStore, err := ristretto.NewCache(httpCacheConfig)
+	if err != nil {
+		logger.Fatal(err.Error())
+	}
+
 	// Initialize optional S3 cache
 	s3cache, s3err := routes.NewS3Cache(
 		config.S3Enabled,
@@ -135,6 +149,14 @@ func main() {
 	app.Use(etag.New())
 	app.Use(cache.New(cache.Config{
 		Expiration: time.Minute * 10,
+		Storage:    storage.NewRistrettoStorage(httpCacheStore),
+		Next: func(c *fiber.Ctx) bool {
+			if strings.HasPrefix(c.Path(), "/videos/") && !strings.HasPrefix(c.Path(), "/videos/preview/") {
+				return true
+			}
+
+			return false
+		},
 		KeyGenerator: func(c *fiber.Ctx) string {
 			if strings.HasPrefix(c.Path(), "/videos/") && !strings.HasPrefix(c.Path(), "/videos/preview/") {
 				return c.Path() + "?" + string(c.Request().Header.Peek("Range"))

storage/ristretto.go

@@ -0,0 +1,49 @@
+package storage
+
+import (
+	"time"
+
+	"github.com/dgraph-io/ristretto/v2"
+)
+
+// RistrettoStorage implements fiber.Storage using Ristretto cache
+type RistrettoStorage struct {
+	cache *ristretto.Cache[string, []byte]
+}
+
+// NewRistrettoStorage creates a new RistrettoStorage
+func NewRistrettoStorage(cache *ristretto.Cache[string, []byte]) *RistrettoStorage {
+	return &RistrettoStorage{cache: cache}
+}
+
+// Get retrieves data from cache
+func (r *RistrettoStorage) Get(key string) ([]byte, error) {
+	if value, found := r.cache.Get(key); found {
+		return value, nil
+	}
+	return nil, nil
+}
+
+// Set stores data in cache
+func (r *RistrettoStorage) Set(key string, val []byte, exp time.Duration) error {
+	r.cache.SetWithTTL(key, val, 1, exp)
+	return nil
+}
+
+// Delete removes data from cache
+func (r *RistrettoStorage) Delete(key string) error {
+	r.cache.Del(key)
+	return nil
+}
+
+// Reset clears all data from cache
+func (r *RistrettoStorage) Reset() error {
+	r.cache.Clear()
+	return nil
+}
+
+// Close closes the storage (no-op for ristretto)
+func (r *RistrettoStorage) Close() error {
+	r.cache.Close()
+	return nil
+}