Improve release frequency for signing keys by checking tx result status

Author: m-Peter

bootstrap/bootstrap.go

@@ -243,7 +243,7 @@ func (b *Bootstrap) StartAPIServer(ctx context.Context) error {
 		}
 	}
 
-	b.keystore = keystore.New(accountKeys)
+	b.keystore = keystore.New(ctx, accountKeys, b.client, b.logger)
 
 	// create transaction pool
 	var txPool requester.TxPool

services/ingestion/event_subscriber.go

@@ -176,7 +176,7 @@ func (r *RPCEventSubscriber) subscribe(ctx context.Context, height uint64) <-cha
 				for _, evt := range blockEvents.Events {
 					r.keyLock.NotifyTransaction(evt.TransactionID)
 				}
-				r.keyLock.NotifyBlock(blockEvents.Height)
+				r.keyLock.NotifyBlock(blockEvents.BlockID)
 
 				eventsChan <- evmEvents
 

services/ingestion/event_subscriber_test.go

@@ -44,9 +44,13 @@ func Test_Subscribing(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	subscriber := NewRPCEventSubscriber(zerolog.Nop(), client, flowGo.Previewnet, keystore.New(nil), 1)
+	ctx := context.Background()
+	logger := zerolog.Nop()
+	ks := keystore.New(ctx, nil, client, logger)
 
-	events := subscriber.Subscribe(context.Background())
+	subscriber := NewRPCEventSubscriber(logger, client, flowGo.Previewnet, ks, 1)
+
+	events := subscriber.Subscribe(ctx)
 
 	var prevHeight uint64
 
@@ -84,9 +88,13 @@ func Test_MissingBlockEvent(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	subscriber := NewRPCEventSubscriber(zerolog.Nop(), client, flowGo.Previewnet, keystore.New(nil), 1)
+	ctx := context.Background()
+	logger := zerolog.Nop()
+	ks := keystore.New(ctx, nil, client, logger)
+
+	subscriber := NewRPCEventSubscriber(logger, client, flowGo.Previewnet, ks, 1)
 
-	events := subscriber.Subscribe(context.Background())
+	events := subscriber.Subscribe(ctx)
 
 	missingHashes := make([]gethCommon.Hash, 0)
 
@@ -186,9 +194,13 @@ func Test_SubscribingWithRetryOnError(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	subscriber := NewRPCEventSubscriber(zerolog.Nop(), client, flowGo.Previewnet, keystore.New(nil), 1)
+	ctx := context.Background()
+	logger := zerolog.Nop()
+	ks := keystore.New(ctx, nil, client, logger)
+
+	subscriber := NewRPCEventSubscriber(logger, client, flowGo.Previewnet, ks, 1)
 
-	events := subscriber.Subscribe(context.Background())
+	events := subscriber.Subscribe(ctx)
 
 	var prevHeight uint64
 
@@ -249,9 +261,13 @@ func Test_SubscribingWithRetryOnErrorMultipleBlocks(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	subscriber := NewRPCEventSubscriber(zerolog.Nop(), client, flowGo.Previewnet, keystore.New(nil), 1)
+	ctx := context.Background()
+	logger := zerolog.Nop()
+	ks := keystore.New(ctx, nil, client, logger)
 
-	events := subscriber.Subscribe(context.Background())
+	subscriber := NewRPCEventSubscriber(logger, client, flowGo.Previewnet, ks, 1)
+
+	events := subscriber.Subscribe(ctx)
 
 	var prevHeight uint64
 
@@ -311,9 +327,13 @@ func Test_SubscribingWithRetryOnErrorEmptyBlocks(t *testing.T) {
 	)
 	require.NoError(t, err)
 
-	subscriber := NewRPCEventSubscriber(zerolog.Nop(), client, flowGo.Previewnet, keystore.New(nil), 1)
+	ctx := context.Background()
+	logger := zerolog.Nop()
+	ks := keystore.New(ctx, nil, client, logger)
+
+	subscriber := NewRPCEventSubscriber(logger, client, flowGo.Previewnet, ks, 1)
 
-	events := subscriber.Subscribe(context.Background())
+	events := subscriber.Subscribe(ctx)
 
 	var prevHeight uint64
 

services/requester/keystore/key_store.go

@@ -1,11 +1,14 @@
 package keystore
 
 import (
+	"context"
 	"fmt"
 	"sync"
 
 	flowsdk "github.com/onflow/flow-go-sdk"
+	"github.com/onflow/flow-go-sdk/access"
 	"github.com/onflow/flow-go/model/flow"
+	"github.com/rs/zerolog"
 )
 
 var ErrNoKeysAvailable = fmt.Errorf("no signing keys available")
@@ -14,21 +17,37 @@ const accountKeyBlockExpiration = flow.DefaultTransactionExpiry
 
 type KeyLock interface {
 	NotifyTransaction(txID flowsdk.Identifier)
-	NotifyBlock(blockHeight uint64)
+	NotifyBlock(blockID flowsdk.Identifier)
 }
 
 type KeyStore struct {
+	client        access.Client
 	availableKeys chan *AccountKey
 	usedKeys      map[flowsdk.Identifier]*AccountKey
 	size          int
 	keyMu         sync.Mutex
+	blockChan     chan flowsdk.Identifier
+	logger        zerolog.Logger
+
+	// Signal channel used to prevent blocking writes
+	// on `blockChan` when the node is shutting down.
+	done chan struct{}
 }
 
 var _ KeyLock = (*KeyStore)(nil)
 
-func New(keys []*AccountKey) *KeyStore {
+func New(
+	ctx context.Context,
+	keys []*AccountKey,
+	client access.Client,
+	logger zerolog.Logger,
+) *KeyStore {
 	ks := &KeyStore{
-		usedKeys: map[flowsdk.Identifier]*AccountKey{},
+		client:    client,
+		usedKeys:  map[flowsdk.Identifier]*AccountKey{},
+		blockChan: make(chan flowsdk.Identifier, 100),
+		logger:    logger,
+		done:      make(chan struct{}),
 	}
 
 	availableKeys := make(chan *AccountKey, len(keys))
@@ -39,6 +58,13 @@ func New(keys []*AccountKey) *KeyStore {
 	ks.size = len(keys)
 	ks.availableKeys = availableKeys
 
+	// For cases where the EVM Gateway is run in an index-mode,
+	// there is no need to release any keys, since transaction
+	// submission is not allowed.
+	if ks.size > 0 {
+		go ks.processLockedKeys(ctx)
+	}
+
 	return ks
 }
 
@@ -69,16 +95,16 @@ func (k *KeyStore) NotifyTransaction(txID flowsdk.Identifier) {
 	k.unsafeUnlockKey(txID)
 }
 
-// NotifyBlock is called to notify the KeyStore of a new ingested block height.
+// NotifyBlock is called to notify the KeyStore of a newly ingested block.
 // Pending transactions older than a threshold number of blocks are removed.
-func (k *KeyStore) NotifyBlock(blockHeight uint64) {
-	k.keyMu.Lock()
-	defer k.keyMu.Unlock()
-
-	for txID, key := range k.usedKeys {
-		if blockHeight-key.lastLockedBlock.Load() >= accountKeyBlockExpiration {
-			k.unsafeUnlockKey(txID)
-		}
+func (k *KeyStore) NotifyBlock(blockID flowsdk.Identifier) {
+	select {
+	case <-k.done:
+		k.logger.Warn().Msg(
+			"received `NotifyBlock` while the server is shutting down",
+		)
+	default:
+		k.blockChan <- blockID
 	}
 }
 
@@ -106,3 +132,50 @@ func (k *KeyStore) setLockMetadata(
 	defer k.keyMu.Unlock()
 	k.usedKeys[txID] = key
 }
+
+func (k *KeyStore) processLockedKeys(ctx context.Context) {
+	for {
+		select {
+		case <-ctx.Done():
+			close(k.done)
+			return
+		case blockID := <-k.blockChan:
+			// TODO: If calling `k.client.GetTransactionResultsByBlockID` for each
+			// new block, seems to be problematic for ANs, we can take an approach
+			// like the one below:
+			// If the available keys ratio is >= 60% of the total signing keys,
+			// then we can skip checking the transaction result statuses.
+			// The signing keys from invalid EVM transactions, will eventually
+			// come again into availability, after `accountKeyBlockExpiration`
+			// blocks have passed.
+			// availablekeysRatio := float64(k.AvailableKeys()) / float64(k.size)
+			// if availablekeysRatio >= 0.6 {
+			// 	continue
+			// }
+
+			txResults, err := k.client.GetTransactionResultsByBlockID(context.Background(), blockID)
+			if err != nil || len(txResults) == 0 {
+				continue
+			}
+
+			k.releasekeys(txResults)
+			fmt.Println("[processLockedKeys]: ", k.AvailableKeys())
+		}
+	}
+}
+
+func (k *KeyStore) releasekeys(txResults []*flowsdk.TransactionResult) {
+	k.keyMu.Lock()
+	defer k.keyMu.Unlock()
+
+	for _, txResult := range txResults {
+		k.unsafeUnlockKey(txResult.TransactionID)
+	}
+
+	blockHeight := txResults[0].BlockHeight
+	for txID, key := range k.usedKeys {
+		if blockHeight-key.lastLockedBlock.Load() >= accountKeyBlockExpiration {
+			k.unsafeUnlockKey(txID)
+		}
+	}
+}

services/requester/keystore/key_store_test.go

@@ -1,14 +1,18 @@
 package keystore
 
 import (
+	"context"
 	"errors"
 	"sync"
 	"testing"
 	"time"
 
+	"github.com/onflow/flow-evm-gateway/services/testutils"
 	sdk "github.com/onflow/flow-go-sdk"
+	"github.com/onflow/flow-go-sdk/access/mocks"
 	"github.com/onflow/flow-go-sdk/test"
 	"github.com/onflow/flow-go/utils/unittest"
+	"github.com/rs/zerolog"
 	"golang.org/x/sync/errgroup"
 
 	"github.com/stretchr/testify/assert"
@@ -25,7 +29,7 @@ func TestTake(t *testing.T) {
 		keys = append(keys, NewAccountKey(*accountKey, addrGenerator.New(), signer))
 	}
 
-	ks := New(keys)
+	ks := New(context.Background(), keys, testutils.SetupClientForRange(1, 100), zerolog.Nop())
 
 	t.Run("Take with no metadata updates", func(t *testing.T) {
 		key, err := ks.Take()
@@ -67,10 +71,34 @@ func TestTake(t *testing.T) {
 	})
 
 	t.Run("Take with expiration", func(t *testing.T) {
+		blockHeight := uint64(10)
+		blockID10 := identifierFixture()
+		client := &testutils.MockClient{
+			Client: &mocks.Client{},
+			GetTransactionResultsByBlockIDFunc: func(ctx context.Context, blockID sdk.Identifier) ([]*sdk.TransactionResult, error) {
+				if blockID == blockID10 {
+					return []*sdk.TransactionResult{}, nil
+				}
+
+				return []*sdk.TransactionResult{
+					{
+						Status:           sdk.TransactionStatusFinalized,
+						Error:            nil,
+						Events:           []sdk.Event{},
+						BlockID:          blockID,
+						BlockHeight:      blockHeight + accountKeyBlockExpiration,
+						TransactionID:    identifierFixture(),
+						CollectionID:     identifierFixture(),
+						ComputationUsage: 104_512,
+					},
+				}, nil
+			},
+		}
+		ks := New(context.Background(), keys, client, zerolog.Nop())
+
 		key, err := ks.Take()
 		require.NoError(t, err)
 
-		blockHeight := uint64(10)
 		txID := identifierFixture()
 		key.SetLockMetadata(txID, blockHeight)
 
@@ -82,13 +110,17 @@ func TestTake(t *testing.T) {
 		assert.Equal(t, key, ks.usedKeys[txID])
 
 		// notify for one block before the expiration block, key should still be reserved
-		ks.NotifyBlock(blockHeight + accountKeyBlockExpiration - 1)
+		ks.NotifyBlock(blockID10)
 
 		assert.True(t, key.inUse.Load())
 		assert.Equal(t, key, ks.usedKeys[txID])
 
 		// notify for the expiration block
-		ks.NotifyBlock(blockHeight + accountKeyBlockExpiration)
+		ks.NotifyBlock(identifierFixture())
+
+		// Give some time to allow the KeyStore to check for the
+		// transaction result statuses in the background.
+		time.Sleep(time.Second * 2)
 
 		// keystore and key should be reset
 		assert.Equal(t, 2, ks.AvailableKeys())
@@ -106,9 +138,14 @@ func TestKeySigning(t *testing.T) {
 	accountKey.Index = 0 // the fixture starts from index 1
 	accountKey.SequenceNumber = 42
 
-	ks := New([]*AccountKey{
-		NewAccountKey(*accountKey, address, signer),
-	})
+	ks := New(
+		context.Background(),
+		[]*AccountKey{
+			NewAccountKey(*accountKey, address, signer),
+		},
+		testutils.SetupClientForRange(1, 100),
+		zerolog.Nop(),
+	)
 
 	key, err := ks.Take()
 	require.NoError(t, err)
@@ -154,7 +191,7 @@ func TestConcurrentUse(t *testing.T) {
 		keys = append(keys, key)
 	}
 
-	ks := New(keys)
+	ks := New(context.Background(), keys, testutils.SetupClientForRange(1, 100), zerolog.Nop())
 
 	g := errgroup.Group{}
 	g.SetLimit(concurrentTxCount)

services/testutils/mock_client.go

@@ -18,6 +18,7 @@ type MockClient struct {
 	GetEventsForHeightRangeFunc      func(
 		ctx context.Context, eventType string, startHeight uint64, endHeight uint64,
 	) ([]flow.BlockEvents, error)
+	GetTransactionResultsByBlockIDFunc func(ctx context.Context, blockID flow.Identifier) ([]*flow.TransactionResult, error)
 }
 
 func (c *MockClient) GetBlockHeaderByHeight(ctx context.Context, height uint64) (*flow.BlockHeader, error) {
@@ -50,6 +51,16 @@ func (c *MockClient) GetEventsForHeightRange(
 	return c.Client.GetEventsForHeightRange(ctx, eventType, startHeight, endHeight)
 }
 
+func (c *MockClient) GetTransactionResultsByBlockID(
+	ctx context.Context,
+	blockID flow.Identifier,
+) ([]*flow.TransactionResult, error) {
+	if c.GetTransactionResultsByBlockIDFunc != nil {
+		return c.GetTransactionResultsByBlockIDFunc(ctx, blockID)
+	}
+	return c.Client.GetTransactionResultsByBlockID(ctx, blockID)
+}
+
 func SetupClientForRange(startHeight uint64, endHeight uint64) *MockClient {
 	client, events := SetupClient(startHeight, endHeight)
 	go func() {