Apply review comments.

Author: afck

linera-faucet/server/src/lib.rs

@@ -50,11 +50,13 @@ pub struct QueryRoot<C> {
 
 /// The root GraphQL mutation type.
 pub struct MutationRoot<C> {
+    /// The chain from which new faucet chains are opened.
     main_chain_id: ChainId,
-    tmp_chain_id: Arc<Mutex<Option<ChainId>>>,
+    /// The chain that is currently used to open requested chains for clients.
+    faucet_chain_id: Arc<Mutex<Option<ChainId>>>,
     context: Arc<Mutex<C>>,
     amount: Amount,
-    max_claims_per_chain: u32,
+    faucet_init_balance: Amount,
     end_timestamp: Timestamp,
     start_timestamp: Timestamp,
     start_balance: Amount,
@@ -123,54 +125,28 @@ where
     C: ClientContext,
 {
     async fn do_claim(&self, owner: AccountOwner) -> Result<ClaimOutcome, Error> {
-        let main_client = self
-            .context
-            .lock()
-            .await
-            .make_chain_client(self.main_chain_id)?;
-        let maybe_tmp_chain_id = *self.tmp_chain_id.lock().await;
-        let main_balance = main_client.local_balance().await?;
-        let tmp_chain_id = match maybe_tmp_chain_id {
-            Some(tmp_chain_id) => tmp_chain_id,
-            None => {
-                let key_pair = main_client.key_pair().await?;
-                let balance = self
-                    .amount
-                    .try_add(MAX_FEE)?
-                    .try_mul(u128::from(self.max_claims_per_chain))?
-                    .try_add(MAX_FEE)? // One more block fee for closing the chain.
-                    .min(main_balance.try_sub(MAX_FEE)?);
-                let ownership = main_client.chain_state_view().await?.ownership().clone();
-                let (message_id, certificate) = main_client
-                    .open_chain(ownership, ApplicationPermissions::default(), balance)
-                    .await?
-                    .try_unwrap()?;
-                let chain_id = ChainId::child(message_id);
-                info!("Switching to a new faucet chain {chain_id:8}");
-                self.context
-                    .lock()
-                    .await
-                    .update_wallet_for_new_chain(
-                        chain_id,
-                        Some(key_pair),
-                        certificate.block().header.timestamp,
-                    )
-                    .await?;
-                *self.tmp_chain_id.lock().await = Some(chain_id);
-                chain_id
-            }
+        let maybe_faucet_chain_id = *self.faucet_chain_id.lock().await;
+        let faucet_chain_id = match maybe_faucet_chain_id {
+            Some(faucet_chain_id) => faucet_chain_id,
+            None => self.open_new_faucet_chain().await?,
+        };
+        let (faucet_client, main_client) = {
+            let guard = self.context.lock().await;
+            (
+                guard.make_chain_client(faucet_chain_id)?,
+                guard.make_chain_client(self.main_chain_id)?,
+            )
         };
-        let tmp_client = self.context.lock().await.make_chain_client(tmp_chain_id)?;
 
         if self.start_timestamp < self.end_timestamp {
-            let local_time = tmp_client.storage_client().clock().current_time();
+            let local_time = faucet_client.storage_client().clock().current_time();
             if local_time < self.end_timestamp {
                 let full_duration = self
                     .end_timestamp
                     .delta_since(self.start_timestamp)
                     .as_micros();
                 let remaining_duration = self.end_timestamp.delta_since(local_time).as_micros();
-                let balance = tmp_client
+                let balance = faucet_client
                     .local_balance()
                     .await?
                     .try_add(main_client.local_balance().await?)?;
@@ -190,27 +166,31 @@ where
         }
 
         let ownership = ChainOwnership::single(owner);
-        let result = tmp_client
+        let (message_id, certificate) = faucet_client
             .open_chain(ownership, ApplicationPermissions::default(), self.amount)
-            .await;
-        self.context.lock().await.update_wallet(&tmp_client).await?;
-        let (message_id, certificate) = result?.try_unwrap()?;
+            .await?
+            .try_unwrap()?;
+        self.context
+            .lock()
+            .await
+            .update_wallet(&faucet_client)
+            .await?;
 
         // Only keep using this chain if there will still be enough balance to close it.
-        if tmp_client.local_balance().await? < self.amount.try_add(MAX_FEE.try_mul(2)?)? {
+        if faucet_client.local_balance().await? < self.amount.try_add(MAX_FEE.try_mul(2)?)? {
             // TODO(#1795): Move the remaining tokens back to the main chain.
-            match tmp_client.close_chain().await {
+            match faucet_client.close_chain().await {
                 Ok(outcome) => {
                     outcome.try_unwrap()?;
                 }
-                Err(err) => tracing::warn!("Failed to close the chain: {err:?}"),
+                Err(err) => tracing::warn!("Failed to close the temporary faucet chain: {err:?}"),
             }
             self.context
                 .lock()
                 .await
-                .forget_chain(&tmp_chain_id)
+                .forget_chain(&faucet_chain_id)
                 .await?;
-            self.tmp_chain_id.lock().await.take();
+            self.faucet_chain_id.lock().await.take();
         }
 
         let chain_id = ChainId::child(message_id);
@@ -220,6 +200,35 @@ where
             certificate_hash: certificate.hash(),
         })
     }
+
+    async fn open_new_faucet_chain(&self) -> Result<ChainId, Error> {
+        let main_client = self
+            .context
+            .lock()
+            .await
+            .make_chain_client(self.main_chain_id)?;
+        let main_balance = main_client.local_balance().await?;
+        let key_pair = main_client.key_pair().await?;
+        let balance = self.faucet_init_balance.min(main_balance.try_sub(MAX_FEE)?);
+        let ownership = main_client.chain_state_view().await?.ownership().clone();
+        let (message_id, certificate) = main_client
+            .open_chain(ownership, ApplicationPermissions::default(), balance)
+            .await?
+            .try_unwrap()?;
+        let chain_id = ChainId::child(message_id);
+        info!("Switching to a new faucet chain {chain_id:8}");
+        self.context
+            .lock()
+            .await
+            .update_wallet_for_new_chain(
+                chain_id,
+                Some(key_pair),
+                certificate.block().header.timestamp,
+            )
+            .await?;
+        *self.faucet_chain_id.lock().await = Some(chain_id);
+        Ok(chain_id)
+    }
 }
 
 impl<C> MutationRoot<C> {
@@ -240,15 +249,15 @@ where
     C: ClientContext,
 {
     main_chain_id: ChainId,
-    tmp_chain_id: Arc<Mutex<Option<ChainId>>>,
+    faucet_chain_id: Arc<Mutex<Option<ChainId>>>,
     context: Arc<Mutex<C>>,
     genesis_config: Arc<GenesisConfig>,
     config: ChainListenerConfig,
     storage: C::Storage,
     port: NonZeroU16,
     amount: Amount,
     end_timestamp: Timestamp,
-    max_claims_per_chain: u32,
+    faucet_init_balance: Amount,
     start_timestamp: Timestamp,
     start_balance: Amount,
 }
@@ -260,14 +269,14 @@ where
     fn clone(&self) -> Self {
         Self {
             main_chain_id: self.main_chain_id,
-            tmp_chain_id: Arc::clone(&self.tmp_chain_id),
+            faucet_chain_id: Arc::clone(&self.faucet_chain_id),
             context: Arc::clone(&self.context),
             genesis_config: Arc::clone(&self.genesis_config),
             config: self.config.clone(),
             storage: self.storage.clone(),
             port: self.port,
             amount: self.amount,
-            max_claims_per_chain: self.max_claims_per_chain,
+            faucet_init_balance: self.faucet_init_balance,
             end_timestamp: self.end_timestamp,
             start_timestamp: self.start_timestamp,
             start_balance: self.start_balance,
@@ -292,21 +301,25 @@ where
         config: ChainListenerConfig,
         storage: C::Storage,
     ) -> anyhow::Result<Self> {
+        let faucet_init_balance = amount
+            .try_add(MAX_FEE)?
+            .try_mul(u128::from(max_claims_per_chain))?
+            .try_add(MAX_FEE)?; // One more block fee for closing the chain.
         let client = context.make_chain_client(chain_id)?;
         let context = Arc::new(Mutex::new(context));
         let start_timestamp = client.storage_client().clock().current_time();
         client.process_inbox().await?;
         let start_balance = client.local_balance().await?;
         Ok(Self {
             main_chain_id: chain_id,
-            tmp_chain_id: Arc::new(Mutex::new(None)),
+            faucet_chain_id: Arc::new(Mutex::new(None)),
             context,
             genesis_config,
             config,
             storage,
             port,
             amount,
-            max_claims_per_chain,
+            faucet_init_balance,
             end_timestamp,
             start_timestamp,
             start_balance,
@@ -316,10 +329,10 @@ where
     pub fn schema(&self) -> Schema<QueryRoot<C>, MutationRoot<C>, EmptySubscription> {
         let mutation_root = MutationRoot {
             main_chain_id: self.main_chain_id,
-            tmp_chain_id: Arc::clone(&self.tmp_chain_id),
+            faucet_chain_id: Arc::clone(&self.faucet_chain_id),
             context: Arc::clone(&self.context),
             amount: self.amount,
-            max_claims_per_chain: self.max_claims_per_chain,
+            faucet_init_balance: self.faucet_init_balance,
             end_timestamp: self.end_timestamp,
             start_timestamp: self.start_timestamp,
             start_balance: self.start_balance,

linera-faucet/server/src/tests.rs

@@ -90,10 +90,10 @@ async fn test_faucet_rate_limiting() {
     let context = Arc::new(Mutex::new(context));
     let root = MutationRoot {
         main_chain_id,
-        tmp_chain_id: Arc::new(Mutex::new(Some(chain_id))),
+        faucet_chain_id: Arc::new(Mutex::new(Some(chain_id))),
         context: context.clone(),
         amount: Amount::from_tokens(1),
-        max_claims_per_chain: 10,
+        faucet_init_balance: Amount::from_tokens(10),
         end_timestamp: Timestamp::from(6000),
         start_timestamp: Timestamp::from(0),
         start_balance: Amount::from_tokens(6),

linera-service/tests/linera_net_tests.rs

@@ -3002,20 +3002,25 @@ async fn test_end_to_end_faucet_chain_limit(config: impl LineraNetConfig) -> Res
     let _guard = INTEGRATION_TEST_GUARD.lock().await;
     tracing::info!("Starting test {}", test_name!());
 
-    let chain_count = 3;
+    let max_claims_per_chain = 3;
 
     let (mut net, faucet_client) = config.instantiate().await?;
 
     let faucet_chain = faucet_client.load_wallet()?.default_chain().unwrap();
 
     let amount = Amount::ONE;
     let mut faucet_service = faucet_client
-        .run_faucet(None, faucet_chain, amount, Some(chain_count as u64))
+        .run_faucet(
+            None,
+            faucet_chain,
+            amount,
+            Some(max_claims_per_chain as u64),
+        )
         .await?;
     let faucet = faucet_service.instance();
 
     // Create one less chain than the configured limit.
-    for _ in 1..chain_count {
+    for _ in 1..max_claims_per_chain {
         let client = net.make_client().await;
         let _ = client
             .wallet_init(&[], FaucetOption::NewChain(&faucet))