Upkeep Balance Monitor (#11180)

* write UpkeepBalanceMonWithBuffer * Update AutomationRegistryInterface2_0.sol Support getMinBalanceForUpkeep * reformat * remove modifier * add typeAndVersion check in constructor * restore AutomationRegistryInterface2_0.sol function * update import paths * update comments * cleanup * get rid of redundant REGISTRY variable * cleanup * rename * cleanup * remove target, min wait period, switch to min/target percentages * refactor getUnderfundedUpkeeps() to return top up amounts * refactor topUp() function * switch to max batch size * add max top up amount * add maxTopUpAmount * whitelist performUpkeep to forwarder * cleanup * rename * bring topUp() back * write initial test suite * fix solhint errors * update test for getUnderfundedUpkeeps() * add tests for owner only functions and events * add topUp and performUpkeep tests * rearrange functions on contract * add support for multiple registries * cleanup * change topUpAmounts to uint96[] * move pausable to topUp(); add length check to topUp() --------- Co-authored-by: De Clercq Wentzel <[email protected]>
smartcontractkit · Dec 14, 2023 · 96fa461 · 96fa461
1 parent d424b11
commit 96fa461
Show file tree

Hide file tree

Showing 2 changed files with 657 additions and 0 deletions.
diff --git a/contracts/src/v0.8/automation/upkeeps/UpkeepBalanceMonitor.sol b/contracts/src/v0.8/automation/upkeeps/UpkeepBalanceMonitor.sol
@@ -0,0 +1,258 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity 0.8.19;
+
+import {ConfirmedOwner} from "../../shared/access/ConfirmedOwner.sol";
+import {IAutomationRegistryConsumer} from "../interfaces/IAutomationRegistryConsumer.sol";
+import {LinkTokenInterface} from "../../shared/interfaces/LinkTokenInterface.sol";
+import {Pausable} from "@openzeppelin/contracts/security/Pausable.sol";
+import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
+
+/// @title The UpkeepBalanceMonitor contract
+/// @notice A keeper-compatible contract that monitors and funds Chainlink Automation upkeeps.
+contract UpkeepBalanceMonitor is ConfirmedOwner, Pausable {
+  using EnumerableSet for EnumerableSet.AddressSet;
+
+  event ConfigSet(Config config);
+  event ForwarderSet(address forwarderAddress);
+  event FundsWithdrawn(uint256 amountWithdrawn, address payee);
+  event TopUpFailed(uint256 indexed upkeepId);
+  event TopUpSucceeded(uint256 indexed upkeepId, uint96 amount);
+  event WatchListSet(address registryAddress);
+
+  error InvalidConfig();
+  error InvalidTopUpData();
+  error OnlyForwarderOrOwner();
+
+  /// @member maxBatchSize is the maximum number of upkeeps to fund in a single transaction
+  /// @member minPercentage is the percentage of the upkeep's minBalance at which top-up occurs
+  /// @member targetPercentage is the percentage of the upkeep's minBalance to top-up to
+  /// @member maxTopUpAmount is the maximum amount of LINK to top-up an upkeep with
+  struct Config {
+    uint8 maxBatchSize;
+    uint24 minPercentage;
+    uint24 targetPercentage;
+    uint96 maxTopUpAmount;
+  }
+
+  // ================================================================
+  // |                           STORAGE                            |
+  // ================================================================
+
+  LinkTokenInterface private immutable LINK_TOKEN;
+
+  mapping(address => uint256[]) s_registryWatchLists;
+  EnumerableSet.AddressSet s_registries;
+  Config private s_config;
+  address private s_forwarderAddress;
+
+  // ================================================================
+  // |                         CONSTRUCTOR                          |
+  // ================================================================
+
+  /// @param linkToken the Link token address
+  /// @param config the initial config for the contract
+  constructor(LinkTokenInterface linkToken, Config memory config) ConfirmedOwner(msg.sender) {
+    require(address(linkToken) != address(0));
+    LINK_TOKEN = linkToken;
+    setConfig(config);
+  }
+
+  // ================================================================
+  // |                      CORE FUNCTIONALITY                      |
+  // ================================================================
+
+  /// @notice Gets a list of upkeeps that are underfunded
+  /// @return needsFunding list of underfunded upkeepIDs
+  /// @return registryAddresses list of registries that the upkeepIDs belong to
+  /// @return topUpAmounts amount to top up each upkeep
+  function getUnderfundedUpkeeps() public view returns (uint256[] memory, address[] memory, uint96[] memory) {
+    Config memory config = s_config;
+    uint256[] memory needsFunding = new uint256[](config.maxBatchSize);
+    address[] memory registryAddresses = new address[](config.maxBatchSize);
+    uint96[] memory topUpAmounts = new uint96[](config.maxBatchSize);
+    uint256 availableFunds = LINK_TOKEN.balanceOf(address(this));
+    uint256 count;
+    for (uint256 i = 0; i < s_registries.length(); i++) {
+      IAutomationRegistryConsumer registry = IAutomationRegistryConsumer(s_registries.at(i));
+      for (uint256 j = 0; j < s_registryWatchLists[address(registry)].length; j++) {
+        uint256 upkeepID = s_registryWatchLists[address(registry)][j];
+        uint96 upkeepBalance = registry.getBalance(upkeepID);
+        uint96 minBalance = registry.getMinBalance(upkeepID);
+        uint96 topUpThreshold = (minBalance * config.minPercentage) / 100;
+        uint96 topUpAmount = ((minBalance * config.targetPercentage) / 100) - upkeepBalance;
+        if (topUpAmount > config.maxTopUpAmount) {
+          topUpAmount = config.maxTopUpAmount;
+        }
+        if (upkeepBalance <= topUpThreshold && availableFunds >= topUpAmount) {
+          needsFunding[count] = upkeepID;
+          topUpAmounts[count] = topUpAmount;
+          registryAddresses[count] = address(registry);
+          count++;
+          availableFunds -= topUpAmount;
+        }
+        if (count == config.maxBatchSize) {
+          break;
+        }
+      }
+      if (count == config.maxBatchSize) {
+        break;
+      }
+    }
+    if (count < config.maxBatchSize) {
+      assembly {
+        mstore(needsFunding, count)
+        mstore(registryAddresses, count)
+        mstore(topUpAmounts, count)
+      }
+    }
+    return (needsFunding, registryAddresses, topUpAmounts);
+  }
+
+  /// @notice Called by the keeper/owner to send funds to underfunded upkeeps
+  /// @param upkeepIDs the list of upkeep ids to fund
+  /// @param registryAddresses the list of registries that the upkeepIDs belong to
+  /// @param topUpAmounts the list of amounts to fund each upkeep with
+  /// @dev We explicitly choose not to verify that input upkeepIDs are included in the watchlist. We also
+  /// explicity permit any amount to be sent via topUpAmounts; it does not have to meet the criteria
+  /// specified in getUnderfundedUpkeeps(). Here, we are relying on the security of automation's OCR to
+  /// secure the output of getUnderfundedUpkeeps() as the input to topUp(), and we are treating the owner
+  /// as a privileged user that can perform arbitrary top-ups to any upkeepID.
+  function topUp(
+    uint256[] memory upkeepIDs,
+    address[] memory registryAddresses,
+    uint96[] memory topUpAmounts
+  ) public whenNotPaused {
+    if (msg.sender != address(s_forwarderAddress) && msg.sender != owner()) revert OnlyForwarderOrOwner();
+    if (upkeepIDs.length != registryAddresses.length || upkeepIDs.length != topUpAmounts.length)
+      revert InvalidTopUpData();
+    for (uint256 i = 0; i < upkeepIDs.length; i++) {
+      try LINK_TOKEN.transferAndCall(registryAddresses[i], topUpAmounts[i], abi.encode(upkeepIDs[i])) returns (
+        bool success
+      ) {
+        if (success) {
+          emit TopUpSucceeded(upkeepIDs[i], topUpAmounts[i]);
+          continue;
+        }
+      } catch {}
+      emit TopUpFailed(upkeepIDs[i]);
+    }
+  }
+
+  // ================================================================
+  // |                    AUTOMATION COMPATIBLE                     |
+  // ================================================================
+
+  /// @notice Gets list of upkeeps ids that are underfunded and returns a keeper-compatible payload.
+  /// @return upkeepNeeded signals if upkeep is needed, performData is an abi encoded list of subscription ids that need funds
+  function checkUpkeep(bytes calldata) external view returns (bool upkeepNeeded, bytes memory performData) {
+    (
+      uint256[] memory needsFunding,
+      address[] memory registryAddresses,
+      uint96[] memory topUpAmounts
+    ) = getUnderfundedUpkeeps();
+    upkeepNeeded = needsFunding.length > 0;
+    if (upkeepNeeded) {
+      performData = abi.encode(needsFunding, registryAddresses, topUpAmounts);
+    }
+    return (upkeepNeeded, performData);
+  }
+
+  /// @notice Called by the keeper to send funds to underfunded addresses.
+  /// @param performData the abi encoded list of addresses to fund
+  function performUpkeep(bytes calldata performData) external {
+    (uint256[] memory upkeepIDs, address[] memory registryAddresses, uint96[] memory topUpAmounts) = abi.decode(
+      performData,
+      (uint256[], address[], uint96[])
+    );
+    topUp(upkeepIDs, registryAddresses, topUpAmounts);
+  }
+
+  // ================================================================
+  // |                            ADMIN                             |
+  // ================================================================
+
+  /// @notice Withdraws the contract balance in LINK.
+  /// @param amount the amount of LINK (in juels) to withdraw
+  /// @param payee the address to pay
+  function withdraw(uint256 amount, address payee) external onlyOwner {
+    require(payee != address(0));
+    LINK_TOKEN.transfer(payee, amount);
+    emit FundsWithdrawn(amount, payee);
+  }
+
+  /// @notice Pause the contract, which prevents executing performUpkeep.
+  function pause() external onlyOwner {
+    _pause();
+  }
+
+  /// @notice Unpause the contract.
+  function unpause() external onlyOwner {
+    _unpause();
+  }
+
+  // ================================================================
+  // |                           SETTERS                            |
+  // ================================================================
+
+  /// @notice Sets the list of upkeeps to watch
+  /// @param registryAddress the registry that this watchlist applies to
+  /// @param watchlist the list of UpkeepIDs to watch
+  function setWatchList(address registryAddress, uint256[] calldata watchlist) external onlyOwner {
+    if (watchlist.length == 0) {
+      s_registries.remove(registryAddress);
+      delete s_registryWatchLists[registryAddress];
+    } else {
+      s_registries.add(registryAddress);
+      s_registryWatchLists[registryAddress] = watchlist;
+    }
+    emit WatchListSet(registryAddress);
+  }
+
+  /// @notice Sets the contract config
+  /// @param config the new config
+  function setConfig(Config memory config) public onlyOwner {
+    if (
+      config.maxBatchSize == 0 ||
+      config.minPercentage < 100 ||
+      config.targetPercentage <= config.minPercentage ||
+      config.maxTopUpAmount == 0
+    ) {
+      revert InvalidConfig();
+    }
+    s_config = config;
+    emit ConfigSet(config);
+  }
+
+  /// @notice Sets the upkeep's forwarder contract
+  /// @param forwarderAddress the new forwarder
+  /// @dev this should only need to be called once, after registering the contract with the registry
+  function setForwarder(address forwarderAddress) external onlyOwner {
+    s_forwarderAddress = forwarderAddress;
+    emit ForwarderSet(forwarderAddress);
+  }
+
+  // ================================================================
+  // |                           GETTERS                            |
+  // ================================================================
+
+  /// @notice Gets the list of upkeeps ids being monitored
+  function getWatchList() external view returns (address[] memory, uint256[][] memory) {
+    address[] memory registryAddresses = s_registries.values();
+    uint256[][] memory upkeepIDs = new uint256[][](registryAddresses.length);
+    for (uint256 i = 0; i < registryAddresses.length; i++) {
+      upkeepIDs[i] = s_registryWatchLists[registryAddresses[i]];
+    }
+    return (registryAddresses, upkeepIDs);
+  }
+
+  /// @notice Gets the contract config
+  function getConfig() external view returns (Config memory) {
+    return s_config;
+  }
+
+  /// @notice Gets the upkeep's forwarder contract
+  function getForwarder() external view returns (address) {
+    return s_forwarderAddress;
+  }
+}