fix: ensure rewards manager looks at subgraph service and staking for…

packages/contracts/contracts/rewards/IRewardsIssuer.sol

@@ -5,9 +5,9 @@ pragma solidity ^0.7.6 || 0.8.26;
 interface IRewardsIssuer {
     /**
      * @dev Get allocation data to calculate rewards issuance
-     * @param allocationId The allocation ID
+     * @param allocationId The allocation Id
      * @return indexer The indexer address
-     * @return subgraphDeploymentID Subgraph deployment id for the allocation
+     * @return subgraphDeploymentId Subgraph deployment id for the allocation
      * @return tokens Amount of allocated tokens
      * @return accRewardsPerAllocatedToken Rewards snapshot
      */
@@ -16,5 +16,19 @@ interface IRewardsIssuer {
     )
         external
         view
-        returns (address indexer, bytes32 subgraphDeploymentID, uint256 tokens, uint256 accRewardsPerAllocatedToken);
+        returns (address indexer, bytes32 subgraphDeploymentId, uint256 tokens, uint256 accRewardsPerAllocatedToken);
+
+    /**
+     * @notice Return the total amount of tokens allocated to subgraph.
+     * @param _subgraphDeploymentId Deployment Id for the subgraph
+     * @return Total tokens allocated to subgraph
+     */
+    function getSubgraphAllocatedTokens(bytes32 _subgraphDeploymentId) external view returns (uint256);
+
+    /**
+     * @notice Wether or not an allocation is active (i.e open)
+     * @param _allocationId Allocation Id
+     * @return Wether or not the allocation is active
+     */
+    function isActiveAllocation(address _allocationId) external view returns (bool);
 }

packages/contracts/contracts/rewards/IRewardsManager.sol

@@ -19,7 +19,7 @@ interface IRewardsManager {
 
     function setMinimumSubgraphSignal(uint256 _minimumSubgraphSignal) external;
 
-    function setRewardsIssuer(address _rewardsIssuer, bool _allowed) external;
+    function setSubgraphService(address _subgraphService) external;
 
     // -- Denylist --
 

packages/contracts/contracts/rewards/RewardsManager.sol

@@ -52,9 +52,9 @@ contract RewardsManager is RewardsManagerV5Storage, GraphUpgradeable, IRewardsMa
     event RewardsDenylistUpdated(bytes32 indexed subgraphDeploymentID, uint256 sinceBlock);
 
     /**
-     * @dev Emitted when a rewards issuer is updated.
+     * @dev Emitted when the subgraph service is set
      */
-    event RewardsIssuerSet(address indexed rewardsIssuer, bool allowed);
+    event SubgraphServiceSet(address indexed oldSubgraphService, address indexed newSubgraphService);
 
     // -- Modifiers --
 
@@ -121,9 +121,10 @@ contract RewardsManager is RewardsManagerV5Storage, GraphUpgradeable, IRewardsMa
         emit ParameterUpdated("minimumSubgraphSignal");
     }
 
-    function setRewardsIssuer(address _rewardsIssuer, bool _allowed) external override onlyGovernor {
-        rewardsIssuers[_rewardsIssuer] = _allowed;
-        emit RewardsIssuerSet(_rewardsIssuer, _allowed);
+    function setSubgraphService(address _subgraphService) external override onlyGovernor {
+        address oldSubgraphService = address(subgraphService);
+        subgraphService = IRewardsIssuer(_subgraphService);
+        emit SubgraphServiceSet(oldSubgraphService, _subgraphService);
     }
 
     // -- Denylist --
@@ -258,7 +259,19 @@ contract RewardsManager is RewardsManagerV5Storage, GraphUpgradeable, IRewardsMa
             subgraph.accRewardsForSubgraphSnapshot
         );
 
