[Merged by Bors] - atx: cache poet proofs with lru

activation/activation.go

@@ -54,6 +54,7 @@ type PoetConfig struct {
 	InfoCacheTTL                   time.Duration `mapstructure:"info-cache-ttl"`
 	PowParamsCacheTTL              time.Duration `mapstructure:"pow-params-cache-ttl"`
 	MaxRequestRetries              int           `mapstructure:"retry-max"`
+	PoetProofsCache                int           `mapstructure:"poet-proofs-cache"`
 }
 
 func DefaultPoetConfig() PoetConfig {
@@ -62,6 +63,7 @@ func DefaultPoetConfig() PoetConfig {
 		MaxRequestRetries: 10,
 		InfoCacheTTL:      5 * time.Minute,
 		PowParamsCacheTTL: 5 * time.Minute,
+		PoetProofsCache:   200,
 	}
 }
 

activation/poetdb.go

@@ -5,11 +5,13 @@
 	"encoding/hex"
 	"fmt"
 
+	lru "github.com/hashicorp/golang-lru/v2"
 	"github.com/spacemeshos/merkle-tree"
 	"github.com/spacemeshos/poet/hash"
 	"github.com/spacemeshos/poet/shared"
 	"github.com/spacemeshos/poet/verifier"
 	"go.uber.org/zap"
+	"golang.org/x/sync/singleflight"
 
 	"github.com/spacemeshos/go-spacemesh/codec"
 	"github.com/spacemeshos/go-spacemesh/common/types"
@@ -19,19 +21,59 @@
 	"github.com/spacemeshos/go-spacemesh/sql/poets"
 )
 
+// PoetDbOptions are options for PoetDb.
+type PoetDbOptions struct {
+	cacheSize int
+}
+
+type PoetDbOption func(*PoetDbOptions)
+
+// WithCacheSize sets the cache size for PoetDb.
+func WithCacheSize(size int) PoetDbOption {
+	return func(opts *PoetDbOptions) {
+		opts.cacheSize = size
+	}
+}
+
 // PoetDb is a database for PoET proofs.
 type PoetDb struct {
-	sqlDB  sql.StateDatabase
-	logger *zap.Logger
+	sqlDB               sql.StateDatabase
+	poetProofsDbRequest singleflight.Group
+	poetProofsLru       *lru.Cache[types.PoetProofRef, *types.PoetProofMessage]
+	logger              *zap.Logger
 }
 
 // NewPoetDb returns a new PoET handler.
