fix: merge close stale and over allocated functions

packages/subgraph-service/contracts/SubgraphService.sol

@@ -307,28 +307,15 @@ contract SubgraphService is
     /**
      * @notice See {ISubgraphService.closeStaleAllocation}
      */
-    function closeStaleAllocation(address allocationId) external override {
+    function forceCloseAllocation(address allocationId) external override {
         Allocation.State memory allocation = allocations.get(allocationId);
-        require(
-            allocation.isStale(maxPOIStaleness),
-            SubgraphServiceAllocationNotStale(allocationId, allocation.lastPOIPresentedAt)
-        );
+        bool isStale = allocation.isStale(maxPOIStaleness);
+        bool isOverAllocated_ = _isOverAllocated(allocation.indexer, delegationRatio);
+        require(isStale || isOverAllocated_, SubgraphServiceCannotForceCloseAllocation(allocationId));
         require(!allocation.isAltruistic(), SubgraphServiceAllocationIsAltruistic(allocationId));
         _closeAllocation(allocationId);
     }
 
-    /**
-     * @notice See {ISubgraphService.closeOverAllocatedAllocation}
-     */
-    function closeOverAllocatedAllocation(address allocationId) external override {
-        Allocation.State memory allocation = allocations.get(allocationId);
-        require(
-            _isAllocationOverAllocated(allocation.indexer, delegationRatio),
-            SubgraphServiceAllocationNotOverAllocated(allocationId)
-        );
-        _closeAllocation(allocationId);
-    }
-
     /**
      * @notice See {ISubgraphService.resizeAllocation}
      */
@@ -483,6 +470,20 @@ contract SubgraphService is
         return _encodeAllocationProof(indexer, allocationId);
     }
 
+    /**
+     * @notice See {ISubgraphService.isStaleAllocation}
+     */
+    function isStaleAllocation(address allocationId) external view override returns (bool) {
+        return allocations.get(allocationId).isStale(maxPOIStaleness);
+    }
+
+    /**
+     * @notice See {ISubgraphService.isOverAllocated}
+     */
+    function isOverAllocated(address indexer) external view override returns (bool) {
+        return _isOverAllocated(indexer, delegationRatio);
+    }
+
     // -- Data service parameter getters --
     /**
      * @notice Getter for the accepted thawing period range for provisions

packages/subgraph-service/contracts/interfaces/ISubgraphService.sol

@@ -112,9 +112,8 @@ interface ISubgraphService is IDataServiceFees {
     /**
      * @notice Thrown when trying to force close an allocation that is not stale
      * @param allocationId The id of the allocation
-     * @param lastPOIPresentedAt The timestamp when the last POI was presented
      */
-    error SubgraphServiceAllocationNotStale(address allocationId, uint256 lastPOIPresentedAt);
+    error SubgraphServiceCannotForceCloseAllocation(address allocationId);
 
     /**
      * @notice Thrown when trying to force close an altruistic allocation
@@ -134,34 +133,22 @@ interface ISubgraphService is IDataServiceFees {
     error SubgraphServiceInvalidZeroStakeToFeesRatio();
 
     /**
-     * @notice Close a stale allocation
-     * @dev This function can be permissionlessly called when the allocation is stale.
-     * This ensures rewards for other allocations are not diluted by an inactive allocation
+     * @notice Force close an allocation
+     * @dev This function can be permissionlessly called when the allocation is stale or
+     * if the indexer is over-allocated. This ensures that rewards for other allocations are
+     * not diluted by an inactive allocation, and that over-allocated indexers stop accumulating
+     * rewards with tokens they no longer have allocated.
      *
      * Requirements:
      * - Allocation must exist and be open
-     * - Allocation must be stale
+     * - Allocation must be stale or indexer must be over-allocated
      * - Allocation cannot be altruistic
      *
      * Emits a {AllocationClosed} event.
      *
      * @param allocationId The id of the allocation
      */
-    function closeStaleAllocation(address allocationId) external;
-
-    /**
-     * @notice Close an over-allocated allocation
-     * @dev This function can be permissionlessly called when the allocation is over-allocated.
-     * This ensures rewards for other allocations are not diluted by an over-allocated allocation
-     *
-     * Requirements:
-     * - Allocation must exist and be open
-     * - Allocation must be over-allocated
-     *
-     * Emits a {AllocationClosed} event.
-     * @param allocationId The id of the allocation
-     */
-    function closeOverAllocatedAllocation(address allocationId) external;
+    function forceCloseAllocation(address allocationId) external;
 
     /**
      * @notice Change the amount of tokens in an allocation
@@ -257,4 +244,18 @@ interface ISubgraphService is IDataServiceFees {
      * @param allocationId The id of the allocation
      */
     function encodeAllocationProof(address indexer, address allocationId) external view returns (bytes32);
