chore: reduce locks (#194)

CHANGELOG.md

@@ -37,6 +37,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## Unreleased
 
+### Improvements
+
+* [#194](https://github.com/babylonlabs-io/vigilante/pull/194) fix: reduce locks
+
 ## v0.19.7
 
 ### Improvements

btcstaking-tracker/stakingeventwatcher/stakingeventwatcher.go

@@ -79,9 +79,6 @@ type StakingEventWatcher struct {
 	unbondingTracker *TrackedDelegations
 	// keeps track of verified delegations to be activated, periodically iterate over them and try to activate them
 	pendingTracker *TrackedDelegations
-	// keeps track of delegations that are verified, and we are trying to activate
-	inProgressTracker *TrackedDelegations
-
 	// used for metrics purposes, keeps track of verified delegations that are not in consumer chain
 	verifiedNotInChainTracker *TrackedDelegations
 	// used for metrics purposes, keeps track of verified delegations in chain but without sufficient confirmations
@@ -113,7 +110,6 @@ func NewStakingEventWatcher(
 		metrics:                         metrics,
 		unbondingTracker:                NewTrackedDelegations(),
 		pendingTracker:                  NewTrackedDelegations(),
-		inProgressTracker:               NewTrackedDelegations(),
 		verifiedInsufficientConfTracker: NewTrackedDelegations(),
 		verifiedNotInChainTracker:       NewTrackedDelegations(),
 		verifiedSufficientConfTracker:   NewTrackedDelegations(),
@@ -620,11 +616,11 @@ func (sew *StakingEventWatcher) checkBtcForStakingTx() {
 	}
 
 	for del := range sew.pendingTracker.DelegationsIter(1000) {
-		if inProgDel, _ := sew.inProgressTracker.GetDelegation(del.StakingTx.TxHash()); inProgDel != nil {
+		if del.ActivationInProgress {
 			continue
 		}
-
 		txHash := del.StakingTx.TxHash()
+
 		details, status, err := sew.btcClient.TxDetails(&txHash, del.StakingTx.TxOut[del.StakingOutputIdx].PkScript)
 		if err != nil {
 			sew.logger.Debugf("error getting tx %v", txHash)
@@ -661,8 +657,9 @@ func (sew *StakingEventWatcher) checkBtcForStakingTx() {
 			continue
 		}
 
-		if err = sew.inProgressTracker.AddEmptyDelegation(del.StakingTx.TxHash()); err != nil {
-			sew.logger.Debugf("add del: %s", err)
+		if err := sew.pendingTracker.UpdateActivation(txHash, true); err != nil {
+			sew.logger.Debugf("error updating activation in pending tracker tx: %v", txHash)
+			sew.activationLimiter.Release(1) // in probable edge case, insure we release the sem
 
 			continue
 		}
@@ -688,7 +685,11 @@ func (sew *StakingEventWatcher) activateBtcDelegation(
 	defer cancel()
 
 	defer sew.latency("activateBtcDelegation")()
-	defer sew.inProgressTracker.RemoveDelegation(stakingTxHash)
+	defer func() {
+		if err := sew.pendingTracker.UpdateActivation(stakingTxHash, false); err != nil {
+			sew.logger.Debugf("err updating activation in pending tracker tx: %v", stakingTxHash)
+		}
+	}()
 
 	if err := sew.waitForRequiredDepth(ctx, stakingTxHash, &inclusionBlockHash, requiredDepth); err != nil {
 		sew.logger.Warnf("exceeded waiting for required depth for tx: %s, will try later: err %v", stakingTxHash.String(), err)

btcstaking-tracker/stakingeventwatcher/stakingeventwatcher_test.go

@@ -40,13 +40,12 @@ func TestHandlingDelegations(t *testing.T) {
 		btcClient:                       mockBTCClient,
 		unbondingTracker:                NewTrackedDelegations(),
 		pendingTracker:                  NewTrackedDelegations(),
-		inProgressTracker:               NewTrackedDelegations(),
 		verifiedInsufficientConfTracker: NewTrackedDelegations(),
 		verifiedNotInChainTracker:       NewTrackedDelegations(),
 		verifiedSufficientConfTracker:   NewTrackedDelegations(),
 		unbondingDelegationChan:         make(chan *newDelegation),
 		unbondingRemovalChan:            make(chan *delegationInactive),
-		activationLimiter:               semaphore.NewWeighted(10),
+		activationLimiter:               semaphore.NewWeighted(30),
 		metrics:                         bsMetrics.UnbondingWatcherMetrics,
 	}
 

btcstaking-tracker/stakingeventwatcher/tracked_delegations.go

@@ -2,11 +2,11 @@ package stakingeventwatcher
 
 import (
 	"fmt"
-	"iter"
-	"sync"
-
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/sasha-s/go-deadlock"
+	"go.uber.org/atomic"
+	"iter"
 )
 
 type TrackedDelegation struct {
@@ -18,9 +18,10 @@ type TrackedDelegation struct {
 }
 
 type TrackedDelegations struct {
-	mu sync.RWMutex
+	mu deadlock.RWMutex
 	// key: staking tx hash
 	mapping map[chainhash.Hash]*TrackedDelegation
+	count   atomic.Int32
 }
 
 func NewTrackedDelegations() *TrackedDelegations {
@@ -168,6 +169,7 @@ func (td *TrackedDelegations) AddDelegation(
 	}
 
 	td.mapping[stakingTxHash] = delegation
+	td.count.Inc()
 
 	return delegation, nil
 }
@@ -181,6 +183,7 @@ func (td *TrackedDelegations) AddEmptyDelegation(txHash chainhash.Hash) error {
 	}
 
 	td.mapping[txHash] = nil
+	td.count.Inc()
 
 	return nil
 }
@@ -189,7 +192,10 @@ func (td *TrackedDelegations) RemoveDelegation(stakingTxHash chainhash.Hash) {
 	td.mu.Lock()
 	defer td.mu.Unlock()
 
-	delete(td.mapping, stakingTxHash)
+	if _, exists := td.mapping[stakingTxHash]; exists {
+		delete(td.mapping, stakingTxHash)
+		td.count.Dec()
+	}
 }
 
 func (td *TrackedDelegations) HasDelegationChanged(
@@ -220,7 +226,6 @@ func (td *TrackedDelegations) UpdateActivation(tx chainhash.Hash, inProgress boo
 	defer td.mu.Unlock()
 
 	delegation, ok := td.mapping[tx]
-
 	if !ok {
 		return fmt.Errorf("delegation with tx hash %s not found", tx.String())
 	}
@@ -231,8 +236,5 @@ func (td *TrackedDelegations) UpdateActivation(tx chainhash.Hash, inProgress boo
 }
 
 func (td *TrackedDelegations) Count() int {
-	td.mu.RLock()
-	defer td.mu.RUnlock()
-
-	return len(td.mapping)
+	return int(td.count.Load())
 }

go.mod

@@ -33,6 +33,7 @@ require (
 	github.com/ory/dockertest/v3 v3.10.0
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.20.5
+	github.com/sasha-s/go-deadlock v0.3.5
 	github.com/spf13/cobra v1.8.1
 	github.com/spf13/viper v1.19.0
 	github.com/stretchr/testify v1.9.0
@@ -266,7 +267,6 @@ require (
 	github.com/rs/zerolog v1.33.0 // indirect
 	github.com/sagikazarmark/locafero v0.4.0 // indirect
 	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
-	github.com/sasha-s/go-deadlock v0.3.5 // indirect
 	github.com/shamaton/msgpack/v2 v2.2.0 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/soheilhy/cmux v0.1.5 // indirect