-        uint256 subgraphAllocatedTokens = staking().getSubgraphAllocatedTokens(_subgraphDeploymentID);
+        // There are two contributors to subgraph allocated tokens:
+        // - the legacy allocations on the legacy staking contract
+        // - the new allocations on the subgraph service
+        uint256 subgraphAllocatedTokens = 0;
+        address[2] memory rewardsIssuers = [address(staking()), address(subgraphService)];
+        for (uint256 i = 0; i < rewardsIssuers.length; i++) {
+            if (rewardsIssuers[i] != address(0)) {
+                subgraphAllocatedTokens += IRewardsIssuer(rewardsIssuers[i]).getSubgraphAllocatedTokens(
+                    _subgraphDeploymentID
+                );
+            }
+        }
+
         if (subgraphAllocatedTokens == 0) {
             return (0, accRewardsForSubgraph);
         }
@@ -321,19 +334,35 @@ contract RewardsManager is RewardsManagerV5Storage, GraphUpgradeable, IRewardsMa
 
     /**
      * @dev Calculate current rewards for a given allocation on demand.
+     * The allocation could be a legacy allocation or a new subgraph service allocation.
      * @param _allocationID Allocation
      * @return Rewards amount for an allocation
      */
     function getRewards(address _allocationID) external view override returns (uint256) {
-        IStaking.AllocationState allocState = staking().getAllocationState(_allocationID);
-        if (allocState != IStakingBase.AllocationState.Active) {
+        address rewardsIssuer = address(0);
+
+        // Check both the legacy and new allocations
+        address[2] memory rewardsIssuers = [address(staking()), address(subgraphService)];
+        for (uint256 i = 0; i < rewardsIssuers.length; i++) {
+            if (rewardsIssuers[i] != address(0)) {
+                if (IRewardsIssuer(rewardsIssuers[i]).isActiveAllocation(_allocationID)) {
+                    rewardsIssuer = address(rewardsIssuers[i]);
+                    break;
+                }
+            }
+        }
+
+        // Bail if allo does not exist
+        if (rewardsIssuer == address(0)) {
             return 0;
         }
 
-        IStaking.Allocation memory alloc = staking().getAllocation(_allocationID);
+        (, bytes32 subgraphDeploymentId, uint256 tokens, uint256 alloAccRewardsPerAllocatedToken) = IRewardsIssuer(
+            rewardsIssuer
+        ).getAllocationData(_allocationID);
 
-        (uint256 accRewardsPerAllocatedToken, ) = getAccRewardsPerAllocatedToken(alloc.subgraphDeploymentID);
-        return _calcRewards(alloc.tokens, alloc.accRewardsPerAllocatedToken, accRewardsPerAllocatedToken);
+        (uint256 accRewardsPerAllocatedToken, ) = getAccRewardsPerAllocatedToken(subgraphDeploymentId);
+        return _calcRewards(tokens, alloAccRewardsPerAllocatedToken, accRewardsPerAllocatedToken);
     }
 
     /**
@@ -354,14 +383,18 @@ contract RewardsManager is RewardsManagerV5Storage, GraphUpgradeable, IRewardsMa
 
     /**
      * @dev Pull rewards from the contract for a particular allocation.
-     * This function can only be called by the Staking contract.
+     * This function can only be called by an authorized rewards issuer which are
+     * the staking contract (for legacy allocations), and the subgraph service (for new allocations).
      * This function will mint the necessary tokens to reward based on the inflation calculation.
      * @param _allocationID Allocation
      * @return Assigned rewards amount
      */
     function takeRewards(address _allocationID) external override returns (uint256) {
         address rewardsIssuer = msg.sender;
-        require(rewardsIssuers[rewardsIssuer], "Caller must be a rewards issuer");
+        require(
+            rewardsIssuer == address(staking()) || rewardsIssuer == address(subgraphService),
+            "Caller must be a rewards issuer"
+        );
 
         (
             address indexer,

packages/contracts/contracts/rewards/RewardsManagerStorage.sol

@@ -1,9 +1,10 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
-pragma solidity ^0.7.6;
+pragma solidity ^0.7.6 || 0.8.26;
 
 import "./IRewardsManager.sol";
 import "../governance/Managed.sol";
+import { IRewardsIssuer } from "./IRewardsIssuer.sol";
 
 contract RewardsManagerV1Storage is Managed {
     // -- State --
@@ -38,6 +39,6 @@ contract RewardsManagerV4Storage is RewardsManagerV3Storage {
 }
 
 contract RewardsManagerV5Storage is RewardsManagerV4Storage {
-    // List of addresses that are allowed to issue rewards
-    mapping(address => bool) public rewardsIssuers;
+    // Address of the subgraph service
+    IRewardsIssuer public subgraphService;
 }

packages/contracts/contracts/staking/IStakingBase.sol

@@ -361,6 +361,20 @@ interface IStakingBase is IStakingData {
      */
     function getAllocation(address _allocationID) external view returns (Allocation memory);
 
+    /**
+     * @dev New function to get the allocation data for the rewards manager
+     * @dev Note that this is only to make tests pass, as the staking contract with
+     * this changes will never get deployed. HorizonStaking is taking it's place.
+     */
+    function getAllocationData(address _allocationID) external view returns (address, bytes32, uint256, uint256);
+
+    /**
+     * @dev New function to get the allocation active status for the rewards manager
+     * @dev Note that this is only to make tests pass, as the staking contract with
+     * this changes will never get deployed. HorizonStaking is taking it's place.
+     */
+    function isActiveAllocation(address _allocationID) external view returns (bool);
+
     /**
      * @notice Return the current state of an allocation
      * @param _allocationID Allocation identifier

packages/contracts/contracts/staking/Staking.sol

@@ -472,6 +472,27 @@ abstract contract Staking is StakingV4Storage, GraphUpgradeable, IStakingBase, M
         return __allocations[_allocationID];
     }
 
+    /**
+     * @dev New function to get the allocation data for the rewards manager
+     * @dev Note that this is only to make tests pass, as the staking contract with
+     * this changes will never get deployed. HorizonStaking is taking it's place.
+     */
+    function getAllocationData(
+        address _allocationID
+    ) external view override returns (address, bytes32, uint256, uint256) {
+        Allocation memory alloc = __allocations[_allocationID];
+        return (alloc.indexer, alloc.subgraphDeploymentID, alloc.tokens, alloc.accRewardsPerAllocatedToken);
+    }
+
+    /**
+     * @dev New function to get the allocation active status for the rewards manager
+     * @dev Note that this is only to make tests pass, as the staking contract with
+     * this changes will never get deployed. HorizonStaking is taking it's place.
+     */
+    function isActiveAllocation(address _allocationID) external view override returns (bool) {
+        return _getAllocationState(_allocationID) == AllocationState.Active;
+    }
+
     /**
      * @notice Return the current state of an allocation
      * @param _allocationID Allocation identifier

packages/contracts/test/unit/rewards/rewards.test.ts

@@ -653,7 +653,6 @@ describe('Rewards', () => {
         const event = rewardsManager.interface.parseLog(receipt.logs[1]).args
         expect(event.indexer).eq(indexer1.address)
         expect(event.allocationID).eq(allocationID1)
-        expect(event.epoch).eq(await epochManager.currentEpoch())
         expect(toRound(event.amount)).eq(toRound(expectedIndexingRewards))
 
         // After state
@@ -692,7 +691,7 @@ describe('Rewards', () => {
         const tx = staking.connect(indexer1).closeAllocation(allocationID1, randomHexBytes())
         await expect(tx)
           .emit(rewardsManager, 'RewardsAssigned')
-          .withArgs(indexer1.address, allocationID1, await epochManager.currentEpoch(), toBN(0))
+          .withArgs(indexer1.address, allocationID1, toBN(0))
       })
 
       it('does not revert with an underflow if the minimum signal changes, and signal came after allocation', async function () {
@@ -710,7 +709,7 @@ describe('Rewards', () => {
         const tx = staking.connect(indexer1).closeAllocation(allocationID1, randomHexBytes())
         await expect(tx)
           .emit(rewardsManager, 'RewardsAssigned')
-          .withArgs(indexer1.address, allocationID1, await epochManager.currentEpoch(), toBN(0))
+          .withArgs(indexer1.address, allocationID1, toBN(0))
       })
 
       it('does not revert if signal was already under minimum', async function () {
@@ -727,7 +726,7 @@ describe('Rewards', () => {
 
         await expect(tx)
           .emit(rewardsManager, 'RewardsAssigned')
-          .withArgs(indexer1.address, allocationID1, await epochManager.currentEpoch(), toBN(0))
+          .withArgs(indexer1.address, allocationID1, toBN(0))
       })
 
       it('should distribute rewards on closed allocation and send to destination', async function () {
@@ -761,7 +760,6 @@ describe('Rewards', () => {
         const event = rewardsManager.interface.parseLog(receipt.logs[1]).args
         expect(event.indexer).eq(indexer1.address)
         expect(event.allocationID).eq(allocationID1)
-        expect(event.epoch).eq(await epochManager.currentEpoch())
         expect(toRound(event.amount)).eq(toRound(expectedIndexingRewards))
 
         // After state
@@ -853,7 +851,7 @@ describe('Rewards', () => {
         const tx = staking.connect(indexer1).closeAllocation(allocationID1, randomHexBytes())
         await expect(tx)
           .emit(rewardsManager, 'RewardsDenied')
-          .withArgs(indexer1.address, allocationID1, await epochManager.currentEpoch())
+          .withArgs(indexer1.address, allocationID1)
       })
     })
   })

packages/horizon/contracts/interfaces/internal/IHorizonStakingExtension.sol

@@ -158,13 +158,6 @@ interface IHorizonStakingExtension {
      */
     function isAllocation(address allocationID) external view returns (bool);
 
-    /**
-     * @notice Return the total amount of tokens allocated to subgraph.
-     * @param subgraphDeploymentID Deployment ID for the subgraph
-     * @return Total tokens allocated to subgraph
-     */
-    function getSubgraphAllocatedTokens(bytes32 subgraphDeploymentID) external view returns (uint256);
-
     /**
      * @notice Retrun the time in blocks to unstake
      * @return Thawing period in blocks

packages/horizon/contracts/staking/HorizonStakingExtension.sol

@@ -268,6 +268,15 @@ contract HorizonStakingExtension is HorizonStakingBase, IRewardsIssuer, IL2Staki
         return __DEPRECATED_subgraphAllocations[subgraphDeploymentID];
     }
 
+    /**
+     * @notice Return true if allocation is active.
+     * @param allocationID Allocation identifier
+     * @return True if allocation is active
+     */
+    function isActiveAllocation(address allocationID) external view override returns (bool) {
+        return _getAllocationState(allocationID) == AllocationState.Active;
+    }
+
     /**
      * @notice Get the total amount of tokens staked by the indexer.
      * @param indexer Address of the indexer

packages/subgraph-service/contracts/SubgraphService.sol

@@ -33,6 +33,7 @@ contract SubgraphService is
 {
     using PPMMath for uint256;
     using Allocation for mapping(address => Allocation.State);
+    using Allocation for Allocation.State;
 
     /**
      * @notice Checks that an indexer is registered
@@ -372,6 +373,29 @@ contract SubgraphService is
         );
     }
 
+    /**
+     * @notice Return the total amount of tokens allocated to subgraph.
+     * @dev Implements {IRewardsIssuer.getSubgraphAllocatedTokens}
+     * @dev To be used by the {RewardsManager}.
+     * @param subgraphDeploymentId Deployment Id for the subgraph
+     * @return Total tokens allocated to subgraph
+     */
+    function getSubgraphAllocatedTokens(bytes32 subgraphDeploymentId) external view override returns (uint256) {
+        return subgraphAllocatedTokens[subgraphDeploymentId];
+    }
+
+    /**
+     * @notice Check if an allocation is open
+     * @dev Implements {IRewardsIssuer.isAllocationActive}
+     * @dev To be used by the {RewardsManager}.
+     *
+     * @param allocationId The allocation Id
+     * @return Wether or not the allocation is active
+     */
+    function isActiveAllocation(address allocationId) external view override returns (bool) {
+        return allocations[allocationId].isOpen();
+    }
+
     /**
      * @notice See {ISubgraphService.getLegacyAllocation}
      */

packages/subgraph-service/test/mocks/MockRewardsManager.sol

@@ -12,7 +12,7 @@ contract MockRewardsManager is IRewardsManager {
 
     function setMinimumSubgraphSignal(uint256) external {}
 
-    function setRewardsIssuer(address, bool) external {}
+    function setSubgraphService(address) external {}
 
     // -- Denylist --