From 7cf206454ddca55df81969980e0616002607994a Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 27 Feb 2024 17:28:13 +1100
Subject: [PATCH] Cache sqlite storage capacity for a short while

---
 request-strategy-impls.go         |  2 +-
 storage/possum/possum-provider.go |  2 +-
 storage/sqlite/direct.go          | 40 ++++++++++++++++++++++++-------
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/request-strategy-impls.go b/request-strategy-impls.go
index 9d779ded..661769f3 100644
--- a/request-strategy-impls.go
+++ b/request-strategy-impls.go
@@ -51,7 +51,7 @@ func (cl *Client) getRequestStrategyInputCommon() requestStrategyInputCommon {
 
 // Returns what is necessary to run request_strategy.GetRequestablePieces for primaryTorrent.
 func (cl *Client) getRequestStrategyInput(primaryTorrent *Torrent) (input request_strategy.Input) {
-	if primaryTorrent.storage.Capacity == nil {
+	if !primaryTorrent.hasStorageCap() {
 		return requestStrategyInputSingleTorrent{
 			requestStrategyInputCommon: cl.getRequestStrategyInputCommon(),
 			t:                          primaryTorrent,
diff --git a/storage/possum/possum-provider.go b/storage/possum/possum-provider.go
index fd7810d3..ebd659e6 100644
--- a/storage/possum/possum-provider.go
+++ b/storage/possum/possum-provider.go
@@ -17,7 +17,7 @@ import (
 )
 
 // Extends possum resource.Provider with an efficient implementation of torrent
-// storage.ConsecutiveChunkReader.
+// storage.ConsecutiveChunkReader. TODO: This doesn't expose Capacity
 type Provider struct {
 	possumResource.Provider
 	Logger log.Logger
diff --git a/storage/sqlite/direct.go b/storage/sqlite/direct.go
index 8e0a4a8d..9758ddb8 100644
--- a/storage/sqlite/direct.go
+++ b/storage/sqlite/direct.go
@@ -5,6 +5,8 @@ package sqliteStorage
 
 import (
 	"io"
+	"sync"
+	"time"
 
 	"github.com/anacrolix/squirrel"
 
@@ -20,26 +22,48 @@ func NewDirectStorage(opts NewDirectStorageOpts) (_ storage.ClientImplCloser, er
 		return
 	}
 	return &client{
-		cache,
-		cache.GetCapacity,
+		cache: cache,
 	}, nil
 }
 
+// Creates a storage.ClientImpl from a provided squirrel.Cache. The caller is responsible for
+// closing the squirrel.Cache.
 func NewWrappingClient(cache *squirrel.Cache) storage.ClientImpl {
 	return &client{
-		cache,
-		cache.GetCapacity,
+		cache: cache,
 	}
 }
 
 type client struct {
-	*squirrel.Cache
-	capacity func() (int64, bool)
+	cache           *squirrel.Cache
+	capacityMu      sync.Mutex
+	capacityFetched time.Time
+	capacityCap     int64
+	capacityCapped  bool
+}
+
+func (c *client) Close() error {
+	return c.cache.Close()
+}
+
+func (c *client) capacity() (cap int64, capped bool) {
+	c.capacityMu.Lock()
+	defer c.capacityMu.Unlock()
+	if !c.capacityFetched.IsZero() && time.Since(c.capacityFetched) < time.Second {
+		cap, capped = c.capacityCap, c.capacityCapped
+		return
+	}
+	c.capacityCap, c.capacityCapped = c.cache.GetCapacity()
+	// Should this go before or after the capacityCap and capacityCapped assignments?
+	c.capacityFetched = time.Now()
+	cap, capped = c.capacityCap, c.capacityCapped
+	return
 }
 
 func (c *client) OpenTorrent(*metainfo.Info, metainfo.Hash) (storage.TorrentImpl, error) {
-	t := torrent{c.Cache}
-	return storage.TorrentImpl{Piece: t.Piece, Close: t.Close, Capacity: &c.capacity}, nil
+	t := torrent{c.cache}
+	capFunc := c.capacity
+	return storage.TorrentImpl{Piece: t.Piece, Close: t.Close, Capacity: &capFunc}, nil
 }
 
 type torrent struct {
-- 
2.44.0