Added heartbeat failure counter per wallet

pkg/tbtc/node.go

@@ -65,10 +65,11 @@ type node struct {
 	// dkgExecutor MUST NOT be used outside this struct.
 	dkgExecutor *dkgExecutor
 
-	// heartbeatFailureCounter is the counter keeping track of consecutive
-	// heartbeat failure. It reset to zero after each successful heartbeat
-	// procedure.
-	heartbeatFailureCounter uint
+	heartbeatFailureCountersMutex sync.Mutex
+	// heartbeatFailureCounters holds counters keeping track of consecutive
+	// heartbeat failures. Each wallet has a separate counter. The key used in
+	// the map is the uncompressed public key (with 04 prefix) of the wallet.
+	heartbeatFailureCounters map[string]*uint
 
 	signingExecutorsMutex sync.Mutex
 	// signingExecutors is the cache holding signing executors for specific wallets.
@@ -111,16 +112,17 @@ func newNode(
 	scheduler.RegisterProtocol(latch)
 
 	node := &node{
-		groupParameters:       groupParameters,
-		chain:                 chain,
-		btcChain:              btcChain,
-		netProvider:           netProvider,
-		walletRegistry:        walletRegistry,
-		walletDispatcher:      newWalletDispatcher(),
-		protocolLatch:         latch,
-		signingExecutors:      make(map[string]*signingExecutor),
-		coordinationExecutors: make(map[string]*coordinationExecutor),
-		proposalGenerator:     proposalGenerator,
+		groupParameters:          groupParameters,
+		chain:                    chain,
+		btcChain:                 btcChain,
+		netProvider:              netProvider,
+		walletRegistry:           walletRegistry,
+		walletDispatcher:         newWalletDispatcher(),
+		protocolLatch:            latch,
+		heartbeatFailureCounters: make(map[string]*uint),
+		signingExecutors:         make(map[string]*signingExecutor),
+		coordinationExecutors:    make(map[string]*coordinationExecutor),
+		proposalGenerator:        proposalGenerator,
 	}
 
 	// Only the operator address is known at this point and can be pre-fetched.
@@ -213,6 +215,29 @@ func (n *node) validateDKG(
 	n.dkgExecutor.executeDkgValidation(seed, submissionBlock, result, resultHash)
 }
 
+func (n *node) getHeartbeatCounter(
+	walletPublicKey *ecdsa.PublicKey,
+) (*uint, error) {
+	n.heartbeatFailureCountersMutex.Lock()
+	defer n.heartbeatFailureCountersMutex.Unlock()
+
+	walletPublicKeyBytes, err := marshalPublicKey(walletPublicKey)
+	if err != nil {
+		return nil, fmt.Errorf("cannot marshal wallet public key: [%v]", err)
+	}
+
+	counterKey := hex.EncodeToString(walletPublicKeyBytes)
+
+	if counter, exists := n.heartbeatFailureCounters[counterKey]; exists {
+		return counter, nil
+	}
+
+	counterInitialValue := new(uint) // The value is zero-initialized.
+	n.heartbeatFailureCounters[counterKey] = counterInitialValue
+
+	return counterInitialValue, nil
+}
+
 // getSigningExecutor gets the signing executor responsible for executing
 // signing related to a specific wallet whose part is controlled by this node.
 // The second boolean return value indicates whether the node controls at least
@@ -461,6 +486,12 @@ func (n *node) handleHeartbeatProposal(
 		return
 	}
 
+	heartbeatFailureCounter, err := n.getHeartbeatCounter(wallet.publicKey)
+	if err != nil {
+		logger.Errorf("cannot get heartbeat failure counter: [%v]", err)
+		return
+	}
+
 	inactivityNotifier, err := n.getInactivityNotifier(wallet.publicKey)
 	if err != nil {
 		logger.Errorf("cannot get inactivity operator: [%v]", err)
@@ -488,7 +519,7 @@ func (n *node) handleHeartbeatProposal(
 		wallet,
 		signingExecutor,
 		proposal,
-		&n.heartbeatFailureCounter,
+		heartbeatFailureCounter,
 		inactivityNotifier,
 		startBlock,
 		expiryBlock,

pkg/tbtc/node_test.go

@@ -19,6 +19,94 @@ import (
 	"github.com/keep-network/keep-core/pkg/tecdsa"
 )
 
+func TestNode_GetHeartbeatCounter(t *testing.T) {
+	groupParameters := &GroupParameters{
+		GroupSize:       5,
+		GroupQuorum:     4,
+		HonestThreshold: 3,
+	}
+
+	localChain := Connect()
+	localProvider := local.Connect()
+
+	signer := createMockSigner(t)
+
+	// Populate the mock keystore with the mock signer's data. This is
+	// required to make the node controlling the signer's wallet.
+	keyStorePersistence := createMockKeyStorePersistence(t, signer)
+
+	node, err := newNode(
+		groupParameters,
+		localChain,
+		newLocalBitcoinChain(),
+		localProvider,
+		keyStorePersistence,
+		&mockPersistenceHandle{},
+		generator.StartScheduler(),
+		&mockCoordinationProposalGenerator{},
+		Config{},
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	walletPublicKey := signer.wallet.publicKey
+
+	testutils.AssertIntsEqual(
+		t,
+		"cache size",
+		0,
+		len(node.heartbeatFailureCounters),
+	)
+
+	counter, err := node.getHeartbeatCounter(walletPublicKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	testutils.AssertIntsEqual(
+		t,
+		"cache size",
+		1,
+		len(node.heartbeatFailureCounters),
+	)
+
+	testutils.AssertUintsEqual(t, "counter value", 0, uint64(*counter))
+
+	// Increment the counter and check the value again
+	*counter++
+	testutils.AssertUintsEqual(t, "counter value", 1, uint64(*counter))
+
+	// Construct an arbitrary public key representing a different wallet.
+	x, y := walletPublicKey.Curve.Double(walletPublicKey.X, walletPublicKey.Y)
+	anotherWalletPublicKey := &ecdsa.PublicKey{
+		Curve: walletPublicKey.Curve,
+		X:     x,
+		Y:     y,
+	}
+
+	anotherCounter, err := node.getHeartbeatCounter(anotherWalletPublicKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	testutils.AssertIntsEqual(
+		t,
+		"cache size",
+		2,
+		len(node.heartbeatFailureCounters),
+	)
+
+	testutils.AssertUintsEqual(t, "counter value", 0, uint64(*anotherCounter))
+
+	// Increment one counter and reset another.
+	*anotherCounter++
+	*counter = 0
+
+	testutils.AssertUintsEqual(t, "counter value", 0, uint64(*counter))
+	testutils.AssertUintsEqual(t, "counter value", 1, uint64(*anotherCounter))
+}
+
 func TestNode_GetSigningExecutor(t *testing.T) {
 	groupParameters := &GroupParameters{
 		GroupSize:       5,