From e10b8f53d6a7fb71322ff9baa26f80df208961e1 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 10 Feb 2022 14:58:14 +0300
Subject: [PATCH] [#1151] morph/client: Cache notary transaction heights

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>
---
 pkg/morph/client/client.go      |  6 ++++--
 pkg/morph/client/constructor.go | 10 ++++++++++
 pkg/morph/client/notary.go      | 15 +++++++++++++--
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/pkg/morph/client/client.go b/pkg/morph/client/client.go
index b6c5e4973..84076ce3a 100644
--- a/pkg/morph/client/client.go
+++ b/pkg/morph/client/client.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"time"
 
+	lru "github.com/hashicorp/golang-lru"
 	"github.com/nspcc-dev/neo-go/pkg/core/native/nativenames"
 	"github.com/nspcc-dev/neo-go/pkg/core/native/noderoles"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
@@ -42,8 +43,9 @@ type Client struct {
 }
 
 type cache struct {
-	nnsHash  util.Uint160
-	groupKey *keys.PublicKey
+	nnsHash   util.Uint160
+	groupKey  *keys.PublicKey
+	txHeights *lru.Cache
 }
 
 type singleClient struct {
diff --git a/pkg/morph/client/constructor.go b/pkg/morph/client/constructor.go
index d1c6e13c1..410c6383d 100644
--- a/pkg/morph/client/constructor.go
+++ b/pkg/morph/client/constructor.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"time"
 
+	lru "github.com/hashicorp/golang-lru"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
 	"github.com/nspcc-dev/neo-go/pkg/rpc/client"
@@ -83,6 +84,7 @@ func New(key *keys.PrivateKey, endpoint string, opts ...Option) (*Client, error)
 
 	if cfg.singleCli != nil {
 		return &Client{
+			cache:        initClientCache(),
 			singleClient: blankSingleClient(cfg.singleCli, wallet.NewAccountFromPrivateKey(key), cfg),
 		}, nil
 	}
@@ -90,6 +92,7 @@ func New(key *keys.PrivateKey, endpoint string, opts ...Option) (*Client, error)
 	endpoints := append(cfg.extraEndpoints, endpoint)
 
 	return &Client{
+		cache: initClientCache(),
 		multiClient: &multiClient{
 			cfg:       *cfg,
 			account:   wallet.NewAccountFromPrivateKey(key),
@@ -99,6 +102,13 @@ func New(key *keys.PrivateKey, endpoint string, opts ...Option) (*Client, error)
 	}, nil
 }
 
+func initClientCache() cache {
+	c, _ := lru.New(100) // returns error only if size is negative
+	return cache{
+		txHeights: c,
+	}
+}
+
 // WithContext returns a client constructor option that
 // specifies the neo-go client context.
 //
diff --git a/pkg/morph/client/notary.go b/pkg/morph/client/notary.go
index 3e84e5de6..5be1e435d 100644
--- a/pkg/morph/client/notary.go
+++ b/pkg/morph/client/notary.go
@@ -876,13 +876,24 @@ func (c *Client) CalculateNonceAndVUB(hash util.Uint256) (nonce uint32, vub uint
 		return 0, 0, nil
 	}
 
-	// TODO: #1151 cache values since some operations uses same TX as triggers
 	nonce = binary.LittleEndian.Uint32(hash.BytesLE())
 
-	height, err := c.client.GetTransactionHeight(hash)
+	height, err := c.getTransactionHeight(hash)
 	if err != nil {
 		return 0, 0, fmt.Errorf("could not get transaction height: %w", err)
 	}
 
 	return nonce, height + c.notary.txValidTime, nil
 }
+
+func (c *Client) getTransactionHeight(h util.Uint256) (uint32, error) {
+	if rh, ok := c.txHeights.Get(h); ok {
+		return rh.(uint32), nil
+	}
+	height, err := c.client.GetTransactionHeight(h)
+	if err != nil {
+		return 0, err
+	}
+	c.txHeights.Add(h, height)
+	return height, nil
+}