proof: add lru cache for trees

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=