-func NewPoetDb(db sql.StateDatabase, log *zap.Logger) *PoetDb {
-	return &PoetDb{sqlDB: db, logger: log}
+func NewPoetDb(db sql.StateDatabase, log *zap.Logger, opts ...PoetDbOption) *PoetDb {
+	options := PoetDbOptions{
+		// in last epochs there are 45 proofs per epoch, with each of them nearly 140KB
+		// 200 is set not to keep multiple epochs, but to account for unexpected growth
+		// select round_id, count(*), max(length(poet)) from poets group by round_id;
+		// 25|42|146200
+		// 26|45|145936
+		// 27|45|145738
+		// 28|45|145903
+		cacheSize: 200,
+	}
+	for _, opt := range opts {
+		opt(&options)
+	}
+	poetProofsLru, err := lru.New[types.PoetProofRef, *types.PoetProofMessage](options.cacheSize)
+	if err != nil {
+		log.Panic("failed to create PoET proofs LRU cache", zap.Error(err))
+	}
+	return &PoetDb{
+		sqlDB:         db,
+		poetProofsLru: poetProofsLru,
+		logger:        log,
+	}
 }
 
 // HasProof returns true if the database contains a proof with the given reference, or false otherwise.
 func (db *PoetDb) HasProof(proofRef types.PoetProofRef) bool {
+	if db.poetProofsLru.Contains(proofRef) {
+		return true
+	}
 	has, err := poets.Has(db.sqlDB, proofRef)
 	return err == nil && has
 }
@@ -99,6 +141,7 @@
 
 // StoreProof saves the poet proof in local db.
 func (db *PoetDb) StoreProof(ctx context.Context, ref types.PoetProofRef, proofMessage *types.PoetProofMessage) error {
+	db.poetProofsLru.Add(ref, proofMessage)
 	messageBytes, err := codec.Encode(proofMessage)
 	if err != nil {
 		return fmt.Errorf("could not marshal proof message: %w", err)
@@ -146,15 +189,27 @@
 
 // Proof returns full proof.
 func (db *PoetDb) Proof(proofRef types.PoetProofRef) (*types.PoetProof, *types.Hash32, error) {
-	proofMessageBytes, err := db.GetProofMessage(proofRef)
+	response, err, _ := db.poetProofsDbRequest.Do(string(proofRef[:]), func() (any, error) {
+		cachedProof, ok := db.poetProofsLru.Get(proofRef)
+		if ok && cachedProof != nil {
+			return cachedProof, nil
+		}
+		proofMessageBytes, err := db.GetProofMessage(proofRef)
+		if err != nil {
+			return nil, fmt.Errorf("could not fetch poet proof for ref %x: %w", proofRef, err)
+		}
+		var proofMessage types.PoetProofMessage
+		if err := codec.Decode(proofMessageBytes, &proofMessage); err != nil {
+			return nil, fmt.Errorf("failed to unmarshal poet proof for ref %x: %w", proofRef, err)
+		}
+		db.poetProofsLru.Add(proofRef, &proofMessage)
+		return &proofMessage, nil
+	})
 	if err != nil {
-		return nil, nil, fmt.Errorf("could not fetch poet proof for ref %x: %w", proofRef, err)
-	}
-	var proofMessage types.PoetProofMessage
-	if err := codec.Decode(proofMessageBytes, &proofMessage); err != nil {
-		return nil, nil, fmt.Errorf("failed to unmarshal poet proof for ref %x: %w", proofRef, err)
+		return nil, nil, err
 	}
-	return &proofMessage.PoetProof, &proofMessage.Statement, nil
+	proof := response.(*types.PoetProofMessage)
+	return &proof.PoetProof, &proof.Statement, nil
 }
 
 func (db *PoetDb) ProofForRound(poetID []byte, roundID string) (*types.PoetProof, error) {

config/mainnet.go

@@ -173,6 +173,7 @@ func MainnetConfig() Config {
 			RequestTimeout:    1100 * time.Second,
 			RequestRetryDelay: 10 * time.Second,
 			MaxRequestRetries: 10,
+			PoetProofsCache:   200,
 		},
 		POST: activation.PostConfig{
 			MinNumUnits:   4,

config/presets/testnet.go

@@ -126,6 +126,7 @@ func testnet() config.Config {
 
 			InfoCacheTTL:      5 * time.Minute,
 			PowParamsCacheTTL: 5 * time.Minute,
+			PoetProofsCache:   200,
 		},
 		POST: activation.PostConfig{
 			MinNumUnits:   2,

node/node.go

@@ -580,7 +580,10 @@ func (app *App) initServices(ctx context.Context) error {
 	layersPerEpoch := types.GetLayersPerEpoch()
 	lg := app.log
 
-	poetDb := activation.NewPoetDb(app.db, app.addLogger(PoetDbLogger, lg).Zap())
+	poetDb := activation.NewPoetDb(
+		app.db,
+		app.addLogger(PoetDbLogger, lg).Zap(),
+		activation.WithCacheSize(app.Config.POET.PoetProofsCache))
 	postStates := activation.NewPostStates(app.addLogger(PostLogger, lg).Zap())
 	opts := []activation.PostVerifierOpt{
 		activation.WithVerifyingOpts(app.Config.SMESHING.VerifyingOpts),