From f5f58a7e91774fad2ce2c65d0f288478d4bb3a4b Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Sat, 29 Aug 2020 22:37:31 +0300
Subject: [PATCH] keys: add simple LRU key cache for 1024 elements

The cost of Y calculation from X is comparable with signature check, so it
reduces witness check overhead by ~30% for cached keys and gives ~5% overall
boost in TPS.
---
 go.mod                            |  1 +
 go.sum                            |  2 ++
 pkg/crypto/keys/publickey.go      | 21 ++++++++++++++++++++-
 pkg/crypto/keys/publickey_test.go |  4 ++++
 4 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 82a2b9b41..093933e3e 100644
--- a/go.mod
+++ b/go.mod
@@ -8,6 +8,7 @@ require (
 	github.com/frankban/quicktest v1.10.0 // indirect
 	github.com/go-redis/redis v6.10.2+incompatible
 	github.com/gorilla/websocket v1.4.2
+	github.com/hashicorp/golang-lru v0.5.4
 	github.com/mr-tron/base58 v1.1.2
 	github.com/nspcc-dev/dbft v0.0.0-20200711144034-c526ccc6f570
 	github.com/nspcc-dev/rfc6979 v0.2.0
diff --git a/go.sum b/go.sum
index a3d1764be..a95292a4e 100644
--- a/go.sum
+++ b/go.sum
@@ -110,6 +110,8 @@ github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
+github.com/hashicorp/golang-lru v0.5.4 h1:YDjusn29QI/Das2iO9M0BHnIbxPeyuCHsjMW+lJfyTc=
+github.com/hashicorp/golang-lru v0.5.4/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
 github.com/hpcloud/tail v1.0.0 h1:nfCOvKYfkgYP8hkirhJocXT2+zOD8yUNjXaWfTlyFKI=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
diff --git a/pkg/crypto/keys/publickey.go b/pkg/crypto/keys/publickey.go
index 1d50d51d0..f10f4e628 100644
--- a/pkg/crypto/keys/publickey.go
+++ b/pkg/crypto/keys/publickey.go
@@ -11,6 +11,7 @@ import (
 	"math/big"
 
 	"github.com/btcsuite/btcd/btcec"
+	lru "github.com/hashicorp/golang-lru"
 	"github.com/nspcc-dev/neo-go/pkg/core/interop/interopnames"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/nspcc-dev/neo-go/pkg/encoding/address"
@@ -105,13 +106,31 @@ func NewPublicKeyFromString(s string) (*PublicKey, error) {
 	return NewPublicKeyFromBytes(b, elliptic.P256())
 }
 
+// keycache is a simple lru cache for P256 keys that avoids Y calculation overhead
+// for known keys.
+var keycache *lru.Cache
+
+func init() {
+	// Less than 100K, probably enough for our purposes.
+	keycache, _ = lru.New(1024)
+}
+
 // NewPublicKeyFromBytes returns public key created from b using given EC.
 func NewPublicKeyFromBytes(b []byte, curve elliptic.Curve) (*PublicKey, error) {
-	pubKey := new(PublicKey)
+	var pubKey *PublicKey
+	cachedKey, ok := keycache.Get(string(b))
+	if ok {
+		pubKey = cachedKey.(*PublicKey)
+		if pubKey.Curve == curve {
+			return pubKey, nil
+		}
+	}
+	pubKey = new(PublicKey)
 	pubKey.Curve = curve
 	if err := pubKey.DecodeBytes(b); err != nil {
 		return nil, err
 	}
+	keycache.Add(string(b), pubKey)
 	return pubKey, nil
 }
 
diff --git a/pkg/crypto/keys/publickey_test.go b/pkg/crypto/keys/publickey_test.go
index c4b03c5d1..105a83712 100644
--- a/pkg/crypto/keys/publickey_test.go
+++ b/pkg/crypto/keys/publickey_test.go
@@ -63,6 +63,10 @@ func TestNewPublicKeyFromBytes(t *testing.T) {
 	pub, err := NewPublicKeyFromBytes(b, elliptic.P256())
 	require.NoError(t, err)
 	require.Equal(t, priv.PublicKey(), pub)
+	// Test cached access
+	pub2, err := NewPublicKeyFromBytes(b, elliptic.P256())
+	require.NoError(t, err)
+	require.Same(t, pub, pub2)
 }
 
 func TestDecodeFromString(t *testing.T) {