perf: add cache to address codec

CHANGELOG.md

@@ -65,6 +65,7 @@ Every module contains its own CHANGELOG.md. Please refer to the module you are i
 
 ### Improvements
 
+* (codec) [#20122](https://github.com/cosmos/cosmos-sdk/pull/20122) Added a cache to address codec.
 * (types) [#19869](https://github.com/cosmos/cosmos-sdk/pull/19869) Removed `Any` type from `codec/types` and replaced it with an alias for `cosmos/gogoproto/types/any`.
 * (server) [#19854](https://github.com/cosmos/cosmos-sdk/pull/19854) Add customizability to start command.
     * Add `StartCmdOptions` in `server.AddCommands` instead of `servertypes.ModuleInitFlags`. To set custom flags set them in the `StartCmdOptions` struct on the `AddFlags` field.

codec/address/bech32_codec.go

@@ -2,30 +2,79 @@ package address
 
 import (
 	"errors"
+	"fmt"
 	"strings"
+	"sync"
+
+	"github.com/hashicorp/golang-lru/simplelru"
 
 	"cosmossdk.io/core/address"
 	errorsmod "cosmossdk.io/errors"
 
+	"github.com/cosmos/cosmos-sdk/internal/conv"
 	sdkAddress "github.com/cosmos/cosmos-sdk/types/address"
 	"github.com/cosmos/cosmos-sdk/types/bech32"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 )
 
+var errEmptyAddress = errors.New("empty address string is not allowed")
+
+type Option func(*Options)
+
+type Options struct {
+	mu    *sync.Mutex
+	cache *simplelru.LRU
+}
+
+func WithLRU(cache *simplelru.LRU) Option {
+	return func(o *Options) {
+		o.cache = cache
+	}
+}
+
+func WithMutex(mu *sync.Mutex) Option {
+	return func(o *Options) {
+		o.mu = mu
+	}
+}
+
 type Bech32Codec struct {
 	Bech32Prefix string
 }
 
-var _ address.Codec = &Bech32Codec{}
+type cachedBech32Codec struct {
+	codec Bech32Codec
+	mu    *sync.Mutex
+	cache *simplelru.LRU
+}
+
+var (
+	_ address.Codec = &Bech32Codec{}
+	_ address.Codec = &cachedBech32Codec{}
+)
+
+func NewBech32Codec(prefix string, opts ...Option) address.Codec {
+	options := Options{}
+	for _, optionFn := range opts {
+		optionFn(&options)
+	}
 
-func NewBech32Codec(prefix string) address.Codec {
-	return Bech32Codec{prefix}
+	ac := Bech32Codec{prefix}
+	if options.mu == nil || options.cache == nil {
+		return ac
+	}
+
+	return cachedBech32Codec{
+		codec: ac,
+		cache: options.cache,
+		mu:    options.mu,
+	}
 }
 
 // StringToBytes encodes text to bytes
 func (bc Bech32Codec) StringToBytes(text string) ([]byte, error) {
 	if len(strings.TrimSpace(text)) == 0 {
-		return []byte{}, errors.New("empty address string is not allowed")
+		return []byte{}, errEmptyAddress
 	}
 
 	hrp, bz, err := bech32.DecodeAndConvert(text)
@@ -61,3 +110,53 @@ func (bc Bech32Codec) BytesToString(bz []byte) (string, error) {
 
 	return text, nil
 }
+
+func (cbc cachedBech32Codec) BytesToString(bz []byte) (string, error) {
+	if len(bz) == 0 {
+		return "", nil
+	}
+
+	key := conv.UnsafeBytesToStr(bz)
+	cbc.mu.Lock()
+	defer cbc.mu.Unlock()
+
+	addrs, ok := cbc.cache.Get(key)
+	if !ok {
+		addrs = make(map[string]string)
+		cbc.cache.Add(key, addrs)
+	}
+
+	addrMap, ok := addrs.(map[string]string)
+	if !ok {
+		return "", fmt.Errorf("cache contains non-map[string]string value for key %s", key)
+	}
+
+	addr, ok := addrMap[cbc.codec.Bech32Prefix]
+	if !ok {
+		var err error
+		addr, err = cbc.codec.BytesToString(bz)
+		if err != nil {
+			return "", err
+		}
+		addrMap[cbc.codec.Bech32Prefix] = addr
+	}
+
+	return addr, nil
+}
+
+func (cbc cachedBech32Codec) StringToBytes(text string) ([]byte, error) {
+	cbc.mu.Lock()
+	defer cbc.mu.Unlock()
+
+	if addr, ok := cbc.cache.Get(text); ok {
+		return addr.([]byte), nil
+	}
+
+	addr, err := cbc.codec.StringToBytes(text)
+	if err != nil {
+		return nil, err
+	}
+	cbc.cache.Add(text, addr)
+
+	return addr, nil
+}

codec/address/bech32_codec_test.go

@@ -0,0 +1,137 @@
+package address
+
+import (
+	"crypto/rand"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/hashicorp/golang-lru/simplelru"
+	"github.com/stretchr/testify/require"
+
+	"github.com/cosmos/cosmos-sdk/internal/conv"
+)
+
+var (
+	lru, _ = simplelru.NewLRU(500, nil)
+	mu     = &sync.Mutex{}
+)
+
+func generateAddresses(totalAddresses int) ([][]byte, error) {
+	keys := make([][]byte, totalAddresses)
+	addr := make([]byte, 32)
+	for i := 0; i < totalAddresses; i++ {
+		_, err := rand.Read(addr)
+		if err != nil {
+			return nil, err
+		}
+		keys[i] = addr
+	}
+
+	return keys, nil
+}
+
+func TestNewBech32Codec(t *testing.T) {
+	tests := []struct {
+		name    string
+		prefix  string
+		lru     *simplelru.LRU
+		address string
+	}{
+		{
+			name:    "create accounts cached bech32 codec",
+			prefix:  "cosmos",
+			lru:     lru,
+			address: "cosmos1p8s0p6gqc6c9gt77lgr2qqujz49huhu6a80smx",
+		},
+		{
+			name:    "create validator cached bech32 codec",
+			prefix:  "cosmosvaloper",
+			lru:     lru,
+			address: "cosmosvaloper1sjllsnramtg3ewxqwwrwjxfgc4n4ef9u2lcnj0",
+		},
+		{
+			name:    "create consensus cached bech32 codec",
+			prefix:  "cosmosvalcons",
+			lru:     lru,
+			address: "cosmosvalcons1ntk8eualewuprz0gamh8hnvcem2nrcdsgz563h",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ac := NewBech32Codec(tt.prefix, WithLRU(lru), WithMutex(mu))
+			cached, ok := ac.(cachedBech32Codec)
+			require.True(t, ok)
+			require.Equal(t, cached.cache, tt.lru)
+
+			addr, err := ac.StringToBytes(tt.address)
+			require.NoError(t, err)
+
+			cachedAddr, ok := tt.lru.Get(tt.address)
+			require.True(t, ok)
+			require.Equal(t, addr, cachedAddr)
+
+			accAddr, err := ac.BytesToString(addr)
+			require.NoError(t, err)
+
+			cachedStrAddr, ok := tt.lru.Get(conv.UnsafeBytesToStr(addr))
+			require.True(t, ok)
+			cachedStrAddrMap, ok := cachedStrAddr.(map[string]string)
+			require.True(t, ok)
+			require.Equal(t, accAddr, cachedStrAddrMap[tt.prefix])
+		})
+	}
+}
+
+func TestMultipleBech32Codec(t *testing.T) {
+	cosmosAc, ok := NewBech32Codec("cosmos", WithLRU(lru), WithMutex(mu)).(cachedBech32Codec)
+	require.True(t, ok)
+	stakeAc := NewBech32Codec("stake", WithLRU(lru), WithMutex(mu)).(cachedBech32Codec)
+	require.True(t, ok)
+	require.Equal(t, cosmosAc.cache, stakeAc.cache)
+
+	addr := make([]byte, 32)
+	_, err := rand.Read(addr)
+	require.NoError(t, err)
+
+	cosmosAddr, err := cosmosAc.BytesToString(addr)
+	require.NoError(t, err)
+	stakeAddr, err := stakeAc.BytesToString(addr)
+	require.NoError(t, err)
+	require.True(t, cosmosAddr != stakeAddr)
+
+	cachedCosmosAddr, err := cosmosAc.BytesToString(addr)
+	require.NoError(t, err)
+	require.Equal(t, cosmosAddr, cachedCosmosAddr)
+
+	cachedStakeAddr, err := stakeAc.BytesToString(addr)
+	require.NoError(t, err)
+	require.Equal(t, stakeAddr, cachedStakeAddr)
+}
+
+func TestBech32CodecRace(t *testing.T) {
+	ac := NewBech32Codec("cosmos", WithLRU(lru), WithMutex(mu))
+	myAddrBz := []byte{0x1, 0x2, 0x3, 0x4, 0x5}
+
+	var (
+		wgStart, wgDone sync.WaitGroup
+		errCount        atomic.Uint32
+	)
+	const n = 3
+	wgStart.Add(n)
+	wgDone.Add(n)
+	for i := 0; i < n; i++ {
+		go func() {
+			wgStart.Done()
+			wgStart.Wait() // wait for all routines started
+
+			got, err := ac.BytesToString(myAddrBz)
+			if err != nil || got != "cosmos1qypqxpq9dc9msf" {
+				errCount.Add(1)
+			}
+			wgDone.Done()
+		}()
+	}
+	wgDone.Wait() // wait for all routines completed
+	require.Equal(t, errCount.Load(), uint32(0))
+}

codec/address/bench_test.go

@@ -0,0 +1,34 @@
+package address
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"cosmossdk.io/core/address"
+)
+
+func BenchmarkCodecWithCache(b *testing.B) {
+	cdc := NewBech32Codec("cosmos", WithLRU(lru), WithMutex(mu))
+	bytesToString(b, cdc)
+}
+
+func BenchmarkCodecWithoutCache(b *testing.B) {
+	cdc := Bech32Codec{Bech32Prefix: "cosmos"}
+	bytesToString(b, cdc)
+}
+
+func bytesToString(b *testing.B, cdc address.Codec) {
+	b.Helper()
+	addresses, err := generateAddresses(10)
+	require.NoError(b, err)
+
+	b.Helper()
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		_, err := cdc.BytesToString(addresses[i%len(addresses)])
+		require.NoError(b, err)
+	}
+}

codec/address/fuzz_test.go

@@ -0,0 +1,66 @@
+package address
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"cosmossdk.io/core/address"
+
+	sdkAddress "github.com/cosmos/cosmos-sdk/types/address"
+)
+
+func FuzzCachedAddressCodec(f *testing.F) {
+	if testing.Short() {
+		f.Skip()
+	}
+
+	addresses, err := generateAddresses(2)
+	require.NoError(f, err)
+
+	for _, addr := range addresses {
+		f.Add(addr)
+	}
+	cdc := NewBech32Codec("cosmos", WithLRU(lru), WithMutex(mu))
+
+	f.Fuzz(func(t *testing.T, addr []byte) {
+		checkAddress(t, addr, cdc)
+	})
+}
+
+func FuzzAddressCodec(f *testing.F) {
+	if testing.Short() {
+		f.Skip()
+	}
+	addresses, err := generateAddresses(2)
+	require.NoError(f, err)
+
+	for _, addr := range addresses {
+		f.Add(addr)
+	}
+
+	cdc := Bech32Codec{Bech32Prefix: "cosmos"}
+
+	f.Fuzz(func(t *testing.T, addr []byte) {
+		checkAddress(t, addr, cdc)
+	})
+}
+
+func checkAddress(t *testing.T, addr []byte, cdc address.Codec) {
+	t.Helper()
+	if len(addr) > sdkAddress.MaxAddrLen {
+		return
+	}
+	strAddr, err := cdc.BytesToString(addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	b, err := cdc.StringToBytes(strAddr)
+	if err != nil {
+		if !errors.Is(errEmptyAddress, err) {
+			t.Fatal(err)
+		}
+	}
+	require.Equal(t, len(addr), len(b))
+}