proof: add lru cache for trees (#109)

Author: worm-emoji

api/api.go

@@ -10,7 +10,9 @@ import (
 	"time"
 
 	"github.com/contextwtf/lanyard/api/tracing"
+	"github.com/ethereum/go-ethereum/common"
 
+	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/rs/cors"
 	"github.com/rs/zerolog"
@@ -20,12 +22,18 @@ import (
 )
 
 type Server struct {
-	db *pgxpool.Pool
+	db   *pgxpool.Pool
+	tlru *lru.Cache[common.Hash, cachedTree]
 }
 
 func New(db *pgxpool.Pool) *Server {
+	l, err := lru.New[common.Hash, cachedTree](1000)
+	if err != nil {
+		log.Fatal().Err(err).Msg("failed to create lru cache")
+	}
 	return &Server{
-		db: db,
+		db:   db,
+		tlru: l,
 	}
 }
 

api/proof.go

@@ -2,6 +2,7 @@ package api
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"net/http"
 
@@ -16,10 +17,41 @@ type getProofResp struct {
 	Proof        []hexutil.Bytes `json:"proof"`
 }
 
+type cachedTree struct {
+	r getTreeResp
+	t merkle.Tree
+}
+
+func (s *Server) getCachedTree(ctx context.Context, root common.Hash) (cachedTree, error) {
+	r, ok := s.tlru.Get(root)
+	if ok {
+		return r, nil
+	}
+
+	td, err := getTree(ctx, s.db, root.Bytes())
+	if err != nil {
+		return cachedTree{}, err
+	}
+
+	leaves := [][]byte{}
+	for _, l := range td.UnhashedLeaves {
+		leaves = append(leaves, l[:])
+	}
+
+	t := merkle.New(leaves)
+	ct := cachedTree{
+		r: td,
+		t: t,
+	}
+
+	s.tlru.Add(root, ct)
+	return ct, nil
+}
+
 func (s *Server) GetProof(w http.ResponseWriter, r *http.Request) {
 	var (
 		ctx  = r.Context()
-		root = common.FromHex(r.URL.Query().Get("root"))
+		root = common.HexToHash(r.URL.Query().Get("root"))
 		leaf = common.FromHex(r.URL.Query().Get("unhashedLeaf"))
 		addr = common.FromHex(r.URL.Query().Get("address"))
 	)
@@ -33,7 +65,7 @@ func (s *Server) GetProof(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	td, err := getTree(ctx, s.db, root)
+	ct, err := s.getCachedTree(ctx, root)
 	if errors.Is(err, pgx.ErrNoRows) {
 		s.sendJSONError(r, w, nil, http.StatusNotFound, "tree not found")
 		w.Header().Set("Cache-Control", "public, max-age=60")
@@ -44,22 +76,19 @@ func (s *Server) GetProof(w http.ResponseWriter, r *http.Request) {
 	}
 
 	var (
-		leaves [][]byte
 		target []byte
 	)
 	// check if leaf is in tree and error if not
-	for _, l := range td.UnhashedLeaves {
+	for _, l := range ct.r.UnhashedLeaves {
 		if len(target) == 0 {
 			if len(leaf) > 0 {
 				if bytes.Equal(l, leaf) {
 					target = l
 				}
-			} else if bytes.Equal(leaf2Addr(l, td.Ltd, td.Packed), addr) {
+			} else if bytes.Equal(leaf2Addr(l, ct.r.Ltd, ct.r.Packed), addr) {
 				target = l
 			}
 		}
-
-		leaves = append(leaves, l)
 	}
 
 	if len(target) == 0 {
@@ -68,8 +97,7 @@ func (s *Server) GetProof(w http.ResponseWriter, r *http.Request) {
 	}
 
 	var (
-		mt   = merkle.New(leaves)
-		p    = mt.Proof(mt.Index(target))
+		p    = ct.t.Proof(ct.t.Index(target))
 		phex = []hexutil.Bytes{}
 	)
 

go.mod

@@ -5,6 +5,7 @@ go 1.19
 require (
 	github.com/contextwtf/migrate v0.0.1
 	github.com/ethereum/go-ethereum v1.10.21
+	github.com/hashicorp/golang-lru/v2 v2.0.4
 	github.com/jackc/pgx/v4 v4.16.1
 	github.com/lib/pq v1.10.9
 	github.com/opentracing/opentracing-go v1.2.0

go.sum

@@ -311,6 +311,8 @@ github.com/hashicorp/go.net v0.0.1/go.mod h1:hjKkEWcCURg++eb33jQU7oqQcI9XDCnUzHA
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
+github.com/hashicorp/golang-lru/v2 v2.0.4 h1:7GHuZcgid37q8o5i3QI9KMT4nCWQQ3Kx3Ov6bb9MfK0=
+github.com/hashicorp/golang-lru/v2 v2.0.4/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/hashicorp/logutils v1.0.0/go.mod h1:QIAnNjmIWmVIIkWDTG1z5v++HQmx9WQRO+LraFDTW64=
 github.com/hashicorp/mdns v1.0.0/go.mod h1:tL+uN++7HEJ6SQLQ2/p+z2pH24WQKWjBPkE0mNTz8vQ=