+
+    /**
+     * @notice Checks if an allocation is stale
+     * @param allocationId The id of the allocation
+     * @return True if the allocation is stale, false otherwise
+     */
+    function isStaleAllocation(address allocationId) external view returns (bool);
+
+    /**
+     * @notice Checks if an indexer is over-allocated
+     * @param allocationId The id of the allocation
+     * @return True if the indexer is over-allocated, false otherwise
+     */
+    function isOverAllocated(address allocationId) external view returns (bool);
 }

packages/subgraph-service/contracts/utilities/AllocationManager.sol

@@ -328,7 +328,7 @@ abstract contract AllocationManager is EIP712Upgradeable, GraphDirectory, Alloca
         );
 
         // Check if the indexer is over-allocated and close the allocation if necessary
-        if (_isAllocationOverAllocated(allocation.indexer, _delegationRatio)) {
+        if (_isOverAllocated(allocation.indexer, _delegationRatio)) {
             _closeAllocation(_allocationId);
         }
 
@@ -492,7 +492,7 @@ abstract contract AllocationManager is EIP712Upgradeable, GraphDirectory, Alloca
      * @param _delegationRatio The delegation ratio to consider when locking tokens
      * @return True if the allocation is over-allocated, false otherwise
      */
-    function _isAllocationOverAllocated(address _indexer, uint32 _delegationRatio) internal view returns (bool) {
+    function _isOverAllocated(address _indexer, uint32 _delegationRatio) internal view returns (bool) {
         return !allocationProvisionTracker.check(_graphStaking(), _indexer, _delegationRatio);
     }
 }

packages/subgraph-service/test/subgraphService/SubgraphService.t.sol

@@ -151,7 +151,7 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
         assertEq(afterSubgraphAllocatedTokens, _tokens);
     }
 
