feat(pulse-scheduler): optimize gas when querying active subscriptions (#2604)

tejasbadadare · web-flow · commit ff06b5e65d1e · 2025-04-23T11:00:53.000-07:00
* feat: optimize gas when querying active subscriptions

* test: expand testActivateDeactivateSubscription

* doc: update docs for getActiveSubscriptions
diff --git a/target_chains/ethereum/contracts/contracts/pulse/scheduler/IScheduler.sol b/target_chains/ethereum/contracts/contracts/pulse/scheduler/IScheduler.sol
@@ -48,7 +48,7 @@ interface IScheduler is SchedulerEvents {
 
     /**
      * @notice Updates price feeds for a subscription.
-     * Verifies the updateData using the Pyth contract and validates that all feeds have the same timestamp.
+     * @dev Internally, this verifies the updateData using the Pyth contract and validates update conditions.
      * @param subscriptionId The ID of the subscription
      * @param updateData The price update data from Pyth
      * @param priceIds The IDs of the price feeds to update
@@ -73,7 +73,8 @@ interface IScheduler is SchedulerEvents {
         bytes32[] calldata priceIds
     ) external view returns (PythStructs.Price[] memory prices);
 
-    /** @notice Returns the exponentially-weighted moving average price of a price feed without any sanity checks.
+    /**
+     * @notice Returns the exponentially-weighted moving average price of a price feed without any sanity checks.
      * @dev This function returns the same price as `getEmaPrice` in the case where the price is available.
      * However, if the price is not recent this function returns the latest available price.
      *
@@ -114,10 +115,12 @@ interface IScheduler is SchedulerEvents {
     ) external view returns (uint256 minimumBalanceInWei);
 
     /**
-     * @notice Gets all active subscriptions with their parameters
-     * @dev This function has no access control to allow keepers to discover active subscriptions
-     * @param startIndex The starting index for pagination
-     * @param maxResults The maximum number of results to return
+     * @notice Gets all active subscriptions with their parameters, paginated.
+     * @dev This function has no access control to allow keepers to discover active subscriptions.
+     * @dev Note that the order of subscription IDs returned may not be sequential and can change
+     *      when subscriptions are deactivated or reactivated.
+     * @param startIndex The starting index within the list of active subscriptions (NOT the subscription ID).
+     * @param maxResults The maximum number of results to return starting from startIndex.
      * @return subscriptionIds Array of active subscription IDs
      * @return subscriptionParams Array of subscription parameters for each active subscription
      * @return totalCount Total number of active subscriptions
diff --git a/target_chains/ethereum/contracts/contracts/pulse/scheduler/Scheduler.sol b/target_chains/ethereum/contracts/contracts/pulse/scheduler/Scheduler.sol
@@ -56,6 +56,8 @@ abstract contract Scheduler is IScheduler, SchedulerState {
         // Map subscription ID to manager
         _state.subscriptionManager[subscriptionId] = msg.sender;
 
+        _addToActiveSubscriptions(subscriptionId);
+
         emit SubscriptionCreated(subscriptionId, msg.sender);
         return subscriptionId;
     }
@@ -102,10 +104,12 @@ abstract contract Scheduler is IScheduler, SchedulerState {
             }
 
             currentParams.isActive = true;
+            _addToActiveSubscriptions(subscriptionId);
             emit SubscriptionActivated(subscriptionId);
         } else if (wasActive && !willBeActive) {
             // Deactivating a subscription
             currentParams.isActive = false;
+            _removeFromActiveSubscriptions(subscriptionId);
             emit SubscriptionDeactivated(subscriptionId);
         }
 
@@ -583,14 +587,7 @@ abstract contract Scheduler is IScheduler, SchedulerState {
             uint256 totalCount
         )
     {
-        // Count active subscriptions first to determine total count
-        // TODO: Optimize this. store numActiveSubscriptions or something.
-        totalCount = 0;
-        for (uint256 i = 1; i < _state.subscriptionNumber; i++) {
-            if (_state.subscriptionParams[i].isActive) {
-                totalCount++;
-            }
-        }
+        totalCount = _state.activeSubscriptionIds.length;
 
         // If startIndex is beyond the total count, return empty arrays
         if (startIndex >= totalCount) {
@@ -607,25 +604,13 @@ abstract contract Scheduler is IScheduler, SchedulerState {
         subscriptionIds = new uint256[](resultCount);
         subscriptionParams = new SubscriptionParams[](resultCount);
 
-        // Find and populate the requested page of active subscriptions
-        uint256 activeIndex = 0;
-        uint256 resultIndex = 0;
-
-        for (
-            uint256 i = 1;
-            i < _state.subscriptionNumber && resultIndex < resultCount;
-            i++
-        ) {
-            if (_state.subscriptionParams[i].isActive) {
-                if (activeIndex >= startIndex) {
-                    subscriptionIds[resultIndex] = i;
-                    subscriptionParams[resultIndex] = _state.subscriptionParams[
-                        i
-                    ];
-                    resultIndex++;
-                }
-                activeIndex++;
-            }
+        // Populate the arrays with the requested page of active subscriptions
+        for (uint256 i = 0; i < resultCount; i++) {
+            uint256 subscriptionId = _state.activeSubscriptionIds[
+                startIndex + i
+            ];
+            subscriptionIds[i] = subscriptionId;
+            subscriptionParams[i] = _state.subscriptionParams[subscriptionId];
         }
 
         return (subscriptionIds, subscriptionParams, totalCount);
@@ -682,4 +667,47 @@ abstract contract Scheduler is IScheduler, SchedulerState {
         }
         _;
     }
+
+    /**
+     * @notice Adds a subscription to the active subscriptions list.
+     * @param subscriptionId The ID of the subscription to add.
+     */
+    function _addToActiveSubscriptions(uint256 subscriptionId) internal {
+        // Only add if not already in the list
+        if (_state.activeSubscriptionIndex[subscriptionId] == 0) {
+            _state.activeSubscriptionIds.push(subscriptionId);
+
+            // Store the index as 1-based, 0 means not in the list
+            _state.activeSubscriptionIndex[subscriptionId] = _state
+                .activeSubscriptionIds
+                .length;
+        }
+    }
+
+    /**
+     * @notice Removes a subscription from the active subscriptions list.
+     * @param subscriptionId The ID of the subscription to remove.
+     */
+    function _removeFromActiveSubscriptions(uint256 subscriptionId) internal {
+        uint256 index = _state.activeSubscriptionIndex[subscriptionId];
+
+        // Only remove if it's in the list
+        if (index > 0) {
+            // Adjust index to be 0-based instead of 1-based
+            index = index - 1;
+
+            // If it's not the last element, move the last element to its position
+            if (index < _state.activeSubscriptionIds.length - 1) {
+                uint256 lastId = _state.activeSubscriptionIds[
+                    _state.activeSubscriptionIds.length - 1
+                ];
+                _state.activeSubscriptionIds[index] = lastId;
+                _state.activeSubscriptionIndex[lastId] = index + 1; // 1-based index
+            }
+
+            // Remove the last element
+            _state.activeSubscriptionIds.pop();
+            _state.activeSubscriptionIndex[subscriptionId] = 0;
+        }
+    }
 }
diff --git a/target_chains/ethereum/contracts/contracts/pulse/scheduler/SchedulerState.sol b/target_chains/ethereum/contracts/contracts/pulse/scheduler/SchedulerState.sol
@@ -31,6 +31,12 @@ contract SchedulerState {
         mapping(uint256 => mapping(bytes32 => PythStructs.PriceFeed)) priceUpdates;
         /// Sub ID -> manager address
         mapping(uint256 => address) subscriptionManager;
+        /// Array of active subscription IDs.
+        /// Gas optimization to avoid scanning through all subscriptions when querying for all active ones.
+        uint256[] activeSubscriptionIds;
+        /// Sub ID -> index in activeSubscriptionIds array + 1 (0 means not in array).
+        /// This lets us avoid a linear scan of `activeSubscriptionIds` when deactivating a subscription.
+        mapping(uint256 => uint256) activeSubscriptionIndex;
     }
     State internal _state;
 
diff --git a/target_chains/ethereum/contracts/forge-test/PulseScheduler.t.sol b/target_chains/ethereum/contracts/forge-test/PulseScheduler.t.sol
@@ -331,51 +331,172 @@ contract SchedulerTest is Test, SchedulerEvents, PulseSchedulerTestUtils {
     }
 
     function testActivateDeactivateSubscription() public {
-        uint256 subscriptionId = addTestSubscription(
-            scheduler,
-            address(reader)
-        );
-
-        // Get current params
-        (SchedulerState.SubscriptionParams memory params, ) = scheduler
-            .getSubscription(subscriptionId);
-
-        // Deactivate subscription using updateSubscription
-        params.isActive = false;
+        // Add multiple subscriptions
+        uint256 subId1 = addTestSubscription(scheduler, address(reader)); // ID 1
+        uint256 subId2 = addTestSubscription(scheduler, address(reader)); // ID 2
+        uint256 subId3 = addTestSubscription(scheduler, address(reader)); // ID 3
 
+        // --- Verify initial state ---
+        (uint256[] memory activeIds, , uint256 totalCount) = scheduler
+            .getActiveSubscriptions(0, 10);
+        assertEq(totalCount, 3, "Initial: Total count should be 3");
+        assertEq(activeIds.length, 3, "Initial: Active IDs length should be 3");
+        assertEq(activeIds[0], subId1, "Initial: ID 1 should be active");
+        assertEq(activeIds[1], subId2, "Initial: ID 2 should be active");
+        assertEq(activeIds[2], subId3, "Initial: ID 3 should be active");
+
+        // --- Deactivate the middle subscription (ID 2) ---
+        (SchedulerState.SubscriptionParams memory params2, ) = scheduler
+            .getSubscription(subId2);
+        params2.isActive = false;
         vm.expectEmit();
-        emit SubscriptionDeactivated(subscriptionId);
+        emit SubscriptionDeactivated(subId2);
         vm.expectEmit();
-        emit SubscriptionUpdated(subscriptionId);
+        emit SubscriptionUpdated(subId2);
+        scheduler.updateSubscription(subId2, params2);
 
-        scheduler.updateSubscription(subscriptionId, params);
+        // Verify state after deactivating ID 2
+        (activeIds, , totalCount) = scheduler.getActiveSubscriptions(0, 10);
+        assertEq(totalCount, 2, "After Deact 2: Total count should be 2");
+        assertEq(
+            activeIds.length,
+            2,
+            "After Deact 2: Active IDs length should be 2"
+        );
+        assertEq(activeIds[0], subId1, "After Deact 2: ID 1 should be active");
+        assertEq(
+            activeIds[1],
+            subId3,
+            "After Deact 2: ID 3 should be active (moved)"
+        ); // ID 3 takes the place of ID 2
+
+        // --- Deactivate the last subscription (ID 3, now at index 1) ---
+        (SchedulerState.SubscriptionParams memory params3, ) = scheduler
+            .getSubscription(subId3);
+        params3.isActive = false;
+        vm.expectEmit();
+        emit SubscriptionDeactivated(subId3);
+        vm.expectEmit();
+        emit SubscriptionUpdated(subId3);
+        scheduler.updateSubscription(subId3, params3);
 
-        // Verify subscription was deactivated
-        (
-            SchedulerState.SubscriptionParams memory storedParams,
-            SchedulerState.SubscriptionStatus memory status
-        ) = scheduler.getSubscription(subscriptionId);
+        // Verify state after deactivating ID 3
+        (activeIds, , totalCount) = scheduler.getActiveSubscriptions(0, 10);
+        assertEq(totalCount, 1, "After Deact 3: Total count should be 1");
+        assertEq(
+            activeIds.length,
+            1,
+            "After Deact 3: Active IDs length should be 1"
+        );
+        assertEq(
+            activeIds[0],
+            subId1,
+            "After Deact 3: Only ID 1 should be active"
+        );
 
-        assertFalse(storedParams.isActive, "Subscription should be inactive");
+        // --- Reactivate the middle subscription (ID 2) ---
+        params2.isActive = true; // Use the params struct from earlier
+        vm.expectEmit();
+        emit SubscriptionActivated(subId2);
+        vm.expectEmit();
+        emit SubscriptionUpdated(subId2);
+        scheduler.updateSubscription(subId2, params2);
 
-        // Reactivate subscription using updateSubscription
-        params.isActive = true;
+        // Verify state after reactivating ID 2
+        (activeIds, , totalCount) = scheduler.getActiveSubscriptions(0, 10);
+        assertEq(totalCount, 2, "After React 2: Total count should be 2");
+        assertEq(
+            activeIds.length,
+            2,
+            "After React 2: Active IDs length should be 2"
+        );
+        assertEq(activeIds[0], subId1, "After React 2: ID 1 should be active");
+        assertEq(activeIds[1], subId2, "After React 2: ID 2 should be active"); // ID 2 is added back to the end
 
+        // --- Reactivate the last subscription (ID 3) ---
+        params3.isActive = true; // Use the params struct from earlier
         vm.expectEmit();
-        emit SubscriptionActivated(subscriptionId);
+        emit SubscriptionActivated(subId3);
         vm.expectEmit();
-        emit SubscriptionUpdated(subscriptionId);
+        emit SubscriptionUpdated(subId3);
+        scheduler.updateSubscription(subId3, params3);
+
+        // Verify final state (all active)
+        (activeIds, , totalCount) = scheduler.getActiveSubscriptions(0, 10);
+        assertEq(totalCount, 3, "Final: Total count should be 3");
+        assertEq(activeIds.length, 3, "Final: Active IDs length should be 3");
+        assertEq(activeIds[0], subId1, "Final: ID 1 should be active");
+        assertEq(activeIds[1], subId2, "Final: ID 2 should be active");
+        assertEq(activeIds[2], subId3, "Final: ID 3 should be active"); // ID 3 is added back to the end
+
+        // --- Deactivate all remaining subscriptions ---
+        // Deactivate ID 1 (first element)
+        (SchedulerState.SubscriptionParams memory params1, ) = scheduler
+            .getSubscription(subId1);
+        params1.isActive = false;
+        vm.expectEmit();
+        emit SubscriptionDeactivated(subId1);
+        vm.expectEmit();
+        emit SubscriptionUpdated(subId1);
+        scheduler.updateSubscription(subId1, params1);
 
-        scheduler.updateSubscription(subscriptionId, params);
+        // Verify state after deactivating ID 1
+        (activeIds, , totalCount) = scheduler.getActiveSubscriptions(0, 10);
+        assertEq(totalCount, 2, "After Deact 1: Total count should be 2");
+        assertEq(
+            activeIds.length,
+            2,
+            "After Deact 1: Active IDs length should be 2"
+        );
+        assertEq(
+            activeIds[0],
+            subId3,
+            "After Deact 1: ID 3 should be at index 0"
+        ); // ID 3 moved to front
+        assertEq(
+            activeIds[1],
+            subId2,
+            "After Deact 1: ID 2 should be at index 1"
+        );
 
-        // Verify subscription was reactivated
-        (storedParams, status) = scheduler.getSubscription(subscriptionId);
+        // Deactivate ID 2 (now last element)
+        params2.isActive = false; // Use existing params struct
+        vm.expectEmit();
+        emit SubscriptionDeactivated(subId2);
+        vm.expectEmit();
+        emit SubscriptionUpdated(subId2);
+        scheduler.updateSubscription(subId2, params2);
 
-        assertTrue(storedParams.isActive, "Subscription should be active");
-        assertTrue(
-            storedParams.isActive,
-            "Subscription params should show active"
+        // Verify state after deactivating ID 2
+        (activeIds, , totalCount) = scheduler.getActiveSubscriptions(0, 10);
+        assertEq(
+            totalCount,
+            1,
+            "After Deact 2 (again): Total count should be 1"
+        );
+        assertEq(
+            activeIds.length,
+            1,
+            "After Deact 2 (again): Active IDs length should be 1"
+        );
+        assertEq(
+            activeIds[0],
+            subId3,
+            "After Deact 2 (again): Only ID 3 should be active"
         );
+
+        // Deactivate ID 3 (last remaining element)
+        params3.isActive = false; // Use existing params struct
+        vm.expectEmit();
+        emit SubscriptionDeactivated(subId3);
+        vm.expectEmit();
+        emit SubscriptionUpdated(subId3);
+        scheduler.updateSubscription(subId3, params3);
+
+        // Verify final empty state
+        (activeIds, , totalCount) = scheduler.getActiveSubscriptions(0, 10);
+        assertEq(totalCount, 0, "Empty: Total count should be 0");
+        assertEq(activeIds.length, 0, "Empty: Active IDs length should be 0");
     }
 
     function testAddFunds() public {
