From bf5370b31f94e9847e206704b6e2d7fc8e86ac59 Mon Sep 17 00:00:00 2001
From: Martin Hutchinson <mhutchinson@gmail.com>
Date: Tue, 7 May 2024 17:42:15 +0100
Subject: [PATCH] Added a local cache for each reader (#133)

This caches only the last tile that was fetched. This allows the full
reader in particular to act like a real client, which will consume all
of the leaves from a tile they downloaded, instead of requesting the
same tile for every leaf in it.
---
 hammer/clients.go | 25 +++++++++++++++++++++++++
 hammer/hammer.go  | 17 +++++++++++------
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/hammer/clients.go b/hammer/clients.go
index 82fdad7..bb45d93 100644
--- a/hammer/clients.go
+++ b/hammer/clients.go
@@ -59,6 +59,7 @@ type LeafReader struct {
 	throttle   <-chan bool
 	errchan    chan<- error
 	cancel     func()
+	c          leafBundleCache
 }
 
 // Run runs the log reader. This should be called in a goroutine.
@@ -94,6 +95,10 @@ func (r *LeafReader) getLeaf(ctx context.Context, i uint64, logSize uint64) ([]b
 	if i >= logSize {
 		return nil, fmt.Errorf("requested leaf %d >= log size %d", i, logSize)
 	}
+	if cached := r.c.get(i); cached != nil {
+		klog.V(2).Infof("Using cached result for index %d", i)
+		return cached, nil
+	}
 	bi := i / uint64(r.bundleSize)
 	br := uint64(0)
 	// Check for partial leaf bundle
@@ -115,6 +120,10 @@ func (r *LeafReader) getLeaf(ctx context.Context, i uint64, logSize uint64) ([]b
 	if l := len(bs); uint64(l) <= br {
 		return nil, fmt.Errorf("huh, short leaf bundle with %d entries, want %d", l, br)
 	}
+	r.c = leafBundleCache{
+		start:  bi * uint64(r.bundleSize),
+		leaves: bs,
+	}
 
 	return base64.StdEncoding.DecodeString(string(bs[br]))
 }
@@ -127,6 +136,22 @@ func (r *LeafReader) Kill() {
 	}
 }
 
+// leafBundleCache stores the results of the last fetched tile. This allows
+// readers that read contiguous blocks of leaves to act more like real
+// clients and fetch a tile of 256 leaves once, instead of 256 times.
+type leafBundleCache struct {
+	start  uint64
+	leaves [][]byte
+}
+
+func (tc leafBundleCache) get(i uint64) []byte {
+	end := tc.start + uint64(len(tc.leaves))
+	if i >= tc.start && i < end {
+		return tc.leaves[i-tc.start]
+	}
+	return nil
+}
+
 // RandomNextLeaf returns a function that fetches a random leaf available in the tree.
 func RandomNextLeaf() func(uint64) uint64 {
 	return func(size uint64) uint64 {
diff --git a/hammer/hammer.go b/hammer/hammer.go
index 0396d80..dd7806b 100644
--- a/hammer/hammer.go
+++ b/hammer/hammer.go
@@ -402,6 +402,16 @@ func readHTTP(ctx context.Context, u *url.URL) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
+	defer func() {
+		if err := resp.Body.Close(); err != nil {
+			klog.Errorf("resp.Body.Close(): %v", err)
+		}
+	}()
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read body: %v", err)
+	}
+
 	switch resp.StatusCode {
 	case 404:
 		klog.Infof("Not found: %q", u.String())
@@ -411,10 +421,5 @@ func readHTTP(ctx context.Context, u *url.URL) ([]byte, error) {
 	default:
 		return nil, fmt.Errorf("unexpected http status %q", resp.Status)
 	}
-	defer func() {
-		if err := resp.Body.Close(); err != nil {
-			klog.Errorf("resp.Body.Close(): %v", err)
-		}
-	}()
-	return io.ReadAll(resp.Body)
+	return body, nil
 }