Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 23, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c8d40f2
feat: optimize gas when querying active subscriptions
tejasbadadare Apr 23, 2025
15e3901
test: expand testActivateDeactivateSubscription
tejasbadadare Apr 23, 2025
29b1a95
doc: update docs for getActiveSubscriptions
tejasbadadare Apr 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions target_chains/ethereum/contracts/contracts/pulse/scheduler/IScheduler.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
*
Expand Down Expand Up @@ -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
Expand Down
82 changes: 55 additions & 27 deletions target_chains/ethereum/contracts/contracts/pulse/scheduler/Scheduler.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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) {
Expand All @@ -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);
Expand Down Expand Up @@ -790,4 +775,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;
Copy link
Contributor

@jayantk jayantk Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe leave a comment here noting that this array is 1-indexed. I thought this was a bug until i looked further

that said, is there a test that adds & removes a bunch of subscriptions to validate that this logic works?

Copy link
Contributor Author

@tejasbadadare tejasbadadare Apr 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, the 1-based indexing is documented in other places, but it's definitely not intuitive so i'll mention it every place the map is directly used

There are tests that cover adding and removing subscriptions, and the gas benchmark adds and removes hundreds. I'll augment the existing unit tests to stress test the implementation some more.

}
}

/**
* @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;
}
}
}
6 changes: 6 additions & 0 deletions target_chains/ethereum/contracts/contracts/pulse/scheduler/SchedulerState.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
183 changes: 152 additions & 31 deletions target_chains/ethereum/contracts/forge-test/PulseScheduler.t.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading