fix fp stats

babylonlabs-io · Nov 30, 2024 · 4929932 · 4929932
1 parent bb96c06
commit 4929932
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 22 deletions.
diff --git a/internal/v2/db/client/interface.go b/internal/v2/db/client/interface.go
@@ -27,6 +27,6 @@ type V2DBClient interface {
 		ctx context.Context, stakingTxHashHex, stakerPkHex string, amount uint64,
 	) error
 	IncrementFinalityProviderStats(
-		ctx context.Context, stakingTxHashHex, fpPkHex string, amount uint64,
+		ctx context.Context, stakingTxHashHex string, fpPkHexes []string, amount uint64,
 	) error
 }
diff --git a/internal/v2/db/client/stats.go b/internal/v2/db/client/stats.go
@@ -309,49 +309,72 @@ func (v2dbclient *V2Database) GetStakerStats(
 }
 
 func (v2dbclient *V2Database) IncrementFinalityProviderStats(
-	ctx context.Context, stakingTxHashHex, fpPkHex string, amount uint64,
+	ctx context.Context,
+	stakingTxHashHex string,
+	fpPkHexes []string,
+	amount uint64,
 ) error {
-	upsertUpdate := bson.M{
-		"$inc": bson.M{
-			"active_tvl":         int64(amount),
-			"total_tvl":          int64(amount),
-			"active_delegations": 1,
-			"total_delegations":  1,
-		},
+	// Create bulk write operations for each FP
+	var operations []mongo.WriteModel
+	for _, fpPkHex := range fpPkHexes {
+		operation := mongo.NewUpdateOneModel().
+			SetFilter(bson.M{"_id": fpPkHex}).
+			SetUpdate(bson.M{
+				"$inc": bson.M{
+					"active_tvl":         int64(amount),
+					"total_tvl":          int64(amount),
+					"active_delegations": 1,
+					"total_delegations":  1,
+				},
+			}).
+			SetUpsert(true)
+		operations = append(operations, operation)
 	}
-	return v2dbclient.updateFinalityProviderStats(ctx, types.Active.ToString(), stakingTxHashHex, fpPkHex, upsertUpdate)
+
+	return v2dbclient.updateFinalityProviderStats(
+		ctx,
+		types.Active.ToString(),
+		stakingTxHashHex,
+		operations,
+	)
 }
 
-func (v2dbclient *V2Database) updateFinalityProviderStats(ctx context.Context, state, stakingTxHashHex, fpPkHex string, upsertUpdate primitive.M) error {
+func (v2dbclient *V2Database) updateFinalityProviderStats(
+	ctx context.Context,
+	state string,
+	stakingTxHashHex string,
+	operations []mongo.WriteModel,
+) error {
 	client := v2dbclient.Client.Database(v2dbclient.DbName).Collection(dbmodel.V2FinalityProviderStatsCollection)
 
-	// Start a session
 	session, sessionErr := v2dbclient.Client.StartSession()
 	if sessionErr != nil {
 		return sessionErr
 	}
 	defer session.EndSession(ctx)
 
 	transactionWork := func(sessCtx mongo.SessionContext) (interface{}, error) {
-		err := v2dbclient.updateStatsLockByFieldName(sessCtx, stakingTxHashHex, state, "finality_provider_stats")
+		// Single lock for the entire operation
+		err := v2dbclient.updateStatsLockByFieldName(
+			sessCtx,
+			stakingTxHashHex,
+			state,
+			"finality_provider_stats",
+		)
 		if err != nil {
 			return nil, err
 		}
 
-		upsertFilter := bson.M{"_id": fpPkHex}
-
-		_, err = client.UpdateOne(sessCtx, upsertFilter, upsertUpdate, options.Update().SetUpsert(true))
+		// Execute all updates in a single bulk write
+		opts := options.BulkWrite().SetOrdered(true)
+		_, err = client.BulkWrite(sessCtx, operations, opts)
 		if err != nil {
 			return nil, err
 		}
+
 		return nil, nil
 	}
 
-	// Execute the transaction
 	_, txErr := session.WithTransaction(ctx, transactionWork)
-	if txErr != nil {
-		return txErr
-	}
-
-	return nil
+	return txErr
 }