Robert/addr memory leak (cosmos#8717)

* address: adding cache address.String() cache benchmark

* address: use LRU cache for .String()

* optimize address.Empty

* move cache initialization to init function

* Use UnsafeBytesToStr convertion with Addr cache

* add cache for other address String() methods

* fix linter issue

* add a non trivial address for Address.String benchmark

* add comment about cache size and update cashe size to 60k

* fix syntax

Co-authored-by: Alessio Treglia <[email protected]>

types/address.go

@@ -13,9 +13,11 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/codec/legacy"
 	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
+	"github.com/cosmos/cosmos-sdk/internal/conv"
 	"github.com/cosmos/cosmos-sdk/types/address"
 	"github.com/cosmos/cosmos-sdk/types/bech32"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	"github.com/hashicorp/golang-lru/simplelru"
 )
 
 const (
@@ -72,6 +74,33 @@ const (
 	Bech32PrefixConsPub = Bech32MainPrefix + PrefixValidator + PrefixConsensus + PrefixPublic
 )
 
+// cache variables
+var (
+	// AccAddress.String() is expensive and if unoptimized dominantly showed up in profiles,
+	// yet has no mechanisms to trivially cache the result given that AccAddress is a []byte type.
+	accAddrMu     sync.Mutex
+	accAddrCache  *simplelru.LRU
+	consAddrMu    sync.Mutex
+	consAddrCache *simplelru.LRU
+	valAddrMu     sync.Mutex
+	valAddrCache  *simplelru.LRU
+)
+
+func init() {
+	var err error
+	// in total the cache size is 61k entries. Key is 32 bytes and value is around 50-70 bytes.
+	// That will make around 92 * 61k * 2 (LRU) bytes ~ 11 MB
+	if accAddrCache, err = simplelru.NewLRU(60000, nil); err != nil {
+		panic(err)
+	}
+	if consAddrCache, err = simplelru.NewLRU(500, nil); err != nil {
+		panic(err)
+	}
+	if valAddrCache, err = simplelru.NewLRU(500, nil); err != nil {
+		panic(err)
+	}
+}
+
 // Address is a common interface for different types of addresses used by the SDK
 type Address interface {
 	Equals(Address) bool
@@ -158,12 +187,7 @@ func (aa AccAddress) Equals(aa2 Address) bool {
 
 // Returns boolean for whether an AccAddress is empty
 func (aa AccAddress) Empty() bool {
-	if aa == nil {
-		return true
-	}
-
-	aa2 := AccAddress{}
-	return bytes.Equal(aa.Bytes(), aa2.Bytes())
+	return aa == nil || len(aa) == 0
 }
 
 // Marshal returns the raw address bytes. It is needed for protobuf
@@ -237,40 +261,18 @@ func (aa AccAddress) Bytes() []byte {
 	return aa
 }
 
-// AccAddress.String() is expensive and if unoptimized dominantly showed up in profiles,
-// yet has no mechanisms to trivially cache the result given that AccAddress is a []byte type.
-// With the string interning below, we are able to achieve zero cost allocations for string->[]byte
-// because the Go compiler recognizes the special case of map[string([]byte)] when used exactly
-// in that pattern. See https://github.com/cosmos/cosmos-sdk/issues/8693.
-var addMu sync.Mutex
-var addrStrMemoize = make(map[string]string)
-
 // String implements the Stringer interface.
-func (aa AccAddress) String() (str string) {
-	addMu.Lock()
-	defer addMu.Unlock()
-
-	if str, ok := addrStrMemoize[string(aa)]; ok {
-		return str
-	}
-
-	// Otherwise cache it for later memoization.
-	defer func() {
-		addrStrMemoize[string(aa)] = str
-	}()
-
+func (aa AccAddress) String() string {
 	if aa.Empty() {
 		return ""
 	}
 
-	bech32PrefixAccAddr := GetConfig().GetBech32AccountAddrPrefix()
-
-	bech32Addr, err := bech32.ConvertAndEncode(bech32PrefixAccAddr, aa.Bytes())
-	if err != nil {
-		panic(err)
+	var key = conv.UnsafeBytesToStr(aa)
+	if addr, ok := accAddrCache.Get(key); ok {
+		return addr.(string)
 	}
-
-	return bech32Addr
+	return cacheBech32Addr(GetConfig().GetBech32AccountAddrPrefix(),
+		aa, &accAddrMu, accAddrCache, key)
 }
 
 // Format implements the fmt.Formatter interface.
@@ -332,12 +334,7 @@ func (va ValAddress) Equals(va2 Address) bool {
 
 // Returns boolean for whether an AccAddress is empty
 func (va ValAddress) Empty() bool {
-	if va == nil {
-		return true
-	}
-
-	va2 := ValAddress{}
-	return bytes.Equal(va.Bytes(), va2.Bytes())
+	return va == nil || len(va) == 0
 }
 
 // Marshal returns the raw address bytes. It is needed for protobuf
@@ -418,14 +415,12 @@ func (va ValAddress) String() string {
 		return ""
 	}
 
-	bech32PrefixValAddr := GetConfig().GetBech32ValidatorAddrPrefix()
-
-	bech32Addr, err := bech32.ConvertAndEncode(bech32PrefixValAddr, va.Bytes())
-	if err != nil {
-		panic(err)
+	var key = conv.UnsafeBytesToStr(va)
+	if addr, ok := valAddrCache.Get(key); ok {
+		return addr.(string)
 	}
-
-	return bech32Addr
+	return cacheBech32Addr(GetConfig().GetBech32ValidatorAddrPrefix(),
+		va, &valAddrMu, valAddrCache, key)
 }
 
 // Format implements the fmt.Formatter interface.
@@ -492,12 +487,7 @@ func (ca ConsAddress) Equals(ca2 Address) bool {
 
 // Returns boolean for whether an ConsAddress is empty
 func (ca ConsAddress) Empty() bool {
-	if ca == nil {
-		return true
-	}
-
-	ca2 := ConsAddress{}
-	return bytes.Equal(ca.Bytes(), ca2.Bytes())
+	return ca == nil || len(ca) == 0
 }
 
 // Marshal returns the raw address bytes. It is needed for protobuf
@@ -578,14 +568,12 @@ func (ca ConsAddress) String() string {
 		return ""
 	}
 
-	bech32PrefixConsAddr := GetConfig().GetBech32ConsensusAddrPrefix()
-
-	bech32Addr, err := bech32.ConvertAndEncode(bech32PrefixConsAddr, ca.Bytes())
-	if err != nil {
-		panic(err)
+	var key = conv.UnsafeBytesToStr(ca)
+	if addr, ok := consAddrCache.Get(key); ok {
+		return addr.(string)
 	}
-
-	return bech32Addr
+	return cacheBech32Addr(GetConfig().GetBech32ConsensusAddrPrefix(),
+		ca, &consAddrMu, consAddrCache, key)
 }
 
 // Bech32ifyAddressBytes returns a bech32 representation of address bytes.
@@ -730,3 +718,14 @@ func addressBytesFromHexString(address string) ([]byte, error) {
 
 	return hex.DecodeString(address)
 }
+
+func cacheBech32Addr(prefix string, addr []byte, m sync.Locker, cache *simplelru.LRU, cacheKey string) string {
+	bech32Addr, err := bech32.ConvertAndEncode(prefix, addr)
+	if err != nil {
+		panic(err)
+	}
+	m.Lock()
+	cache.Add(cacheKey, bech32Addr)
+	m.Unlock()
+	return bech32Addr
+}

types/address_bench_test.go

@@ -11,6 +11,25 @@ import (
 	"github.com/cosmos/cosmos-sdk/types"
 )
 
+func BenchmarkAccAddressString(b *testing.B) {
+	b.ReportAllocs()
+	pkBz := make([]byte, ed25519.PubKeySize)
+	pk := &ed25519.PubKey{Key: pkBz}
+	a := pk.Address()
+	pk2 := make([]byte, ed25519.PubKeySize)
+	for i := 1; i <= ed25519.PubKeySize; i++ {
+		pk2[i] = byte(i)
+	}
+	a2 := pk.Address()
+	var str, str2 string
+	for i := 0; i < b.N; i++ {
+		str = a.String()
+		str2 = a2.String()
+	}
+	require.NotEmpty(b, str)
+	require.NotEmpty(b, str2)
+}
+
 func BenchmarkBech32ifyPubKey(b *testing.B) {
 	b.ReportAllocs()
 	pkBz := make([]byte, ed25519.PubKeySize)