Add cache to dedupe table (which adds limit) (#177)

feat: add otter.Cache to dedupeHashTable to reduce unbounded growth.

This should prioritize entries that have been retreived.

Author: NGTmeaty

client.go

@@ -6,6 +6,8 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+
+	"github.com/maypok86/otter"
 )
 
 type Error struct {
@@ -44,7 +46,7 @@ type HTTPClientSettings struct {
 type CustomHTTPClient struct {
 	interfacesWatcherStop    chan bool
 	WaitGroup                *WaitGroupWithCount
-	dedupeHashTable          *sync.Map
+	dedupeHashTable          *otter.Cache[string, revisitRecord]
 	ErrChan                  chan *Error
 	WARCWriter               chan *RecordBatch
 	interfacesWatcherStarted chan bool
@@ -60,6 +62,7 @@ type CustomHTTPClient struct {
 	FullOnDisk             bool
 	DigestAlgorithm        DigestAlgorithm
 	closeDNSCache          func()
+	closeDedupeCache       func()
 	// MaxRAMUsageFraction is the fraction of system RAM above which we'll force spooling to disk. For example, 0.5 = 50%.
 	// If set to <= 0, the default value is DefaultMaxRAMUsageFraction.
 	MaxRAMUsageFraction float64
@@ -99,6 +102,7 @@ func (c *CustomHTTPClient) Close() error {
 	}
 
 	c.closeDNSCache()
+	c.closeDedupeCache()
 
 	return nil
 }
@@ -132,7 +136,23 @@ func NewWARCWritingHTTPClient(HTTPClientSettings HTTPClientSettings) (httpClient
 
 	// Toggle deduplication options and create map for deduplication records.
 	httpClient.dedupeOptions = HTTPClientSettings.DedupeOptions
-	httpClient.dedupeHashTable = new(sync.Map)
+
+	// Set default dedupe cache size to 1M entries if not specified
+	dedupeCacheSize := HTTPClientSettings.DedupeOptions.DedupeCacheSize
+	if dedupeCacheSize == 0 {
+		dedupeCacheSize = 1_000_000
+	}
+
+	dedupeCache, err := otter.MustBuilder[string, revisitRecord](dedupeCacheSize).Build()
+	if err != nil {
+		return nil, err
+	}
+	httpClient.dedupeHashTable = &dedupeCache
+
+	httpClient.closeDedupeCache = func() {
+		httpClient.dedupeHashTable.Close()
+		time.Sleep(1 * time.Second)
+	}
 
 	// Set default deduplication threshold to 2048 bytes
 	if httpClient.dedupeOptions.SizeThreshold == 0 {

dedupe.go

@@ -27,6 +27,7 @@ type DedupeOptions struct {
 	DoppelgangerHost   string
 	CDXCookie          string
 	SizeThreshold      int
+	DedupeCacheSize    int
 	LocalDedupe        bool
 	CDXDedupe          bool
 	DoppelgangerDedupe bool
@@ -40,9 +41,9 @@ type revisitRecord struct {
 }
 
 func (d *customDialer) checkLocalRevisit(digest string) revisitRecord {
-	revisit, exists := d.client.dedupeHashTable.Load(digest)
+	revisit, exists := d.client.dedupeHashTable.Get(digest)
 	if exists {
-		return revisit.(revisitRecord)
+		return revisit
 	}
 
 	return revisitRecord{}
@@ -51,7 +52,7 @@ func (d *customDialer) checkLocalRevisit(digest string) revisitRecord {
 func checkCDXRevisit(CDXURL string, digest string, targetURI string, cookie string) (revisitRecord, error) {
 	// CDX expects no hash header. For now we need to strip it.
 	digest = strings.SplitN(digest, ":", 2)[1]
-	
+
 	req, err := http.NewRequest("GET", CDXURL+"/web/timemap/cdx?url="+url.QueryEscape(targetURI)+"&limit=-1", nil)
 	if err != nil {
 		return revisitRecord{}, err
@@ -95,7 +96,7 @@ func checkCDXRevisit(CDXURL string, digest string, targetURI string, cookie stri
 func checkDoppelgangerRevisit(DoppelgangerHost string, digest string, targetURI string) (revisitRecord, error) {
 	// Doppelganger is not expecting a hash header either but this will all be rewritten ... shortly...
 	digest = strings.SplitN(digest, ":", 2)[1]
-	
+
 	req, err := http.NewRequest("GET", DoppelgangerHost+"/api/records/"+digest+"?uri="+targetURI, nil)
 	if err != nil {
 		return revisitRecord{}, err

dialer.go

@@ -611,7 +611,7 @@ func (d *customDialer) writeWARCFromConnection(ctx context.Context, reqPipe, res
 						}
 						return
 					}
-					d.client.dedupeHashTable.Store(r.Header.Get("WARC-Payload-Digest"), revisitRecord{
+					d.client.dedupeHashTable.Set(r.Header.Get("WARC-Payload-Digest"), revisitRecord{
 						responseUUID: recordIDs[i],
 						size:         getContentLength(r.Content),
 						targetURI:    warcTargetURI,