Skip to main content

linera_client/
client_context.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5
6#[cfg(not(web))]
7use futures::StreamExt as _;
8use futures::{Future, TryStreamExt as _};
9use linera_base::{
10    crypto::{CryptoHash, ValidatorPublicKey},
11    data_types::{ChainDescription, Epoch, Timestamp},
12    identifiers::{Account, AccountOwner, ChainId},
13    ownership::ChainOwnership,
14    time::{Duration, Instant},
15    util::future::FutureSyncExt as _,
16};
17use linera_chain::{manager::LockingBlock, types::ConfirmedBlockCertificate};
18use linera_core::{
19    client::{chain_client, ChainClient, Client, ListeningMode},
20    data_types::{ChainInfo, ChainInfoQuery, ClientOutcome},
21    join_set_ext::JoinSet,
22    node::ValidatorNode,
23    wallet, Environment, JoinSetExt as _, Wallet as _,
24};
25use linera_rpc::node_provider::{NodeOptions, NodeProvider};
26use linera_storage::Storage as _;
27use linera_version::VersionInfo;
28use thiserror_context::Context;
29use tracing::{debug, info, warn};
30#[cfg(not(web))]
31use {
32    crate::{
33        benchmark::{fungible_transfer, Benchmark, BenchmarkError},
34        client_metrics::ClientMetrics,
35    },
36    futures::stream,
37    linera_base::{
38        crypto::AccountPublicKey,
39        data_types::{Amount, BlockHeight},
40        identifiers::{ApplicationId, BlobType},
41    },
42    linera_execution::{
43        system::{OpenChainConfig, SystemOperation},
44        Operation,
45    },
46    std::{collections::HashSet, path::Path},
47    tokio::{sync::mpsc, task},
48};
49#[cfg(feature = "fs")]
50use {
51    linera_base::{
52        data_types::{BlobContent, Bytecode},
53        identifiers::ModuleId,
54        vm::VmRuntime,
55    },
56    linera_core::client::create_bytecode_blobs,
57    std::{fs, path::PathBuf},
58};
59
60use crate::{
61    chain_listener::{self, ClientContext as _},
62    client_options::{ChainOwnershipConfig, Options},
63    config::GenesisConfig,
64    error, util, Error,
65};
66
67/// Results from querying a validator about version, network description, and chain info.
68pub struct ValidatorQueryResults {
69    /// The validator's version information.
70    pub version_info: Result<VersionInfo, Error>,
71    /// The validator's genesis config hash.
72    pub genesis_config_hash: Result<CryptoHash, Error>,
73    /// The validator's chain info (if valid and signature check passed).
74    pub chain_info: Result<ChainInfo, Error>,
75}
76
77impl ValidatorQueryResults {
78    /// Returns a vector of references to all errors in the query results.
79    pub fn errors(&self) -> Vec<&Error> {
80        let mut errors = Vec::new();
81        if let Err(e) = &self.version_info {
82            errors.push(e);
83        }
84        if let Err(e) = &self.genesis_config_hash {
85            errors.push(e);
86        }
87        if let Err(e) = &self.chain_info {
88            errors.push(e);
89        }
90        errors
91    }
92
93    /// Prints validator information to stdout.
94    ///
95    /// Prints public key, address, and optionally weight, version info, and chain info.
96    /// If `reference` is provided, only prints fields that differ from the reference.
97    pub fn print(
98        &self,
99        public_key: Option<&ValidatorPublicKey>,
100        address: Option<&str>,
101        weight: Option<u64>,
102        reference: Option<&ValidatorQueryResults>,
103    ) {
104        if let Some(key) = public_key {
105            println!("Public key: {key}");
106        }
107        if let Some(address) = address {
108            println!("Address: {address}");
109        }
110        if let Some(w) = weight {
111            println!("Weight: {w}");
112        }
113
114        let ref_version = reference.and_then(|ref_results| ref_results.version_info.as_ref().ok());
115        match &self.version_info {
116            Ok(version_info) => {
117                if ref_version.is_none_or(|ref_v| ref_v.crate_version != version_info.crate_version)
118                {
119                    println!("Linera protocol: v{}", version_info.crate_version);
120                }
121                if ref_version.is_none_or(|ref_v| ref_v.rpc_hash != version_info.rpc_hash) {
122                    println!("RPC API hash: {}", version_info.rpc_hash);
123                }
124                if ref_version.is_none_or(|ref_v| ref_v.graphql_hash != version_info.graphql_hash) {
125                    println!("GraphQL API hash: {}", version_info.graphql_hash);
126                }
127                if ref_version.is_none_or(|ref_v| ref_v.wit_hash != version_info.wit_hash) {
128                    println!("WIT API hash: v{}", version_info.wit_hash);
129                }
130                if ref_version.is_none_or(|ref_v| {
131                    (&ref_v.git_commit, ref_v.git_dirty)
132                        != (&version_info.git_commit, version_info.git_dirty)
133                }) {
134                    println!(
135                        "Source code: {}/tree/{}{}",
136                        env!("CARGO_PKG_REPOSITORY"),
137                        version_info.git_commit,
138                        if version_info.git_dirty {
139                            " (dirty)"
140                        } else {
141                            ""
142                        }
143                    );
144                }
145            }
146            Err(err) => println!("Error getting version info: {err}"),
147        }
148
149        let ref_genesis_hash =
150            reference.and_then(|ref_results| ref_results.genesis_config_hash.as_ref().ok());
151        match &self.genesis_config_hash {
152            Ok(hash) if ref_genesis_hash.is_some_and(|ref_hash| ref_hash == hash) => {}
153            Ok(hash) => println!("Genesis config hash: {hash}"),
154            Err(err) => println!("Error getting genesis config: {err}"),
155        }
156
157        let ref_info = reference.and_then(|ref_results| ref_results.chain_info.as_ref().ok());
158        match &self.chain_info {
159            Ok(info) => {
160                if ref_info.is_none_or(|ref_info| info.block_hash != ref_info.block_hash) {
161                    if let Some(hash) = info.block_hash {
162                        println!("Block hash: {hash}");
163                    } else {
164                        println!("Block hash: None");
165                    }
166                }
167                if ref_info
168                    .is_none_or(|ref_info| info.next_block_height != ref_info.next_block_height)
169                {
170                    println!("Next height: {}", info.next_block_height);
171                }
172                if ref_info.is_none_or(|ref_info| info.timestamp != ref_info.timestamp) {
173                    println!("Timestamp: {}", info.timestamp);
174                }
175                if ref_info.is_none_or(|ref_info| info.epoch != ref_info.epoch) {
176                    println!("Epoch: {}", info.epoch);
177                }
178                if ref_info.is_none_or(|ref_info| {
179                    info.manager.current_round != ref_info.manager.current_round
180                }) {
181                    println!("Round: {}", info.manager.current_round);
182                }
183                if let Some(leader) = info.manager.leader {
184                    println!("Leader: {leader}");
185                }
186                if let Some(locking) = &info.manager.requested_locking {
187                    match &**locking {
188                        LockingBlock::Fast(proposal) => {
189                            println!(
190                                "Locking fast block from {}",
191                                proposal.content.block.timestamp
192                            );
193                        }
194                        LockingBlock::Regular(validated) => {
195                            println!(
196                                "Locking block {} in {} from {}",
197                                validated.hash(),
198                                validated.round,
199                                validated.block().header.timestamp
200                            );
201                        }
202                    }
203                }
204            }
205            Err(err) => println!("Error getting chain info: {err}"),
206        }
207        println!();
208    }
209}
210
211/// The state shared by the client commands: the core client, wallet configuration, and
212/// network timeouts.
213pub struct ClientContext<Env: Environment> {
214    /// The core client used to interact with chains and validators.
215    pub client: Arc<Client<Env>>,
216    /// The genesis configuration of the network.
217    // TODO(#5083): this doesn't really need to be stored
218    pub genesis_config: crate::config::GenesisConfig,
219    /// The timeout for sending requests to validators.
220    pub send_timeout: Duration,
221    /// The timeout for receiving responses from validators.
222    pub recv_timeout: Duration,
223    /// The delay before retrying a failed request to a validator.
224    pub retry_delay: Duration,
225    /// The maximum number of times to retry a failed request to a validator.
226    pub max_retries: u32,
227    /// The maximum backoff between retries of a failed request to a validator.
228    pub max_backoff: Duration,
229    /// The set of background tasks listening for chain notifications.
230    pub chain_listeners: JoinSet,
231    /// The default chain used when no chain is explicitly specified.
232    // TODO(#5082): move this into the upstream UI layers (maybe just the CLI)
233    pub default_chain: Option<ChainId>,
234    /// The metrics collector, if metrics collection is enabled.
235    #[cfg(not(web))]
236    pub client_metrics: Option<ClientMetrics>,
237}
238
239impl<Env: Environment> chain_listener::ClientContext for ClientContext<Env> {
240    type Environment = Env;
241
242    fn wallet(&self) -> &Env::Wallet {
243        self.client.wallet()
244    }
245
246    fn storage(&self) -> &Env::Storage {
247        self.client.storage_client()
248    }
249
250    fn client(&self) -> &Arc<Client<Env>> {
251        &self.client
252    }
253
254    #[cfg(not(web))]
255    fn timing_sender(
256        &self,
257    ) -> Option<mpsc::UnboundedSender<(u64, linera_core::client::TimingType)>> {
258        self.client_metrics
259            .as_ref()
260            .map(|metrics| metrics.timing_sender.clone())
261    }
262
263    async fn update_wallet_for_new_chain(
264        &mut self,
265        chain_id: ChainId,
266        owner: Option<AccountOwner>,
267        timestamp: Timestamp,
268        epoch: Epoch,
269    ) -> Result<(), Error> {
270        self.update_wallet_for_new_chain(chain_id, owner, timestamp, epoch)
271            .make_sync()
272            .await
273    }
274
275    async fn update_wallet(&mut self, chain_client: &ChainClient<Env>) -> Result<(), Error> {
276        self.update_wallet_from_client(chain_client)
277            .make_sync()
278            .await
279    }
280}
281
282impl<S, Si, W> ClientContext<linera_core::environment::Impl<S, NodeProvider, Si, W>>
283where
284    S: linera_core::environment::Storage,
285    Si: linera_core::environment::Signer,
286    W: linera_core::environment::Wallet,
287{
288    // not worth refactoring this because
289    // https://github.com/linera-io/linera-protocol/issues/5082
290    // https://github.com/linera-io/linera-protocol/issues/5083
291    /// Creates a new client context from the given storage, wallet, signer, and options.
292    #[expect(clippy::too_many_arguments)]
293    pub async fn new(
294        storage: S,
295        wallet: W,
296        signer: Si,
297        options: &Options,
298        default_chain: Option<ChainId>,
299        genesis_config: GenesisConfig,
300        block_cache_size: usize,
301        execution_state_cache_size: usize,
302    ) -> Result<Self, Error> {
303        #[cfg(not(web))]
304        let timing_config = options.to_timing_config();
305        let node_provider = NodeProvider::new(NodeOptions {
306            send_timeout: options.send_timeout,
307            recv_timeout: options.recv_timeout,
308            retry_delay: options.retry_delay,
309            max_retries: options.max_retries,
310            max_backoff: options.max_backoff,
311        });
312        let chain_modes: Vec<_> = wallet
313            .items()
314            .map_ok(|(id, _chain)| (id, ListeningMode::FullChain))
315            .try_collect()
316            .await
317            .map_err(error::Inner::wallet)?;
318        let name = match chain_modes.len() {
319            0 => "Client node".to_string(),
320            1 => format!("Client node for {:.8}", chain_modes[0].0),
321            n => format!(
322                "Client node for {:.8} and {} others",
323                chain_modes[0].0,
324                n - 1
325            ),
326        };
327
328        let client = Client::new(
329            linera_core::environment::Impl {
330                network: node_provider,
331                storage,
332                signer,
333                wallet,
334            },
335            genesis_config.admin_chain_id(),
336            options.long_lived_services,
337            chain_modes,
338            name,
339            util::non_zero_duration(options.chain_worker_ttl),
340            util::non_zero_duration(options.sender_chain_worker_ttl),
341            options.cross_chain_batch_size_limit,
342            options.to_chain_client_options(),
343            &options.to_requests_scheduler_config(),
344            block_cache_size,
345            execution_state_cache_size,
346        );
347
348        #[cfg(not(web))]
349        let client_metrics = if timing_config.enabled {
350            Some(ClientMetrics::new(timing_config))
351        } else {
352            None
353        };
354
355        Ok(ClientContext {
356            client: Arc::new(client),
357            default_chain,
358            genesis_config,
359            send_timeout: options.send_timeout,
360            recv_timeout: options.recv_timeout,
361            retry_delay: options.retry_delay,
362            max_retries: options.max_retries,
363            max_backoff: options.max_backoff,
364            chain_listeners: JoinSet::default(),
365            #[cfg(not(web))]
366            client_metrics,
367        })
368    }
369}
370
371impl<Env: Environment> ClientContext<Env> {
372    // TODO(#5084) this (and other injected dependencies) should not be re-exposed by the
373    // client interface
374    /// Returns a reference to the wallet.
375    pub fn wallet(&self) -> &Env::Wallet {
376        self.client.wallet()
377    }
378
379    /// Returns the ID of the admin chain.
380    pub fn admin_chain_id(&self) -> ChainId {
381        self.client.admin_chain_id()
382    }
383
384    /// Retrieve the default account. Current this is the common account of the default
385    /// chain.
386    pub fn default_account(&self) -> Account {
387        Account::chain(self.default_chain())
388    }
389
390    /// Retrieve the default chain.
391    pub fn default_chain(&self) -> ChainId {
392        self.default_chain
393            .expect("default chain requested but none set")
394    }
395
396    /// Returns the lowest non-admin chain ID in the wallet.
397    pub async fn first_non_admin_chain(&self) -> Result<ChainId, Error> {
398        let admin_chain_id = self.admin_chain_id();
399        let chain_ids = self
400            .wallet()
401            .chain_ids()
402            .try_filter(|chain_id| futures::future::ready(*chain_id != admin_chain_id))
403            .try_collect::<Vec<ChainId>>()
404            .await
405            .map_err(Error::wallet)?;
406        Ok(chain_ids
407            .into_iter()
408            .min()
409            .expect("No non-admin chain specified in wallet with no non-admin chain"))
410    }
411
412    /// Creates a node provider configured with this context's network options.
413    // TODO(#5084) this should match the `NodeProvider` from the `Environment`
414    pub fn make_node_provider(&self) -> NodeProvider {
415        NodeProvider::new(self.make_node_options())
416    }
417
418    fn make_node_options(&self) -> NodeOptions {
419        NodeOptions {
420            send_timeout: self.send_timeout,
421            recv_timeout: self.recv_timeout,
422            retry_delay: self.retry_delay,
423            max_retries: self.max_retries,
424            max_backoff: self.max_backoff,
425        }
426    }
427
428    /// Returns the client metrics, if metrics collection is enabled.
429    #[cfg(not(web))]
430    pub fn client_metrics(&self) -> Option<&ClientMetrics> {
431        self.client_metrics.as_ref()
432    }
433
434    /// Updates the wallet entry for the client's chain from its current chain info.
435    pub async fn update_wallet_from_client<Env_: Environment>(
436        &self,
437        chain_client: &ChainClient<Env_>,
438    ) -> Result<(), Error> {
439        let info = chain_client.chain_info().await?;
440        let existing_owner = self
441            .wallet()
442            .get(info.chain_id)
443            .await
444            .map_err(error::Inner::wallet)?
445            .and_then(|chain| chain.owner);
446
447        // Only persist proposals that were made in the fast round: they need to be
448        // remembered across sessions to make sure there are no conflicting fast proposals.
449        let pending_proposal = chain_client
450            .pending_proposal()
451            .await
452            .filter(|p| p.round.is_some_and(|r| r.is_fast()));
453        self.wallet()
454            .insert(
455                info.chain_id,
456                wallet::Chain {
457                    pending_proposal,
458                    owner: existing_owner,
459                    ..info.as_ref().into()
460                },
461            )
462            .await
463            .map_err(error::Inner::wallet)?;
464
465        Ok(())
466    }
467
468    /// Remembers the new chain and its owner (if any) in the wallet.
469    pub async fn update_wallet_for_new_chain(
470        &mut self,
471        chain_id: ChainId,
472        owner: Option<AccountOwner>,
473        timestamp: Timestamp,
474        epoch: Epoch,
475    ) -> Result<(), Error> {
476        self.wallet()
477            .try_insert(
478                chain_id,
479                linera_core::wallet::Chain::new(owner, epoch, timestamp),
480            )
481            .await
482            .map_err(error::Inner::wallet)?;
483        Ok(())
484    }
485
486    /// Registers a chain from its description: initializes local storage, adds to
487    /// wallet, and starts tracking it for cross-chain message delivery.
488    pub async fn extend_with_chain(
489        &mut self,
490        description: ChainDescription,
491        owner: Option<AccountOwner>,
492    ) -> Result<(), Error> {
493        let chain_id = description.id();
494        self.client
495            .storage_client()
496            .create_chain(description.clone())
497            .await?;
498        self.wallet()
499            .try_insert(
500                chain_id,
501                linera_core::wallet::Chain::new(
502                    owner,
503                    description.config().epoch,
504                    description.timestamp(),
505                ),
506            )
507            .await
508            .map_err(error::Inner::wallet)?;
509        self.client
510            .extend_chain_mode(chain_id, ListeningMode::FullChain);
511        Ok(())
512    }
513
514    /// Processes the chain's inbox, waiting for round timeouts, and updates the wallet.
515    pub async fn process_inbox(
516        &mut self,
517        chain_client: &ChainClient<Env>,
518    ) -> Result<Vec<ConfirmedBlockCertificate>, Error> {
519        let mut certificates = Vec::new();
520        // Try processing the inbox optimistically without waiting for validator notifications.
521        let (new_certificates, maybe_timeout) = {
522            chain_client.synchronize_from_validators().await?;
523            let result = chain_client.process_inbox_without_prepare().await;
524            self.update_wallet_from_client(chain_client).await?;
525            result?
526        };
527        certificates.extend(new_certificates);
528        if maybe_timeout.is_none() {
529            return Ok(certificates);
530        }
531
532        // Start listening for notifications, so we learn about new rounds and blocks.
533        let (listener, _listen_handle, mut notification_stream) = chain_client.listen().await?;
534        self.chain_listeners.spawn_task(listener);
535
536        loop {
537            let (new_certificates, maybe_timeout) = {
538                let result = chain_client.process_inbox().await;
539                self.update_wallet_from_client(chain_client).await?;
540                result?
541            };
542            certificates.extend(new_certificates);
543            if let Some(timestamp) = maybe_timeout {
544                util::wait_for_next_round(&mut notification_stream, timestamp).await
545            } else {
546                return Ok(certificates);
547            }
548        }
549    }
550
551    /// Assigns the given chain to the owner, tracking it and recording it in the wallet.
552    pub async fn assign_new_chain_to_key(
553        &mut self,
554        chain_id: ChainId,
555        owner: AccountOwner,
556    ) -> Result<(), Error> {
557        self.client
558            .extend_chain_mode(chain_id, ListeningMode::FullChain);
559        let chain_client = self.make_chain_client(chain_id).await?;
560
561        // Ensure we have the chain description blob.
562        chain_client.get_chain_description().await?;
563
564        // Synchronize and get chain info.
565        chain_client.synchronize_from_validators().await?;
566        let info = chain_client.chain_info().await?;
567
568        // Validate that the owner can propose on this chain (either as owner or via
569        // open_multi_leader_rounds).
570        if !info
571            .manager
572            .ownership
573            .can_propose_in_multi_leader_round(&owner)
574        {
575            tracing::error!("Chain {chain_id} is not owned by {owner}.");
576            return Err(error::Inner::ChainOwnership.into());
577        }
578
579        // Try to modify existing chain entry, setting the owner.
580        let modified = self
581            .wallet()
582            .modify(chain_id, |chain| chain.owner = Some(owner))
583            .await
584            .map_err(error::Inner::wallet)?;
585        // If the chain didn't exist, insert a new entry.
586        if modified.is_none() {
587            self.wallet()
588                .insert(
589                    chain_id,
590                    wallet::Chain {
591                        owner: Some(owner),
592                        timestamp: info.timestamp,
593                        epoch: Some(info.epoch),
594                        ..Default::default()
595                    },
596                )
597                .await
598                .map_err(error::Inner::wallet)
599                .context("assigning new chain")?;
600        }
601        Ok(())
602    }
603
604    /// Applies the given function to the chain client.
605    ///
606    /// Updates the wallet regardless of the outcome. As long as the function returns a round
607    /// timeout, it will wait and retry.
608    pub async fn apply_client_command<E, F, Fut, T>(
609        &mut self,
610        chain_client: &ChainClient<Env>,
611        mut f: F,
612    ) -> Result<T, Error>
613    where
614        F: FnMut(&ChainClient<Env>) -> Fut,
615        Fut: Future<Output = Result<ClientOutcome<T>, E>>,
616        Error: From<E>,
617    {
618        chain_client.prepare_chain().await?;
619        // Try applying f optimistically without validator notifications. Return if committed.
620        let result = f(chain_client).await;
621        self.update_wallet_from_client(chain_client).await?;
622        match result? {
623            ClientOutcome::Committed(t) => return Ok(t),
624            ClientOutcome::Conflict(certificate) => {
625                return Err(chain_client::Error::Conflict(certificate.hash()).into());
626            }
627            ClientOutcome::WaitForTimeout(_) => {}
628        }
629
630        // Start listening for notifications, so we learn about new rounds and blocks.
631        let (listener, _listen_handle, mut notification_stream) = chain_client.listen().await?;
632        self.chain_listeners.spawn_task(listener);
633
634        loop {
635            // Try applying f. Return if committed.
636            let result = f(chain_client).await;
637            self.update_wallet_from_client(chain_client).await?;
638            let timeout = match result? {
639                ClientOutcome::Committed(t) => return Ok(t),
640                ClientOutcome::Conflict(certificate) => {
641                    return Err(chain_client::Error::Conflict(certificate.hash()).into());
642                }
643                ClientOutcome::WaitForTimeout(timeout) => timeout,
644            };
645            // Otherwise wait and try again in the next round.
646            util::wait_for_next_round(&mut notification_stream, timeout).await;
647        }
648    }
649
650    /// Returns the ownership configuration of the given chain.
651    pub async fn ownership(&mut self, chain_id: Option<ChainId>) -> Result<ChainOwnership, Error> {
652        let chain_id = chain_id.unwrap_or_else(|| self.default_chain());
653        let chain_client = self.make_chain_client(chain_id).await?;
654        let info = chain_client.chain_info().await?;
655        Ok(info.manager.ownership)
656    }
657
658    /// Changes the ownership configuration of the given chain.
659    pub async fn change_ownership(
660        &mut self,
661        chain_id: Option<ChainId>,
662        ownership_config: ChainOwnershipConfig,
663    ) -> Result<(), Error> {
664        let chain_id = chain_id.unwrap_or_else(|| self.default_chain());
665        let chain_client = self.make_chain_client(chain_id).await?;
666        info!(
667            ?ownership_config, %chain_id, preferred_owner=?chain_client.preferred_owner(),
668            "Changing ownership of a chain"
669        );
670        let time_start = Instant::now();
671        let mut ownership = chain_client.query_chain_ownership().await?;
672        ownership_config.update(&mut ownership)?;
673
674        if ownership.super_owners.is_empty() && ownership.owners.is_empty() {
675            tracing::error!("At least one owner or super owner of the chain has to be set.");
676            return Err(error::Inner::ChainOwnership.into());
677        }
678
679        let certificate = self
680            .apply_client_command(&chain_client, |chain_client| {
681                let ownership = ownership.clone();
682                let chain_client = chain_client.clone();
683                async move {
684                    chain_client
685                        .change_ownership(ownership)
686                        .await
687                        .map_err(Error::from)
688                        .context("Failed to change ownership")
689                }
690            })
691            .await?;
692        let time_total = time_start.elapsed();
693        info!("Operation confirmed after {} ms", time_total.as_millis());
694        debug!("{:?}", certificate);
695        Ok(())
696    }
697
698    /// Sets the preferred owner used to propose blocks on the given chain.
699    pub async fn set_preferred_owner(
700        &mut self,
701        chain_id: Option<ChainId>,
702        preferred_owner: AccountOwner,
703    ) -> Result<(), Error> {
704        let chain_id = chain_id.unwrap_or_else(|| self.default_chain());
705        let mut chain_client = self.make_chain_client(chain_id).await?;
706        let old_owner = chain_client.preferred_owner();
707        info!(%chain_id, ?old_owner, %preferred_owner, "Changing preferred owner for chain");
708        chain_client.set_preferred_owner(preferred_owner);
709        self.update_wallet_from_client(&chain_client).await?;
710        info!("New preferred owner set");
711        Ok(())
712    }
713
714    /// Checks that the validator's version info is compatible with the local version.
715    pub async fn check_compatible_version_info(
716        &self,
717        address: &str,
718        node: &impl ValidatorNode,
719    ) -> Result<VersionInfo, Error> {
720        match node.get_version_info().await {
721            Ok(version_info) if version_info.is_compatible_with(&linera_version::VERSION_INFO) => {
722                debug!(
723                    "Version information for validator {address}: {}",
724                    version_info
725                );
726                Ok(version_info)
727            }
728            Ok(version_info) => Err(error::Inner::UnexpectedVersionInfo {
729                remote: Box::new(version_info),
730                local: Box::new(linera_version::VERSION_INFO.clone()),
731            }
732            .into()),
733            Err(error) => Err(error::Inner::UnavailableVersionInfo {
734                address: address.to_string(),
735                error: Box::new(error),
736            }
737            .into()),
738        }
739    }
740
741    /// Checks that the validator's network description matches the local genesis config.
742    pub async fn check_matching_network_description(
743        &self,
744        address: &str,
745        node: &impl ValidatorNode,
746    ) -> Result<CryptoHash, Error> {
747        let network_description = self.genesis_config.network_description();
748        match node.get_network_description().await {
749            Ok(description) => {
750                if description == network_description {
751                    Ok(description.genesis_config_hash)
752                } else {
753                    Err(error::Inner::UnexpectedNetworkDescription {
754                        remote: Box::new(description),
755                        local: Box::new(network_description),
756                    }
757                    .into())
758                }
759            }
760            Err(error) => Err(error::Inner::UnavailableNetworkDescription {
761                address: address.to_string(),
762                error: Box::new(error),
763            }
764            .into()),
765        }
766    }
767
768    /// Queries a validator for the given chain's info and verifies its signature.
769    pub async fn check_validator_chain_info_response(
770        &self,
771        public_key: Option<&ValidatorPublicKey>,
772        address: &str,
773        node: &impl ValidatorNode,
774        chain_id: ChainId,
775    ) -> Result<ChainInfo, Error> {
776        let query = ChainInfoQuery::new(chain_id).with_manager_values();
777        match node.handle_chain_info_query(query).await {
778            Ok(response) => {
779                debug!(
780                    "Validator {address} sees chain {chain_id} at block height {} and epoch {:?}",
781                    response.info.next_block_height, response.info.epoch,
782                );
783                if let Some(public_key) = public_key {
784                    if response.check(*public_key).is_ok() {
785                        debug!("Signature for public key {public_key} is OK.");
786                    } else {
787                        return Err(error::Inner::InvalidSignature {
788                            public_key: *public_key,
789                        }
790                        .into());
791                    }
792                } else {
793                    warn!("Not checking signature as public key was not given");
794                }
795                Ok(*response.info)
796            }
797            Err(error) => Err(error::Inner::UnavailableChainInfo {
798                address: address.to_string(),
799                chain_id,
800                error: Box::new(error),
801            }
802            .into()),
803        }
804    }
805
806    /// Query a validator for version info, network description, and chain info.
807    ///
808    /// Returns a `ValidatorQueryResults` struct with the results of all three queries.
809    pub async fn query_validator(
810        &self,
811        address: &str,
812        node: &impl ValidatorNode,
813        chain_id: ChainId,
814        public_key: Option<&ValidatorPublicKey>,
815    ) -> ValidatorQueryResults {
816        let version_info = self.check_compatible_version_info(address, node).await;
817        let genesis_config_hash = self.check_matching_network_description(address, node).await;
818        let chain_info = self
819            .check_validator_chain_info_response(public_key, address, node, chain_id)
820            .await;
821
822        ValidatorQueryResults {
823            version_info,
824            genesis_config_hash,
825            chain_info,
826        }
827    }
828
829    /// Query the local node for version info, network description, and chain info.
830    ///
831    /// Returns a `ValidatorQueryResults` struct with the local node's information.
832    pub async fn query_local_node(
833        &self,
834        chain_id: ChainId,
835    ) -> Result<ValidatorQueryResults, Error> {
836        let version_info = Ok(linera_version::VERSION_INFO.clone());
837        let genesis_config_hash = Ok(self
838            .genesis_config
839            .network_description()
840            .genesis_config_hash);
841        let chain_info = self
842            .make_chain_client(chain_id)
843            .await?
844            .chain_info_with_manager_values()
845            .await
846            .map(|info| *info)
847            .map_err(|e| e.into());
848
849        Ok(ValidatorQueryResults {
850            version_info,
851            genesis_config_hash,
852            chain_info,
853        })
854    }
855}
856
857#[cfg(feature = "fs")]
858impl<Env: Environment> ClientContext<Env> {
859    /// Publishes a module from its contract and service bytecode files.
860    pub async fn publish_module(
861        &mut self,
862        chain_client: &ChainClient<Env>,
863        contract: PathBuf,
864        service: PathBuf,
865        vm_runtime: VmRuntime,
866    ) -> Result<ModuleId, Error> {
867        let (blobs, module_id) = load_bytecode_blobs(&contract, &service, vm_runtime).await?;
868
869        info!("Publishing module");
870        let (module_id, _) = self
871            .apply_client_command(chain_client, |chain_client| {
872                let blobs = blobs.clone();
873                let chain_client = chain_client.clone();
874                async move {
875                    chain_client
876                        .publish_module_blobs(blobs, module_id)
877                        .await
878                        .context("Failed to publish module")
879                }
880            })
881            .await?;
882
883        info!("{}", "Module published successfully!");
884
885        info!("Synchronizing client and processing inbox");
886        self.process_inbox(chain_client).await?;
887        Ok(module_id)
888    }
889
890    /// Publishes a data blob loaded from the given file.
891    pub async fn publish_data_blob(
892        &mut self,
893        chain_client: &ChainClient<Env>,
894        blob_path: PathBuf,
895    ) -> Result<CryptoHash, Error> {
896        info!("Loading data blob file");
897        let blob_bytes = fs::read(&blob_path).map_err(|e| {
898            std::io::Error::new(
899                e.kind(),
900                format!("failed to load data blob bytes from {blob_path:?}: {e}"),
901            )
902        })?;
903
904        info!("Publishing data blob");
905        self.apply_client_command(chain_client, |chain_client| {
906            let blob_bytes = blob_bytes.clone();
907            let chain_client = chain_client.clone();
908            async move {
909                chain_client
910                    .publish_data_blob(blob_bytes)
911                    .await
912                    .context("Failed to publish data blob")
913            }
914        })
915        .await?;
916
917        info!("{}", "Data blob published successfully!");
918        Ok(CryptoHash::new(&BlobContent::new_data(blob_bytes)))
919    }
920
921    // TODO(#2490): Consider removing or renaming this.
922    /// Verifies that a data blob with the given hash is available.
923    pub async fn read_data_blob(
924        &mut self,
925        chain_client: &ChainClient<Env>,
926        hash: CryptoHash,
927    ) -> Result<(), Error> {
928        info!("Verifying data blob");
929        self.apply_client_command(chain_client, |chain_client| {
930            let chain_client = chain_client.clone();
931            async move {
932                chain_client
933                    .read_data_blob(hash)
934                    .await
935                    .context("Failed to verify data blob")
936            }
937        })
938        .await?;
939
940        info!("{}", "Data blob verified successfully!");
941        Ok(())
942    }
943}
944
945#[cfg(all(feature = "fs", not(web)))]
946impl<Env: Environment> ClientContext<Env> {
947    /// Publishes a module along with the JSON-encoded `Formats` description loaded
948    /// from `formats`. The module publication and the formats-registry write
949    /// happen atomically in a single block.
950    pub async fn publish_module_with_formats(
951        &mut self,
952        chain_client: &ChainClient<Env>,
953        contract: PathBuf,
954        service: PathBuf,
955        vm_runtime: VmRuntime,
956        formats: PathBuf,
957        registry_application_id: ApplicationId,
958    ) -> Result<ModuleId, Error> {
959        let owner = chain_client
960            .preferred_owner()
961            .ok_or(error::Inner::ChainOwnership)?;
962        let (blobs, formats_blob_bytes, module_id, registry_op_bytes) = self
963            .prepare_bcs_publication(owner, &contract, &service, vm_runtime, &formats)
964            .await?;
965
966        // Publish the formats data blob in its own block first, so it is committed
967        // before the registry `Write` operation asserts its existence.
968        info!("Publishing the formats data blob");
969        self.apply_client_command(chain_client, |chain_client| {
970            let formats_blob_bytes = formats_blob_bytes.clone();
971            let chain_client = chain_client.clone();
972            async move {
973                chain_client
974                    .publish_data_blob(formats_blob_bytes)
975                    .await
976                    .context("Failed to publish the formats data blob")
977            }
978        })
979        .await?;
980
981        info!("Publishing module and registering its formats");
982        self.apply_client_command(chain_client, |chain_client| {
983            let blobs = blobs.clone();
984            let registry_op_bytes = registry_op_bytes.clone();
985            let chain_client = chain_client.clone();
986            async move {
987                chain_client
988                    .execute_operations(
989                        vec![
990                            Operation::system(SystemOperation::PublishModule { module_id }),
991                            Operation::User {
992                                application_id: registry_application_id,
993                                bytes: registry_op_bytes,
994                            },
995                        ],
996                        blobs,
997                    )
998                    .await
999                    .context("Failed to publish module and register formats")
1000            }
1001        })
1002        .await?;
1003
1004        info!(
1005            "{}",
1006            "Module published and formats registered successfully!"
1007        );
1008        info!("Synchronizing client and processing inbox");
1009        self.process_inbox(chain_client).await?;
1010        Ok(module_id)
1011    }
1012
1013    /// Publishes a module, registers its `Formats` description in the formats
1014    /// registry, and creates an application from that module — all atomically in
1015    /// a single block.
1016    #[allow(clippy::too_many_arguments)]
1017    pub async fn publish_bcs_application(
1018        &mut self,
1019        chain_client: &ChainClient<Env>,
1020        contract: PathBuf,
1021        service: PathBuf,
1022        vm_runtime: VmRuntime,
1023        formats: PathBuf,
1024        registry_application_id: ApplicationId,
1025        parameters: Vec<u8>,
1026        instantiation_argument: Vec<u8>,
1027        required_application_ids: Vec<ApplicationId>,
1028    ) -> Result<(ApplicationId, ModuleId), Error> {
1029        let owner = chain_client
1030            .preferred_owner()
1031            .ok_or(error::Inner::ChainOwnership)?;
1032        let (blobs, formats_blob_bytes, module_id, registry_op_bytes) = self
1033            .prepare_bcs_publication(owner, &contract, &service, vm_runtime, &formats)
1034            .await?;
1035
1036        // Publish the formats data blob in its own block first, so it is committed
1037        // before the registry `Write` operation asserts its existence.
1038        info!("Publishing the formats data blob");
1039        self.apply_client_command(chain_client, |chain_client| {
1040            let formats_blob_bytes = formats_blob_bytes.clone();
1041            let chain_client = chain_client.clone();
1042            async move {
1043                chain_client
1044                    .publish_data_blob(formats_blob_bytes)
1045                    .await
1046                    .context("Failed to publish the formats data blob")
1047            }
1048        })
1049        .await?;
1050
1051        info!("Publishing module, registering its formats and creating the application");
1052        let application_id = self
1053            .apply_client_command(chain_client, |chain_client| {
1054                let blobs = blobs.clone();
1055                let registry_op_bytes = registry_op_bytes.clone();
1056                let parameters = parameters.clone();
1057                let instantiation_argument = instantiation_argument.clone();
1058                let required_application_ids = required_application_ids.clone();
1059                let chain_client = chain_client.clone();
1060                async move {
1061                    let outcome: ClientOutcome<ConfirmedBlockCertificate> = chain_client
1062                        .execute_operations(
1063                            vec![
1064                                Operation::system(SystemOperation::PublishModule { module_id }),
1065                                Operation::User {
1066                                    application_id: registry_application_id,
1067                                    bytes: registry_op_bytes,
1068                                },
1069                                Operation::system(SystemOperation::CreateApplication {
1070                                    module_id,
1071                                    parameters,
1072                                    instantiation_argument,
1073                                    required_application_ids,
1074                                }),
1075                            ],
1076                            blobs,
1077                        )
1078                        .await?;
1079                    outcome.try_map(|certificate| {
1080                        let mut creation: Vec<_> = certificate
1081                            .block()
1082                            .created_blob_ids()
1083                            .into_iter()
1084                            .filter(|blob_id| blob_id.blob_type == BlobType::ApplicationDescription)
1085                            .collect();
1086                        if creation.len() != 1 {
1087                            return Err(chain_client::Error::InternalError(
1088                                "Unexpected number of application descriptions published",
1089                            ));
1090                        }
1091                        let blob_id = creation.pop().expect("checked length");
1092                        Ok(ApplicationId::new(blob_id.hash))
1093                    })
1094                }
1095            })
1096            .await?;
1097
1098        info!(
1099            "{}",
1100            "Module published, formats registered and application created successfully!"
1101        );
1102        info!("Synchronizing client and processing inbox");
1103        self.process_inbox(chain_client).await?;
1104        Ok((application_id, module_id))
1105    }
1106
1107    /// Loads the bytecode files and the SNAP file, building the bytecode blobs, the
1108    /// BCS-encoded formats data blob, and the formats-registry write operation
1109    /// authorized by `owner`. The formats blob is returned separately from the
1110    /// bytecode blobs because the caller publishes it in its own block first (so it
1111    /// is committed before the registry `Write` operation asserts its existence).
1112    async fn prepare_bcs_publication(
1113        &self,
1114        owner: AccountOwner,
1115        contract: &Path,
1116        service: &Path,
1117        vm_runtime: VmRuntime,
1118        formats: &Path,
1119    ) -> Result<
1120        (
1121            Vec<linera_base::data_types::Blob>,
1122            Vec<u8>,
1123            ModuleId,
1124            Vec<u8>,
1125        ),
1126        Error,
1127    > {
1128        let (blobs, module_id) = load_bytecode_blobs(contract, service, vm_runtime).await?;
1129
1130        info!("Loading formats from {formats:?}");
1131        let parsed = read_formats_from_snap(formats)?;
1132        let formats_blob_bytes = bcs::to_bytes(&parsed)?;
1133        let formats_blob_hash = CryptoHash::new(&BlobContent::new_data(formats_blob_bytes.clone()));
1134        let registry_op = linera_sdk::abis::formats_registry::Operation::Write {
1135            owner,
1136            module_id,
1137            blob_hash: linera_base::identifiers::DataBlobHash(formats_blob_hash),
1138        };
1139        let registry_op_bytes = bcs::to_bytes(&registry_op)?;
1140        Ok((blobs, formats_blob_bytes, module_id, registry_op_bytes))
1141    }
1142}
1143
1144/// Reads the contract and service Wasm bytecode files from disk and turns them
1145/// into the blobs needed for module publication. Shared between
1146/// [`ClientContext::publish_module`] and the formats-aware variants so the two
1147/// code paths can't drift on error messages or blob construction.
1148#[cfg(feature = "fs")]
1149async fn load_bytecode_blobs(
1150    contract: &Path,
1151    service: &Path,
1152    vm_runtime: VmRuntime,
1153) -> Result<(Vec<linera_base::data_types::Blob>, ModuleId), Error> {
1154    info!("Loading bytecode files");
1155    let contract_bytecode = Bytecode::load_from_file(contract).map_err(|e| {
1156        std::io::Error::new(
1157            e.kind(),
1158            format!("failed to load contract bytecode from {contract:?}: {e}"),
1159        )
1160    })?;
1161    let service_bytecode = Bytecode::load_from_file(service).map_err(|e| {
1162        std::io::Error::new(
1163            e.kind(),
1164            format!("failed to load service bytecode from {service:?}: {e}"),
1165        )
1166    })?;
1167    Ok(create_bytecode_blobs(contract_bytecode, service_bytecode, vm_runtime).await)
1168}
1169
1170/// Parses the `Formats` description of an application from a SNAP file (the YAML
1171/// snapshot produced by the `format` test of the example applications). The body
1172/// between the `---` frontmatter delimiters is deserialized as
1173/// [`linera_sdk::formats::Formats`]. BCS-serializing the result yields the data-blob
1174/// payload the formats registry expects, so external tooling can produce a blob file
1175/// ready for `linera publish-data-blob` (see the `extract-formats` binary in the
1176/// `formats-registry` example).
1177#[cfg(all(feature = "fs", not(web)))]
1178pub fn read_formats_from_snap(path: &Path) -> Result<linera_sdk::formats::Formats, Error> {
1179    let content = fs::read_to_string(path).map_err(|e| {
1180        std::io::Error::new(e.kind(), format!("failed to read SNAP file {path:?}: {e}"))
1181    })?;
1182    let body = strip_snap_frontmatter(&content).ok_or_else(|| {
1183        std::io::Error::new(
1184            std::io::ErrorKind::InvalidData,
1185            format!("SNAP file {path:?} is missing the `---` frontmatter delimiters"),
1186        )
1187    })?;
1188    serde_yaml_08::from_str(body).map_err(|e| {
1189        std::io::Error::new(
1190            std::io::ErrorKind::InvalidData,
1191            format!("failed to parse SNAP body in {path:?} as Formats: {e}"),
1192        )
1193        .into()
1194    })
1195}
1196
1197#[cfg(all(feature = "fs", not(web)))]
1198fn strip_snap_frontmatter(content: &str) -> Option<&str> {
1199    let rest = content.strip_prefix("---\n")?;
1200    let end = rest.find("\n---\n")?;
1201    Some(&rest[end + "\n---\n".len()..])
1202}
1203
1204#[cfg(all(test, feature = "fs", not(web)))]
1205mod snap_loader_tests {
1206    use super::read_formats_from_snap;
1207
1208    #[test]
1209    fn parses_fungible_snap() {
1210        let path = std::path::Path::new("../examples/fungible/tests/snapshots/format__format.snap");
1211        read_formats_from_snap(path).expect("fungible snap should parse");
1212    }
1213
1214    #[test]
1215    fn parses_social_snap() {
1216        let path = std::path::Path::new("../examples/social/tests/snapshots/format__format.snap");
1217        read_formats_from_snap(path).expect("social snap should parse");
1218    }
1219
1220    #[test]
1221    fn parses_counter_snap() {
1222        let path = std::path::Path::new("../examples/counter/tests/snapshots/format__format.snap");
1223        read_formats_from_snap(path).expect("counter snap should parse");
1224    }
1225}
1226
1227#[cfg(not(web))]
1228impl<Env: Environment> ClientContext<Env> {
1229    /// Prepares the chains and fungible tokens needed to run a benchmark.
1230    pub async fn prepare_for_benchmark(
1231        &mut self,
1232        num_chains: usize,
1233        tokens_per_chain: Amount,
1234        fungible_application_id: Option<ApplicationId>,
1235        pub_keys: Vec<AccountPublicKey>,
1236        chains_config_path: Option<&Path>,
1237        close_chains: bool,
1238    ) -> Result<Vec<ChainClient<Env>>, Error> {
1239        let start = Instant::now();
1240        // Below all block proposals are supposed to succeed without retries, we
1241        // must make sure that all incoming payments have been accepted on-chain
1242        // and that no validator is missing user certificates.
1243        self.process_inboxes_and_force_validator_updates().await;
1244        info!(
1245            "Processed inboxes and forced validator updates in {} ms",
1246            start.elapsed().as_millis()
1247        );
1248
1249        let start = Instant::now();
1250        let (benchmark_chains, chain_clients) = self
1251            .make_benchmark_chains(
1252                num_chains,
1253                tokens_per_chain,
1254                pub_keys,
1255                chains_config_path.is_some(),
1256                close_chains,
1257            )
1258            .await?;
1259        info!(
1260            "Got {} chains in {} ms",
1261            num_chains,
1262            start.elapsed().as_millis()
1263        );
1264
1265        if let Some(id) = fungible_application_id {
1266            let start = Instant::now();
1267            self.supply_fungible_tokens(&benchmark_chains, id).await?;
1268            info!(
1269                "Supplied fungible tokens in {} ms",
1270                start.elapsed().as_millis()
1271            );
1272            // Need to process inboxes to make sure the chains receive the supplied tokens.
1273            let start = Instant::now();
1274            for chain_client in &chain_clients {
1275                chain_client.process_inbox().await?;
1276            }
1277            info!(
1278                "Processed inboxes after supplying fungible tokens in {} ms",
1279                start.elapsed().as_millis()
1280            );
1281        }
1282
1283        let all_chains = Benchmark::<Env>::get_all_chains(chains_config_path, &benchmark_chains)?;
1284        let known_chain_ids: HashSet<_> = benchmark_chains.iter().map(|(id, _)| *id).collect();
1285        let unknown_chain_ids: Vec<_> = all_chains
1286            .iter()
1287            .filter(|id| !known_chain_ids.contains(id))
1288            .copied()
1289            .collect();
1290        if !unknown_chain_ids.is_empty() {
1291            // The current client won't have the blobs for the chains in the other wallets. Even
1292            // though it will eventually get those blobs, we're getting a head start here and
1293            // fetching those blobs in advance.
1294            for chain_id in &unknown_chain_ids {
1295                self.client.get_chain_description(*chain_id).await?;
1296            }
1297        }
1298
1299        Ok(chain_clients)
1300    }
1301
1302    /// Closes the benchmark chains, or processes their inboxes and updates the wallet.
1303    pub async fn wrap_up_benchmark(
1304        &mut self,
1305        chain_clients: Vec<ChainClient<Env>>,
1306        close_chains: bool,
1307        wrap_up_max_in_flight: usize,
1308    ) -> Result<(), Error> {
1309        if close_chains {
1310            info!("Closing chains...");
1311            let chain_ids: Vec<_> = chain_clients.iter().map(|c| c.chain_id()).collect();
1312            let stream = stream::iter(chain_clients)
1313                .map(|chain_client| async move {
1314                    Benchmark::<Env>::close_benchmark_chain(&chain_client).await?;
1315                    info!("Closed chain {:?}", chain_client.chain_id());
1316                    Ok::<(), BenchmarkError>(())
1317                })
1318                .buffer_unordered(wrap_up_max_in_flight);
1319            stream.try_collect::<Vec<_>>().await?;
1320            // Remove closed chains from wallet (the chain listener may have added them).
1321            for chain_id in chain_ids {
1322                if let Err(error) = self.wallet().remove(chain_id).await {
1323                    warn!(%chain_id, %error, "Failed to remove closed chain from wallet");
1324                }
1325            }
1326        } else {
1327            info!("Processing inbox for all chains...");
1328            let stream = stream::iter(chain_clients.clone())
1329                .map(|chain_client| async move {
1330                    chain_client.process_inbox().await?;
1331                    info!("Processed inbox for chain {:?}", chain_client.chain_id());
1332                    Ok::<(), chain_client::Error>(())
1333                })
1334                .buffer_unordered(wrap_up_max_in_flight);
1335            stream.try_collect::<Vec<_>>().await?;
1336
1337            info!("Updating wallet from chain clients...");
1338            for chain_client in chain_clients {
1339                let info = chain_client.chain_info().await?;
1340                let client_owner = chain_client.preferred_owner();
1341                let pending_proposal = chain_client
1342                    .pending_proposal()
1343                    .await
1344                    .filter(|p| p.round.is_some_and(|r| r.is_fast()));
1345                self.wallet()
1346                    .insert(
1347                        info.chain_id,
1348                        wallet::Chain {
1349                            pending_proposal,
1350                            owner: client_owner,
1351                            ..info.as_ref().into()
1352                        },
1353                    )
1354                    .await
1355                    .map_err(error::Inner::wallet)?;
1356            }
1357        }
1358
1359        Ok(())
1360    }
1361
1362    async fn process_inboxes_and_force_validator_updates(&mut self) {
1363        let mut join_set = task::JoinSet::new();
1364
1365        let chain_clients: Vec<_> = self
1366            .wallet()
1367            .owned_chain_ids()
1368            .map_err(|e| error::Inner::wallet(e).into())
1369            .and_then(|id| self.make_chain_client(id))
1370            .try_collect()
1371            .await
1372            .unwrap();
1373
1374        for chain_client in chain_clients {
1375            join_set.spawn(async move {
1376                Self::process_inbox_without_updating_wallet(&chain_client)
1377                    .await
1378                    .expect("Processing inbox should not fail!");
1379                chain_client
1380            });
1381        }
1382
1383        for chain_client in join_set.join_all().await {
1384            self.update_wallet_from_client(&chain_client).await.unwrap();
1385        }
1386    }
1387
1388    async fn process_inbox_without_updating_wallet(
1389        chain_client: &ChainClient<Env>,
1390    ) -> Result<Vec<ConfirmedBlockCertificate>, Error> {
1391        // Try processing the inbox optimistically without waiting for validator notifications.
1392        chain_client.synchronize_from_validators().await?;
1393        let (certificates, maybe_timeout) = chain_client.process_inbox_without_prepare().await?;
1394        assert!(
1395            maybe_timeout.is_none(),
1396            "Should not timeout within benchmark!"
1397        );
1398
1399        Ok(certificates)
1400    }
1401
1402    /// Creates chains if necessary, and returns a map of exactly `num_chains` chain IDs
1403    /// with key pairs, as well as a map of the chain clients.
1404    ///
1405    /// If `close_chains` is true, chains are not looked up from or stored in the wallet,
1406    /// since they will be closed after the benchmark and shouldn't be reused.
1407    async fn make_benchmark_chains(
1408        &mut self,
1409        num_chains: usize,
1410        balance: Amount,
1411        pub_keys: Vec<AccountPublicKey>,
1412        wallet_only: bool,
1413        close_chains: bool,
1414    ) -> Result<(Vec<(ChainId, AccountOwner)>, Vec<ChainClient<Env>>), Error> {
1415        let mut chains_found_in_wallet = 0;
1416        let mut benchmark_chains = Vec::with_capacity(num_chains);
1417        let mut chain_clients = Vec::with_capacity(num_chains);
1418        let start = Instant::now();
1419
1420        // When close_chains is true and we're creating our own chains (not wallet_only),
1421        // skip wallet lookup to avoid picking up existing chains that would then be closed.
1422        // When wallet_only is true, chains were pre-created by the parent process and must
1423        // be read from the wallet.
1424        if !close_chains || wallet_only {
1425            let mut owned_chain_ids = std::pin::pin!(self.wallet().owned_chain_ids());
1426            while let Some(chain_id) = owned_chain_ids.next().await {
1427                let chain_id = chain_id.map_err(error::Inner::wallet)?;
1428                if chains_found_in_wallet == num_chains {
1429                    break;
1430                }
1431                let chain_client = self.make_chain_client(chain_id).await?;
1432                let ownership = chain_client.chain_info().await?.manager.ownership;
1433                if !ownership.owners.is_empty() || ownership.super_owners.len() != 1 {
1434                    continue;
1435                }
1436                let owner = *ownership.super_owners.first().unwrap();
1437                chain_client.process_inbox().await?;
1438                benchmark_chains.push((chain_id, owner));
1439                chain_clients.push(chain_client);
1440                chains_found_in_wallet += 1;
1441            }
1442            info!(
1443                "Got {} chains from the wallet in {} ms",
1444                benchmark_chains.len(),
1445                start.elapsed().as_millis()
1446            );
1447        }
1448
1449        let num_chains_to_create = num_chains - chains_found_in_wallet;
1450
1451        let default_chain_client = self.make_chain_client(self.default_chain()).await?;
1452
1453        if num_chains_to_create > 0 {
1454            if wallet_only {
1455                return Err(
1456                    error::Inner::Benchmark(BenchmarkError::NotEnoughChainsInWallet(
1457                        num_chains,
1458                        chains_found_in_wallet,
1459                    ))
1460                    .into(),
1461                );
1462            }
1463            let mut pub_keys_iter = pub_keys.into_iter().take(num_chains_to_create);
1464            let operations_per_block = 900; // Over this we seem to hit the block size limits.
1465            for i in (0..num_chains_to_create).step_by(operations_per_block) {
1466                let num_new_chains = operations_per_block.min(num_chains_to_create - i);
1467                // Each chain gets its own unique owner (previously all chains in a batch
1468                // shared one owner, which could cause conflicts during benchmarking).
1469                let owners: Vec<AccountOwner> = (&mut pub_keys_iter)
1470                    .take(num_new_chains)
1471                    .map(|pk| pk.into())
1472                    .collect();
1473
1474                let certificate = Self::execute_open_chains_operations(
1475                    &default_chain_client,
1476                    balance,
1477                    owners.clone(),
1478                )
1479                .await?;
1480                info!("Block executed successfully");
1481
1482                let block = certificate.block();
1483                for (i, owner) in owners.into_iter().enumerate() {
1484                    let chain_id = block.body.blobs[i]
1485                        .iter()
1486                        .find(|blob| blob.id().blob_type == BlobType::ChainDescription)
1487                        .map(|blob| ChainId(blob.id().hash))
1488                        .expect("failed to create a new chain");
1489                    self.client
1490                        .extend_chain_mode(chain_id, ListeningMode::FullChain);
1491
1492                    let mut chain_client = self.client.create_chain_client(
1493                        chain_id,
1494                        None,
1495                        BlockHeight::ZERO,
1496                        &None,
1497                        Some(owner),
1498                        self.timing_sender(),
1499                    );
1500                    chain_client.set_preferred_owner(owner);
1501                    chain_client.process_inbox().await?;
1502                    benchmark_chains.push((chain_id, owner));
1503                    chain_clients.push(chain_client);
1504                }
1505            }
1506
1507            info!(
1508                "Created {} chains in {} ms",
1509                num_chains_to_create,
1510                start.elapsed().as_millis()
1511            );
1512        }
1513
1514        // Only update wallet if chains will be reused (not closed after benchmark)
1515        if !close_chains {
1516            info!("Updating wallet from client");
1517            self.update_wallet_from_client(&default_chain_client)
1518                .await?;
1519        }
1520        info!("Retrying pending outgoing messages");
1521        default_chain_client
1522            .retry_pending_outgoing_messages()
1523            .await
1524            .context("outgoing messages to create the new chains should be delivered")?;
1525        info!("Processing default chain inbox");
1526        default_chain_client.process_inbox().await?;
1527
1528        Ok((benchmark_chains, chain_clients))
1529    }
1530
1531    async fn execute_open_chains_operations(
1532        chain_client: &ChainClient<Env>,
1533        balance: Amount,
1534        owners: Vec<AccountOwner>,
1535    ) -> Result<ConfirmedBlockCertificate, Error> {
1536        let operations: Vec<_> = owners
1537            .iter()
1538            .map(|owner| {
1539                let config = OpenChainConfig {
1540                    ownership: ChainOwnership::single_super(*owner),
1541                    balance,
1542                    application_permissions: Default::default(),
1543                };
1544                Operation::system(SystemOperation::OpenChain(config))
1545            })
1546            .collect();
1547        info!("Executing {} OpenChain operations", operations.len());
1548        Ok(chain_client
1549            .execute_operations(operations, vec![])
1550            .await?
1551            .expect("should execute block with OpenChain operations"))
1552    }
1553
1554    /// Supplies fungible tokens to the chains.
1555    async fn supply_fungible_tokens(
1556        &mut self,
1557        key_pairs: &[(ChainId, AccountOwner)],
1558        application_id: ApplicationId,
1559    ) -> Result<(), Error> {
1560        let default_chain_id = self.default_chain();
1561        let default_key = self
1562            .wallet()
1563            .get(default_chain_id)
1564            .await
1565            .unwrap()
1566            .unwrap()
1567            .owner
1568            .unwrap();
1569        // This should be enough to run the benchmark at 1M TPS for an hour.
1570        let amount = Amount::from_nanos(4);
1571        let operations: Vec<Operation> = key_pairs
1572            .iter()
1573            .map(|(chain_id, owner)| {
1574                fungible_transfer(application_id, *chain_id, default_key, *owner, amount)
1575            })
1576            .collect();
1577        let chain_client = self.make_chain_client(default_chain_id).await?;
1578        // Put at most 1000 fungible token operations in each block.
1579        for operation_chunk in operations.chunks(1000) {
1580            chain_client
1581                .execute_operations(operation_chunk.to_vec(), vec![])
1582                .await?
1583                .expect("should execute block with Transfer operations");
1584        }
1585        self.update_wallet_from_client(&chain_client).await?;
1586
1587        Ok(())
1588    }
1589}