fix: take pending rewards for subgraph service

Signed-off-by: Tomás Migone <[email protected]>
graphprotocol · May 23, 2024 · 04432d2 · 04432d2
1 parent 187daaf
commit 04432d2
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 5 deletions.
diff --git a/packages/subgraph-service/contracts/SubgraphService.sol b/packages/subgraph-service/contracts/SubgraphService.sol
@@ -205,7 +205,12 @@ contract SubgraphService is
         address allocationId
     ) external view override returns (address, bytes32, uint256, uint256) {
         Allocation.State memory allo = allocations[allocationId];
-        return (allo.indexer, allo.subgraphDeploymentId, allo.tokens, allo.accRewardsPerAllocatedToken);
+        return (
+            allo.indexer,
+            allo.subgraphDeploymentId,
+            allo.tokens,
+            allo.accRewardsPerAllocatedToken + allo.accRewardsPending
+        );
     }
 
     function getLegacyAllocation(address allocationId) external view returns (LegacyAllocation.State memory) {

diff --git a/packages/subgraph-service/contracts/libraries/Allocation.sol b/packages/subgraph-service/contracts/libraries/Allocation.sol
@@ -69,6 +69,17 @@ library Allocation {
         return allocation;
     }
 
+    function clearPendingRewards(
+        mapping(address => State) storage self,
+        address allocationId
+    ) internal returns (State memory) {
+        State storage allocation = _get(self, allocationId);
+        if (!allocation.isOpen()) revert AllocationClosed(allocationId, allocation.closedAt);
+        allocation.accRewardsPending = 0;
+
+        return allocation;
+    }
+
     function close(mapping(address => State) storage self, address allocationId) internal returns (State memory) {
         State storage allocation = _get(self, allocationId);
         if (!allocation.isOpen()) revert AllocationClosed(allocationId, allocation.closedAt);

diff --git a/packages/subgraph-service/contracts/utilities/AllocationManager.sol b/packages/subgraph-service/contracts/utilities/AllocationManager.sol
@@ -147,6 +147,9 @@ abstract contract AllocationManager is EIP712, GraphDirectory, AllocationManager
         );
         allocations.presentPOI(allocationId);
 
+        // Any pending rewards should have been collected now
+        allocations.clearPendingRewards(allocationId);
+
         if (tokensRewards == 0) {
             emit IndexingRewardsCollected(
                 allocation.indexer,
@@ -225,12 +228,11 @@ abstract contract AllocationManager is EIP712, GraphDirectory, AllocationManager
         // Update the allocation
         allocations[_allocationId].tokens = _tokens;
         allocations[_allocationId].accRewardsPerAllocatedToken = accRewardsPerAllocatedToken;
-        allocations[_allocationId].accRewardsPending = allocations[_allocationId].accRewardsPending + accRewardsPending;
+        allocations[_allocationId].accRewardsPending += accRewardsPending;
 
         // Update total allocated tokens for the subgraph deployment
-        subgraphAllocatedTokens[allocation.subgraphDeploymentId] =
-            subgraphAllocatedTokens[allocation.subgraphDeploymentId] +
-            (_tokens - oldTokens);
+        // underflow: subgraphAllocatedTokens should at least be oldTokens so it can't underflow
+        subgraphAllocatedTokens[allocation.subgraphDeploymentId] += (_tokens - oldTokens);
 
         emit AllocationResized(allocation.indexer, _allocationId, allocation.subgraphDeploymentId, _tokens, oldTokens);
         return allocations[_allocationId];
@@ -242,6 +244,13 @@ abstract contract AllocationManager is EIP712, GraphDirectory, AllocationManager
         allocations.close(_allocationId);
         allocationProvisionTracker.release(allocation.indexer, allocation.tokens);
 
+        // Take rewards snapshot to prevent other allos from counting tokens from this allo
+        allocations.snapshotRewards(
+            _allocationId,
+            _graphRewardsManager().onSubgraphAllocationUpdate(allocation.subgraphDeploymentId)
+        );
+
+        // Update total allocated tokens for the subgraph deployment
         subgraphAllocatedTokens[allocation.subgraphDeploymentId] =
             subgraphAllocatedTokens[allocation.subgraphDeploymentId] -
             allocation.tokens;