From fdfcfcc0708a39b06825e158e9bfa96b30cf42aa Mon Sep 17 00:00:00 2001
From: Henrique Dias <mail@hacdias.com>
Date: Thu, 25 Jan 2024 11:37:58 +0100
Subject: [PATCH] namesys: add WithMaxCacheTTL

---
 CHANGELOG.md             |  1 +
 namesys/namesys.go       | 19 +++++++++++++++++--
 namesys/namesys_cache.go | 10 +++++++++-
 namesys/namesys_test.go  | 38 ++++++++++++++++++++++++++++----------
 4 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 773cd1807..d0a64b4e7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,7 @@ The following emojis are used to highlight certain changes:
 - `blockservice` now has `ContextWithSession` and `EmbedSessionInContext` functions, which allows to embed a session in a context. Future calls to `BlockGetter.GetBlock`, `BlockGetter.GetBlocks` and `NewSession` will use the session in the context.
 - `blockservice.NewWritethrough` deprecated function has been removed, instead you can do `blockservice.New(..., ..., WriteThrough())` like previously.
 - `gateway`: a new header configuration middleware has been added to replace the existing header configuration, which can be used more generically.
+- `namesys` now has a `WithMaxCacheTTL` option, which allows you to define a maximum TTL that will be used for caching IPNS entries.
 
 ### Changed
 
diff --git a/namesys/namesys.go b/namesys/namesys.go
index 00b1f4d2d..7928dde40 100644
--- a/namesys/namesys.go
+++ b/namesys/namesys.go
@@ -48,8 +48,9 @@ type namesys struct {
 	dnsResolver, ipnsResolver resolver
 	ipnsPublisher             Publisher
 
-	staticMap map[string]*cacheEntry
-	cache     *lru.Cache[string, cacheEntry]
+	staticMap   map[string]*cacheEntry
+	cache       *lru.Cache[string, cacheEntry]
+	maxCacheTTL *time.Duration
 }
 
 var _ NameSystem = &namesys{}
@@ -73,6 +74,20 @@ func WithCache(size int) Option {
 	}
 }
 
+// WithMaxCacheTTL configures the maximum cache TTL. By default, if the cache is
+// enabled, the entry TTL will be used for caching. By setting this option, you
+// can limit how long that TTL is.
+//
+// For example, if you configure a maximum cache TTL of 1 minute:
+//   - Entry TTL is 5 minutes -> Cache TTL is 1 minute
+//   - Entry TTL is 30 seconds -> Cache TTL is 30 seconds
+func WithMaxCacheTTL(dur time.Duration) Option {
+	return func(n *namesys) error {
+		n.maxCacheTTL = &dur
+		return nil
+	}
+}
+
 // WithDNSResolver is an option that supplies a custom DNS resolver to use instead
 // of the system default.
 func WithDNSResolver(rslv madns.BasicResolver) Option {
diff --git a/namesys/namesys_cache.go b/namesys/namesys_cache.go
index fc8842e3b..531220b23 100644
--- a/namesys/namesys_cache.go
+++ b/namesys/namesys_cache.go
@@ -60,12 +60,20 @@ func (ns *namesys) cacheSet(name string, val path.Path, ttl time.Duration, lastM
 		}
 	}
 
+	// The cache TTL is capped at the configured maxCacheTTL. If not
+	// configured, the entry TTL will always be used.
+	cacheTTL := ttl
+	if ns.maxCacheTTL != nil && cacheTTL > *ns.maxCacheTTL {
+		cacheTTL = *ns.maxCacheTTL
+	}
+	cacheEOL := time.Now().Add(cacheTTL)
+
 	// Add automatically evicts previous entry, so it works for updating.
 	ns.cache.Add(name, cacheEntry{
 		val:      val,
 		ttl:      ttl,
 		lastMod:  lastMod,
-		cacheEOL: time.Now().Add(ttl),
+		cacheEOL: cacheEOL,
 	})
 }
 
diff --git a/namesys/namesys_test.go b/namesys/namesys_test.go
index 41fa0ce88..48310f025 100644
--- a/namesys/namesys_test.go
+++ b/namesys/namesys_test.go
@@ -146,20 +146,38 @@ func TestPublishWithTTL(t *testing.T) {
 		"pk":   record.PublicKeyValidator{},
 	})
 
-	ns, err := NewNameSystem(routing, WithDatastore(dst), WithCache(128))
-	require.NoError(t, err)
-
 	// CID is arbitrary.
 	p, err := path.NewPath("/ipfs/QmUNLLsPACCz1vLxQVkXqqLX5R1X345qqfHbsf67hvA3Nn")
 	require.NoError(t, err)
 
-	ttl := 1 * time.Second
-	eol := time.Now().Add(2 * time.Second)
+	ttl := 5 * time.Minute
+	eol := time.Now().Add(time.Hour)
 
-	err = ns.Publish(context.Background(), priv, p, PublishWithEOL(eol), PublishWithTTL(ttl))
-	require.NoError(t, err)
+	t.Run("Without MaxCacheTTL", func(t *testing.T) {
+		ns, err := NewNameSystem(routing, WithDatastore(dst), WithCache(128))
+		require.NoError(t, err)
+
+		err = ns.Publish(context.Background(), priv, p, PublishWithEOL(eol), PublishWithTTL(ttl))
+		require.NoError(t, err)
+
+		entry, ok := ns.(*namesys).cache.Get(ipns.NameFromPeer(pid).String())
+		require.True(t, ok)
+		require.Equal(t, ttl, entry.ttl)
+		require.LessOrEqual(t, time.Until(entry.cacheEOL), ttl)
+	})
+
+	t.Run("With MaxCacheTTL", func(t *testing.T) {
+		cacheTTL := 30 * time.Second
+
+		ns, err := NewNameSystem(routing, WithDatastore(dst), WithCache(128), WithMaxCacheTTL(cacheTTL))
+		require.NoError(t, err)
 
-	entry, ok := ns.(*namesys).cache.Get(ipns.NameFromPeer(pid).String())
-	require.True(t, ok)
-	require.LessOrEqual(t, entry.cacheEOL.Sub(eol), 10*time.Millisecond)
+		err = ns.Publish(context.Background(), priv, p, PublishWithEOL(eol), PublishWithTTL(ttl))
+		require.NoError(t, err)
+
+		entry, ok := ns.(*namesys).cache.Get(ipns.NameFromPeer(pid).String())
+		require.True(t, ok)
+		require.Equal(t, ttl, entry.ttl)
+		require.LessOrEqual(t, time.Until(entry.cacheEOL), cacheTTL)
+	})
 }