-    function _closeStaleAllocation(address _allocationId) internal {
+    function _forceCloseAllocation(address _allocationId) internal {
         assertTrue(subgraphService.isActiveAllocation(_allocationId));
 
         Allocation.State memory allocation = subgraphService.getAllocation(_allocationId);
@@ -168,37 +168,7 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
         );
 
         // close stale allocation
-        subgraphService.closeStaleAllocation(_allocationId);
-
-        // update allocation
-        allocation = subgraphService.getAllocation(_allocationId);
-
-        // check allocation
-        assertEq(allocation.closedAt, block.timestamp);
-
-        // check subgraph deployment allocated tokens
-        uint256 subgraphAllocatedTokens = subgraphService.getSubgraphAllocatedTokens(subgraphDeployment);
-        assertEq(subgraphAllocatedTokens, previousSubgraphAllocatedTokens - allocation.tokens);
-    }
-
-    function _closeOverAllocatedAllocation(address _allocationId) internal {
-        assertTrue(subgraphService.isActiveAllocation(_allocationId));
-
-        Allocation.State memory allocation = subgraphService.getAllocation(_allocationId);
-        uint256 previousSubgraphAllocatedTokens = subgraphService.getSubgraphAllocatedTokens(
-            allocation.subgraphDeploymentId
-        );
-
-        vm.expectEmit(address(subgraphService));
-        emit AllocationManager.AllocationClosed(
-            allocation.indexer,
-            _allocationId,
-            allocation.subgraphDeploymentId,
-            allocation.tokens
-        );
-
-        // close over allocated allocation
-        subgraphService.closeOverAllocatedAllocation(_allocationId);
+        subgraphService.forceCloseAllocation(_allocationId);
 
         // update allocation
         allocation = subgraphService.getAllocation(_allocationId);

packages/subgraph-service/test/subgraphService/allocation/closeStale.t.sol → packages/subgraph-service/test/subgraphService/allocation/forceClose.t.sol

@@ -14,25 +14,25 @@ import { ISubgraphService } from "../../../contracts/interfaces/ISubgraphService
 import { LegacyAllocation } from "../../../contracts/libraries/LegacyAllocation.sol";
 import { SubgraphServiceTest } from "../SubgraphService.t.sol";
 
-contract SubgraphServiceAllocationCloseStaleTest is SubgraphServiceTest {
+contract SubgraphServiceAllocationForceCloseTest is SubgraphServiceTest {
 
     address private permissionlessBob = makeAddr("permissionlessBob");
 
     /*
      * TESTS
      */
 
-    function test_SubgraphService_Allocation_CloseStale(
+    function test_SubgraphService_Allocation_ForceClose_Stale(
         uint256 tokens
-    ) public useIndexer useAllocation(tokens) {      
-        // Skip forward  
+    ) public useIndexer useAllocation(tokens) {
+        // Skip forward
         skip(maxPOIStaleness + 1);
 
         resetPrank(permissionlessBob);
-        _closeStaleAllocation(allocationID);
+        _forceCloseAllocation(allocationID);
     }
 
-    function test_SubgraphService_Allocation_CloseStale_AfterCollecting(
+    function test_SubgraphService_Allocation_ForceClose_Stale_AfterCollecting(
         uint256 tokens
     ) public useIndexer useAllocation(tokens) {      
         // Simulate POIs being submitted
@@ -52,10 +52,43 @@ contract SubgraphServiceAllocationCloseStaleTest is SubgraphServiceTest {
 
         // Close the stale allocation
         resetPrank(permissionlessBob);
-        _closeStaleAllocation(allocationID);
+        _forceCloseAllocation(allocationID);
+    }
+
+    function test_SubgraphService_Allocation_ForceClose_OverAllocated(
+        uint256 tokens
+    ) public useIndexer useAllocation(tokens) { 
+        // thaw some tokens to become over allocated
+        staking.thaw(users.indexer, address(subgraphService), tokens / 2);
+
+        resetPrank(permissionlessBob);
+        _forceCloseAllocation(allocationID);
+    }
+
+    function test_SubgraphService_Allocation_ForceClose_OverAllocated_AfterCollecting(
+        uint256 tokens
+    ) public useIndexer useAllocation(tokens) {      
+        // Simulate POIs being submitted
+        uint8 numberOfPOIs = 5;
+        uint256 timeBetweenPOIs = 5 days;
+
+        for (uint8 i = 0; i < numberOfPOIs; i++) {
+            // Skip forward
+            skip(timeBetweenPOIs);
+
+            bytes memory data = abi.encode(allocationID, bytes32("POI1"));
+            _collect(users.indexer, IGraphPayments.PaymentTypes.IndexingRewards, data);
+        }
+
+        // thaw some tokens to become over allocated
+        staking.thaw(users.indexer, address(subgraphService), tokens / 2);
+
+        // Close the over allocated allocation
+        resetPrank(permissionlessBob);
+        _forceCloseAllocation(allocationID);
     }
 
-    function test_SubgraphService_Allocation_CloseStale_RevertIf_NotStale(
+    function test_SubgraphService_Allocation_ForceClose_RevertIf_NotStaleOrOverAllocated(
         uint256 tokens
     ) public useIndexer useAllocation(tokens) {
         // Simulate POIs being submitted
@@ -74,16 +107,15 @@ contract SubgraphServiceAllocationCloseStaleTest is SubgraphServiceTest {
             resetPrank(permissionlessBob);
             vm.expectRevert(
                 abi.encodeWithSelector(
-                    ISubgraphService.SubgraphServiceAllocationNotStale.selector,
-                    allocationID,
-                    block.timestamp
+                    ISubgraphService.SubgraphServiceCannotForceCloseAllocation.selector,
+                    allocationID
                 )
             );
-            subgraphService.closeStaleAllocation(allocationID);
+            subgraphService.forceCloseAllocation(allocationID);
         }
     }
 
-    function test_SubgraphService_Allocation_CloseStale_RevertIf_Altruistic(
+    function test_SubgraphService_Allocation_ForceClose_RevertIf_Altruistic(
         uint256 tokens
     ) public useIndexer {
         tokens = bound(tokens, minimumProvisionTokens, MAX_TOKENS);
@@ -103,6 +135,6 @@ contract SubgraphServiceAllocationCloseStaleTest is SubgraphServiceTest {
                 allocationID
             )
         );
-        subgraphService.closeStaleAllocation(allocationID);
+        subgraphService.forceCloseAllocation(allocationID);
     }
 }