Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
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
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,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;
Copy link
Contributor

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

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;
}
}
}
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
Loading