From c8d40f284958967ff01cfc2b39ac0c87dc2a4a7b Mon Sep 17 00:00:00 2001 From: Tejas Badadare Date: Tue, 22 Apr 2025 19:57:37 -0700 Subject: [PATCH 1/3] feat: optimize gas when querying active subscriptions --- .../contracts/pulse/scheduler/Scheduler.sol | 80 ++++++++++++------- .../pulse/scheduler/SchedulerState.sol | 6 ++ 2 files changed, 59 insertions(+), 27 deletions(-) diff --git a/target_chains/ethereum/contracts/contracts/pulse/scheduler/Scheduler.sol b/target_chains/ethereum/contracts/contracts/pulse/scheduler/Scheduler.sol index 80c3853f5b..c41195eb00 100644 --- 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; } @@ -100,10 +102,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); } @@ -691,14 +695,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) { @@ -715,25 +712,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); @@ -790,4 +775,45 @@ 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); + _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 index 1e1093e775..0b3a151aac 100644 --- a/target_chains/ethereum/contracts/contracts/pulse/scheduler/SchedulerState.sol +++ b/target_chains/ethereum/contracts/contracts/pulse/scheduler/SchedulerState.sol @@ -35,6 +35,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; From 15e39011ec9a9324a80479aae97abb593ee846be Mon Sep 17 00:00:00 2001 From: Tejas Badadare Date: Wed, 23 Apr 2025 10:48:40 -0700 Subject: [PATCH 2/3] test: expand testActivateDeactivateSubscription --- .../contracts/pulse/scheduler/Scheduler.sol | 2 + .../contracts/forge-test/PulseScheduler.t.sol | 183 +++++++++++++++--- 2 files changed, 154 insertions(+), 31 deletions(-) diff --git a/target_chains/ethereum/contracts/contracts/pulse/scheduler/Scheduler.sol b/target_chains/ethereum/contracts/contracts/pulse/scheduler/Scheduler.sol index c41195eb00..8c6f95181c 100644 --- a/target_chains/ethereum/contracts/contracts/pulse/scheduler/Scheduler.sol +++ b/target_chains/ethereum/contracts/contracts/pulse/scheduler/Scheduler.sol @@ -784,6 +784,8 @@ abstract contract Scheduler is IScheduler, SchedulerState { // 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; diff --git a/target_chains/ethereum/contracts/forge-test/PulseScheduler.t.sol b/target_chains/ethereum/contracts/forge-test/PulseScheduler.t.sol index 3eb3888b5b..5c8462697b 100644 --- a/target_chains/ethereum/contracts/forge-test/PulseScheduler.t.sol +++ b/target_chains/ethereum/contracts/forge-test/PulseScheduler.t.sol @@ -348,51 +348,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 { From 29b1a9579774e3f9fe8d5666f6dbb1047752151e Mon Sep 17 00:00:00 2001 From: Tejas Badadare Date: Wed, 23 Apr 2025 10:59:21 -0700 Subject: [PATCH 3/3] doc: update docs for getActiveSubscriptions --- .../contracts/pulse/scheduler/IScheduler.sol | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/target_chains/ethereum/contracts/contracts/pulse/scheduler/IScheduler.sol b/target_chains/ethereum/contracts/contracts/pulse/scheduler/IScheduler.sol index 3dddbbd5f3..b3a3b7b42f 100644 --- 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