From a98c580cade8cf1fe31568ee498018bef6e46fc3 Mon Sep 17 00:00:00 2001
From: Matt Hess <matt.hess@swirldslabs.com>
Date: Fri, 19 Apr 2024 16:22:55 -0500
Subject: [PATCH] fix: Don't double-count child staking rewards in parent
 finalization (#12905)

Signed-off-by: Matt Hess <matt.hess@swirldslabs.com>
---
 .../app/workflows/handle/HandleWorkflow.java  | 24 ++++++++++++++-----
 .../staking/StakingRewardsHandlerImpl.java    | 12 ++++++++--
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/hedera-node/hedera-app/src/main/java/com/hedera/node/app/workflows/handle/HandleWorkflow.java b/hedera-node/hedera-app/src/main/java/com/hedera/node/app/workflows/handle/HandleWorkflow.java
index 0568c7d9f2ef..32da991891d0 100644
--- a/hedera-node/hedera-app/src/main/java/com/hedera/node/app/workflows/handle/HandleWorkflow.java
+++ b/hedera-node/hedera-app/src/main/java/com/hedera/node/app/workflows/handle/HandleWorkflow.java
@@ -638,12 +638,24 @@ private void handleUserTransaction(
         }
 
         throttleServiceManager.saveThrottleSnapshotsAndCongestionLevelStartsTo(stack);
-        transactionFinalizer.finalizeParentRecord(
-                payer,
-                tokenServiceContext,
-                functionality,
-                extraRewardReceivers(txBody, functionality, recordBuilder),
-                prePaidRewards);
+        try {
+            transactionFinalizer.finalizeParentRecord(
+                    payer,
+                    tokenServiceContext,
+                    functionality,
+                    extraRewardReceivers(txBody, functionality, recordBuilder),
+                    prePaidRewards);
+        } catch (final Exception e) {
+            logger.error(
+                    "Possibly CATASTROPHIC error: failed to finalize parent record for transaction {}",
+                    transactionInfo.transactionID(),
+                    e);
+
+            // Undo any changes made to the state
+            final var userTransactionRecordBuilder = recordListBuilder.userTransactionRecordBuilder();
+            userTransactionRecordBuilder.nullOutSideEffectFields();
+            rollback(true, ResponseCodeEnum.FAIL_INVALID, stack, recordListBuilder);
+        }
 
         // Commit all state changes
         stack.commitFullStack();
diff --git a/hedera-node/hedera-token-service-impl/src/main/java/com/hedera/node/app/service/token/impl/handlers/staking/StakingRewardsHandlerImpl.java b/hedera-node/hedera-token-service-impl/src/main/java/com/hedera/node/app/service/token/impl/handlers/staking/StakingRewardsHandlerImpl.java
index b99bef80c896..3afb6b526e85 100644
--- a/hedera-node/hedera-token-service-impl/src/main/java/com/hedera/node/app/service/token/impl/handlers/staking/StakingRewardsHandlerImpl.java
+++ b/hedera-node/hedera-token-service-impl/src/main/java/com/hedera/node/app/service/token/impl/handlers/staking/StakingRewardsHandlerImpl.java
@@ -95,6 +95,9 @@ public Map<AccountID, Long> applyStakingRewards(
         final var rewardsPaid = rewardsPayer.payRewardsIfPending(
                 rewardReceivers, writableStore, stakingRewardsStore, stakingInfoStore, consensusNow, recordBuilder);
 
+        // Decrease staking reward account balance by rewardPaid amount
+        decreaseStakeRewardAccountBalance(rewardsPaid, stakingRewardAccountId, writableStore);
+
         if (!context.isScheduleDispatch()) {
             // We only manage stake metadata once, at the end of a transaction; but to do
             // this correctly, we need to include information about any rewards paid during
@@ -107,10 +110,15 @@ public Map<AccountID, Long> applyStakingRewards(
             // Adjust stakes for nodes and account's staking metadata
             adjustStakeMetadata(
                     writableStore, stakingInfoStore, stakingRewardsStore, consensusNow, rewardsPaid, rewardReceivers);
+
+            // Don't double-report prepaid rewards in the parent record
+            if (!prePaidRewards.isEmpty()) {
+                for (AccountID accountID : prePaidRewards.keySet()) {
+                    rewardsPaid.remove(accountID);
+                }
+            }
         }
 
-        // Decrease staking reward account balance by rewardPaid amount
-        decreaseStakeRewardAccountBalance(rewardsPaid, stakingRewardAccountId, writableStore);
         return rewardsPaid;
     }