[fortuna] Withdraw fees (#1610)

* withdraw fees command

* hmmmm

* rg

* withdraw fees

* it builds

* add fee manager to config

* withdrawal cli fixes

* cleanup

* pr comments

* cargo bump

* log

apps/fortuna/Cargo.lock

@@ -1502,7 +1502,7 @@ dependencies = [
 
 [[package]]
 name = "fortuna"
-version = "6.2.3"
+version = "6.3.0"
 dependencies = [
  "anyhow",
  "axum",

apps/fortuna/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name    = "fortuna"
-version = "6.2.3"
+version = "6.3.0"
 edition = "2021"
 
 [dependencies]

apps/fortuna/config.sample.yaml

@@ -6,9 +6,12 @@ chains:
     # Keeper configuration for the chain
     reveal_delay_blocks: 0
     gas_limit: 500000
+    min_keeper_balance: 100000000000000000
 
     # Provider configuration
+    # How much to charge in fees
     fee: 1500000000000000
+
     # Historical commitments -- delete this block for local development purposes
     commitments:
       # prettier-ignore
@@ -34,6 +37,10 @@ provider:
     value: abcd
     # For production, you can store the private key in a file.
     # file: secret.txt
+
+  # Set this to the address of your keeper wallet if you would like the keeper wallet to
+  # be able to withdraw fees from the contract.
+  fee_manager: 0xADDRESS
 keeper:
   # An ethereum wallet address and private key for running the keeper service.
   # This does not have to be the same key as the provider's key above.

apps/fortuna/src/chain/ethereum.rs

@@ -104,7 +104,6 @@ impl<M> LegacyTxMiddleware<M> {
     }
 }
 
-
 #[derive(Error, Debug)]
 pub enum LegacyTxMiddlewareError<M: Middleware> {
     #[error("{0}")]
@@ -167,6 +166,16 @@ impl<M: Middleware> Middleware for LegacyTxMiddleware<M> {
 }
 
 impl<T: JsonRpcClient + 'static + Clone> SignablePythContractInner<T> {
+    /// Get the wallet that signs transactions sent to this contract.
+    pub fn wallet(&self) -> LocalWallet {
+        self.client().inner().inner().inner().signer().clone()
+    }
+
+    /// Get the underlying provider that communicates with the blockchain.
+    pub fn provider(&self) -> Provider<T> {
+        self.client().inner().inner().inner().provider().clone()
+    }
+
     /// Submit a request for a random number to the contract.
     ///
     /// This method is a version of the autogenned `request` method that parses the emitted logs

apps/fortuna/src/command.rs

@@ -5,6 +5,7 @@ mod register_provider;
 mod request_randomness;
 mod run;
 mod setup_provider;
+mod withdraw_fees;
 
 pub use {
     generate::generate,
@@ -14,4 +15,5 @@ pub use {
     request_randomness::request_randomness,
     run::run,
     setup_provider::setup_provider,
+    withdraw_fees::withdraw_fees,
 };

apps/fortuna/src/command/setup_provider.rs

@@ -32,7 +32,10 @@ use {
             LocalWallet,
             Signer,
         },
-        types::Bytes,
+        types::{
+            Address,
+            Bytes,
+        },
     },
     std::sync::Arc,
     tracing::Instrument,
@@ -66,7 +69,6 @@ async fn setup_chain_provider(
         "Please specify a provider private key in the config file."
     ))?;
     let provider_address = private_key.clone().parse::<LocalWallet>()?.address();
-    let provider_fee = chain_config.fee;
     // Initialize a Provider to interface with the EVM contract.
     let contract = Arc::new(SignablePythContract::from_config(&chain_config, &private_key).await?);
 
@@ -130,15 +132,28 @@ async fn setup_chain_provider(
             .await
             .map_err(|e| anyhow!("Chain: {} - Failed to register provider: {}", &chain_id, e))?;
         tracing::info!("Registered");
-    } else {
-        sync_fee(&contract, &provider_info, provider_fee)
-            .in_current_span()
-            .await?;
-        let uri = get_register_uri(&provider_config.uri, &chain_id)?;
-        sync_uri(&contract, &provider_info, uri)
-            .in_current_span()
-            .await?;
     }
+
+
+    let provider_info = contract.get_provider_info(provider_address).call().await?;
+
+    sync_fee(&contract, &provider_info, chain_config.fee)
+        .in_current_span()
+        .await?;
+
+    let uri = get_register_uri(&provider_config.uri, &chain_id)?;
+    sync_uri(&contract, &provider_info, uri)
+        .in_current_span()
+        .await?;
+
+    sync_fee_manager(
+        &contract,
+        &provider_info,
+        provider_config.fee_manager.unwrap_or(Address::zero()),
+    )
+    .in_current_span()
+    .await?;
+
     Ok(())
 }
 
@@ -180,3 +195,17 @@ async fn sync_fee(
     }
     Ok(())
 }
+
+async fn sync_fee_manager(
+    contract: &Arc<SignablePythContract>,
+    provider_info: &ProviderInfo,
+    fee_manager: Address,
+) -> Result<()> {
+    if provider_info.fee_manager != fee_manager {
+        tracing::info!("Updating provider fee manager to {:?}", fee_manager);
+        if let Some(receipt) = contract.set_fee_manager(fee_manager).send().await?.await? {
+            tracing::info!("Updated provider fee manager: {:?}", receipt);
+        }
+    }
+    Ok(())
+}

apps/fortuna/src/command/withdraw_fees.rs

@@ -0,0 +1,100 @@
+use {
+    crate::{
+        chain::ethereum::SignablePythContract,
+        config::{
+            Config,
+            WithdrawFeesOptions,
+        },
+    },
+    anyhow::{
+        anyhow,
+        Result,
+    },
+    ethers::{
+        signers::Signer,
+        types::Address,
+    },
+};
+
+
+pub async fn withdraw_fees(opts: &WithdrawFeesOptions) -> Result<()> {
+    let config = Config::load(&opts.config.config)?;
+
+    let private_key_string = if opts.keeper {
+        config.keeper.private_key.load()?.ok_or(anyhow!("Please specify a keeper private key in the config or omit the --keeper option to use the provider private key"))?
+    } else {
+        config.provider.private_key.load()?.ok_or(anyhow!(
+            "Please specify a provider private key in the config or provide the --keeper option to use the keeper private key instead."
+        ))?
+    };
+
+    match opts.chain_id.clone() {
+        Some(chain_id) => {
+            let chain_config = &config.get_chain_config(&chain_id)?;
+            let contract =
+                SignablePythContract::from_config(&chain_config, &private_key_string).await?;
+
+            withdraw_fees_for_chain(
+                contract,
+                config.provider.address.clone(),
+                opts.keeper,
+                opts.retain_balance_wei,
+            )
+            .await?;
+        }
+        None => {
+            for (chain_id, chain_config) in config.chains.iter() {
+                tracing::info!("Withdrawing fees for chain: {}", chain_id);
+                let contract =
+                    SignablePythContract::from_config(&chain_config, &private_key_string).await?;
+
+                withdraw_fees_for_chain(
+                    contract,
+                    config.provider.address.clone(),
+                    opts.keeper,
+                    opts.retain_balance_wei,
+                )
+                .await?;
+            }
+        }
+    }
+    Ok(())
+}
+
+pub async fn withdraw_fees_for_chain(
+    contract: SignablePythContract,
+    provider_address: Address,
+    is_fee_manager: bool,
+    retained_balance: u128,
+) -> Result<()> {
+    tracing::info!("Fetching fees for provider: {:?}", provider_address);
+    let provider_info = contract.get_provider_info(provider_address).call().await?;
+    let fees = provider_info.accrued_fees_in_wei;
+    tracing::info!("Accrued fees: {} wei", fees);
+
+    let withdrawal_amount_wei = fees.saturating_sub(retained_balance);
+    if withdrawal_amount_wei > 0 {
+        tracing::info!(
+            "Withdrawing {} wei to {}...",
+            withdrawal_amount_wei,
+            contract.wallet().address()
+        );
+
+        let call = match is_fee_manager {
+            true => contract.withdraw_as_fee_manager(provider_address, withdrawal_amount_wei),
+            false => contract.withdraw(withdrawal_amount_wei),
+        };
+        let tx_result = call.send().await?.await?;
+
+        match &tx_result {
+            Some(receipt) => {
+                tracing::info!("Withdrawal transaction hash {:?}", receipt.transaction_hash);
+            }
+            None => {
+                tracing::warn!("No transaction receipt. Unclear what happened to the transaction");
+            }
+        }
+    }
+
+    Ok(())
+}

apps/fortuna/src/config.rs

@@ -32,6 +32,7 @@ pub use {
     request_randomness::RequestRandomnessOptions,
     run::RunOptions,
     setup_provider::SetupProviderOptions,
+    withdraw_fees::WithdrawFeesOptions,
 };
 
 mod generate;
@@ -41,6 +42,7 @@ mod register_provider;
 mod request_randomness;
 mod run;
 mod setup_provider;
+mod withdraw_fees;
 
 const DEFAULT_RPC_ADDR: &str = "127.0.0.1:34000";
 const DEFAULT_HTTP_ADDR: &str = "http://127.0.0.1:34000";
@@ -73,6 +75,9 @@ pub enum Options {
 
     /// Get the status of a pending request for a random number.
     GetRequest(GetRequestOptions),
+
+    /// Withdraw any of the provider's accumulated fees from the contract.
+    WithdrawFees(WithdrawFeesOptions),
 }
 
 #[derive(Args, Clone, Debug)]
@@ -140,6 +145,12 @@ pub struct EthereumConfig {
     /// The gas limit to use for entropy callback transactions.
     pub gas_limit: u64,
 
+    /// Minimum wallet balance for the keeper. If the balance falls below this level, the keeper will
+    /// withdraw fees from the contract to top up. This functionality requires the keeper to be the fee
+    /// manager for the provider.
+    #[serde(default)]
+    pub min_keeper_balance: u128,
+
     /// How much the provider charges for a request on this chain.
     #[serde(default)]
     pub fee: u128,
@@ -186,6 +197,10 @@ pub struct ProviderConfig {
     /// compute per request for less RAM use.
     #[serde(default = "default_chain_sample_interval")]
     pub chain_sample_interval: u64,
+
+    /// The address of the fee manager for the provider. Set this value to the keeper wallet address to
+    /// enable keeper balance top-ups.
+    pub fee_manager: Option<Address>,
 }
 
 fn default_chain_sample_interval() -> u64 {

apps/fortuna/src/config/withdraw_fees.rs

@@ -0,0 +1,31 @@
+use {
+    crate::{
+        api::ChainId,
+        config::ConfigOptions,
+    },
+    clap::Args,
+};
+
+#[derive(Args, Clone, Debug)]
+#[command(next_help_heading = "Withdraw Fees Options")]
+#[group(id = "Withdraw Fees")]
+pub struct WithdrawFeesOptions {
+    #[command(flatten)]
+    pub config: ConfigOptions,
+
+    /// Withdraw the fees on this chain, or all chains if not specified.
+    #[arg(long = "chain-id")]
+    pub chain_id: Option<ChainId>,
+
+    /// If provided, run the command using the keeper wallet. By default, the command uses the provider wallet.
+    /// If this option is provided, the keeper wallet must be configured and set as the fee manager for the provider.
+    #[arg(long = "keeper")]
+    #[arg(default_value = "false")]
+    pub keeper: bool,
+
+    /// If specified, only withdraw fees over the given balance from the contract.
+    /// If omitted, all accrued fees are withdrawn.
+    #[arg(long = "retain-balance")]
+    #[arg(default_value = "0")]
+    pub retain_balance_wei: u128,
+}

apps/fortuna/src/keeper.rs

@@ -82,6 +82,8 @@ const BLOCK_BATCH_SIZE: u64 = 100;
 const POLL_INTERVAL: Duration = Duration::from_secs(2);
 /// Track metrics in this interval
 const TRACK_INTERVAL: Duration = Duration::from_secs(10);
+/// Check whether we need to conduct a withdrawal at this interval.
+const WITHDRAW_INTERVAL: Duration = Duration::from_secs(300);
 /// Rety last N blocks
 const RETRY_PREVIOUS_BLOCKS: u64 = 100;
 
@@ -230,7 +232,7 @@ pub async fn run_keeper_threads(
         .await
         .expect("Chain config should be valid"),
     );
-    let keeper_address = contract.client().inner().inner().inner().signer().address();
+    let keeper_address = contract.wallet().address();
 
     let fulfilled_requests_cache = Arc::new(RwLock::new(HashSet::<u64>::new()));
 
@@ -275,6 +277,17 @@ pub async fn run_keeper_threads(
         .in_current_span(),
     );
 
+    // Spawn a thread that watches the keeper wallet balance and submits withdrawal transactions as needed to top-up the balance.
+    spawn(
+        withdraw_fees_wrapper(
+            contract.clone(),
+            chain_state.provider_address.clone(),
+            WITHDRAW_INTERVAL,
+            U256::from(chain_eth_config.min_keeper_balance),
+        )
+        .in_current_span(),
+    );
+
     // Spawn a thread to track the provider info and the balance of the keeper
     spawn(
         async move {
@@ -858,3 +871,72 @@ pub async fn track_provider(
         })
         .set(end_sequence_number as i64);
 }
+
+#[tracing::instrument(name = "withdraw_fees", skip_all, fields())]
+pub async fn withdraw_fees_wrapper(
+    contract: Arc<InstrumentedSignablePythContract>,
+    provider_address: Address,
+    poll_interval: Duration,
+    min_balance: U256,
+) {
+    loop {
+        if let Err(e) = withdraw_fees_if_necessary(contract.clone(), provider_address, min_balance)
+            .in_current_span()
+            .await
+        {
+            tracing::error!("Withdrawing fees. error: {:?}", e);
+        }
+        time::sleep(poll_interval).await;
+    }
+}
+
+/// Withdraws accumulated fees in the contract as needed to maintain the balance of the keeper wallet.
+pub async fn withdraw_fees_if_necessary(
+    contract: Arc<InstrumentedSignablePythContract>,
+    provider_address: Address,
+    min_balance: U256,
+) -> Result<()> {
+    let provider = contract.provider();
+    let wallet = contract.wallet();
+
+    let keeper_balance = provider
+        .get_balance(wallet.address(), None)
+        .await
+        .map_err(|e| anyhow!("Error while getting balance. error: {:?}", e))?;
+
+    let provider_info = contract
+        .get_provider_info(provider_address)
+        .call()
+        .await
+        .map_err(|e| anyhow!("Error while getting provider info. error: {:?}", e))?;
+
+    if provider_info.fee_manager != wallet.address() {
+        return Err(anyhow!("Fee manager for provider {:?} is not the keeper wallet. Fee manager: {:?} Keeper: {:?}", provider, provider_info.fee_manager, wallet.address()));
+    }
+
+    let fees = provider_info.accrued_fees_in_wei;
+
+    if keeper_balance < min_balance && U256::from(fees) > min_balance {
+        tracing::info!("Claiming accrued fees...");
+        let contract_call = contract.withdraw_as_fee_manager(provider_address, fees);
+        let pending_tx = contract_call
+            .send()
+            .await
+            .map_err(|e| anyhow!("Error submitting the withdrawal transaction: {:?}", e))?;
+
+        let tx_result = pending_tx
+            .await
+            .map_err(|e| anyhow!("Error waiting for withdrawal transaction receipt: {:?}", e))?
+            .ok_or_else(|| anyhow!("Can't verify the withdrawal, probably dropped from mempool"))?;
+
+        tracing::info!(
+            transaction_hash = &tx_result.transaction_hash.to_string(),
+            "Withdrew fees to keeper address. Receipt: {:?}",
+            tx_result,
+        );
+    } else if keeper_balance < min_balance {
+        tracing::warn!("Keeper balance {:?} is too low (< {:?}) but provider fees are not sufficient to top-up.", keeper_balance, min_balance)
+    }
+
+    Ok(())
+}

apps/fortuna/src/main.rs

@@ -43,5 +43,6 @@ async fn main() -> Result<()> {
         config::Options::SetupProvider(opts) => command::setup_provider(&opts).await,
         config::Options::RequestRandomness(opts) => command::request_randomness(&opts).await,
         config::Options::Inspect(opts) => command::inspect(&opts).await,
+        config::Options::WithdrawFees(opts) => command::withdraw_fees(&opts).await,
     }
 }