1use crate::{
8 ansible::{
9 inventory::AnsibleInventoryType,
10 provisioning::{PrivateNodeProvisionInventory, ProvisionOptions},
11 },
12 error::{Error, Result},
13 get_anvil_node_data, get_bootstrap_cache_url, get_genesis_multiaddr, get_multiaddr,
14 DeploymentInventory, DeploymentType, EvmNetwork, InfraRunOptions, NodeType, TestnetDeployer,
15};
16use colored::Colorize;
17use evmlib::common::U256;
18use log::debug;
19use std::{collections::HashSet, time::Duration};
20
21#[derive(Clone)]
22pub struct UpscaleOptions {
23 pub ansible_verbose: bool,
24 pub ant_version: Option<String>,
25 pub current_inventory: DeploymentInventory,
26 pub desired_client_vm_count: Option<u16>,
27 pub desired_full_cone_private_node_count: Option<u16>,
28 pub desired_full_cone_private_node_vm_count: Option<u16>,
29 pub desired_node_count: Option<u16>,
30 pub desired_node_vm_count: Option<u16>,
31 pub desired_peer_cache_node_count: Option<u16>,
32 pub desired_peer_cache_node_vm_count: Option<u16>,
33 pub desired_symmetric_private_node_count: Option<u16>,
34 pub desired_symmetric_private_node_vm_count: Option<u16>,
35 pub desired_uploaders_count: Option<u16>,
36 pub enable_download_verifier: bool,
37 pub enable_random_verifier: bool,
38 pub enable_performance_verifier: bool,
39 pub funding_wallet_secret_key: Option<String>,
40 pub gas_amount: Option<U256>,
41 pub interval: Duration,
42 pub infra_only: bool,
43 pub max_archived_log_files: u16,
44 pub max_log_files: u16,
45 pub network_dashboard_branch: Option<String>,
46 pub node_env_variables: Option<Vec<(String, String)>>,
47 pub plan: bool,
48 pub public_rpc: bool,
49 pub provision_only: bool,
50 pub token_amount: Option<U256>,
51}
52
53impl TestnetDeployer {
54 pub async fn upscale(&self, options: &UpscaleOptions) -> Result<()> {
55 let is_bootstrap_deploy = matches!(
56 options
57 .current_inventory
58 .environment_details
59 .deployment_type,
60 DeploymentType::Bootstrap
61 );
62
63 if is_bootstrap_deploy
64 && (options.desired_peer_cache_node_count.is_some()
65 || options.desired_peer_cache_node_vm_count.is_some()
66 || options.desired_client_vm_count.is_some())
67 {
68 return Err(Error::InvalidUpscaleOptionsForBootstrapDeployment);
69 }
70
71 let desired_peer_cache_node_vm_count = options
72 .desired_peer_cache_node_vm_count
73 .unwrap_or(options.current_inventory.peer_cache_node_vms.len() as u16);
74 if desired_peer_cache_node_vm_count
75 < options.current_inventory.peer_cache_node_vms.len() as u16
76 {
77 return Err(Error::InvalidUpscaleDesiredPeerCacheVmCount);
78 }
79 debug!("Using {desired_peer_cache_node_vm_count} for desired Peer Cache node VM count");
80
81 let desired_node_vm_count = options
82 .desired_node_vm_count
83 .unwrap_or(options.current_inventory.node_vms.len() as u16);
84 if desired_node_vm_count < options.current_inventory.node_vms.len() as u16 {
85 return Err(Error::InvalidUpscaleDesiredNodeVmCount);
86 }
87 debug!("Using {desired_node_vm_count} for desired node VM count");
88
89 let desired_full_cone_private_node_vm_count = options
90 .desired_full_cone_private_node_vm_count
91 .unwrap_or(options.current_inventory.full_cone_private_node_vms.len() as u16);
92 if desired_full_cone_private_node_vm_count
93 < options.current_inventory.full_cone_private_node_vms.len() as u16
94 {
95 return Err(Error::InvalidUpscaleDesiredFullConePrivateNodeVmCount);
96 }
97 debug!("Using {desired_full_cone_private_node_vm_count} for desired full cone private node VM count");
98
99 let desired_symmetric_private_node_vm_count = options
100 .desired_symmetric_private_node_vm_count
101 .unwrap_or(options.current_inventory.symmetric_private_node_vms.len() as u16);
102 if desired_symmetric_private_node_vm_count
103 < options.current_inventory.symmetric_private_node_vms.len() as u16
104 {
105 return Err(Error::InvalidUpscaleDesiredSymmetricPrivateNodeVmCount);
106 }
107 debug!("Using {desired_symmetric_private_node_vm_count} for desired full cone private node VM count");
108
109 let desired_client_vm_count = options
110 .desired_client_vm_count
111 .unwrap_or(options.current_inventory.client_vms.len() as u16);
112 if desired_client_vm_count < options.current_inventory.client_vms.len() as u16 {
113 return Err(Error::InvalidUpscaleDesiredClientVmCount);
114 }
115 debug!("Using {desired_client_vm_count} for desired Client VM count");
116
117 let desired_peer_cache_node_count = options
118 .desired_peer_cache_node_count
119 .unwrap_or(options.current_inventory.peer_cache_node_count() as u16);
120 if desired_peer_cache_node_count < options.current_inventory.peer_cache_node_count() as u16
121 {
122 return Err(Error::InvalidUpscaleDesiredPeerCacheNodeCount);
123 }
124 debug!("Using {desired_peer_cache_node_count} for desired peer cache node count");
125
126 let desired_node_count = options
127 .desired_node_count
128 .unwrap_or(options.current_inventory.node_count() as u16);
129 if desired_node_count < options.current_inventory.node_count() as u16 {
130 return Err(Error::InvalidUpscaleDesiredNodeCount);
131 }
132 debug!("Using {desired_node_count} for desired node count");
133
134 let desired_full_cone_private_node_count = options
135 .desired_full_cone_private_node_count
136 .unwrap_or(options.current_inventory.full_cone_private_node_count() as u16);
137 if desired_full_cone_private_node_count
138 < options.current_inventory.full_cone_private_node_count() as u16
139 {
140 return Err(Error::InvalidUpscaleDesiredFullConePrivateNodeCount);
141 }
142 debug!(
143 "Using {desired_full_cone_private_node_count} for desired full cone private node count"
144 );
145
146 let desired_symmetric_private_node_count = options
147 .desired_symmetric_private_node_count
148 .unwrap_or(options.current_inventory.symmetric_private_node_count() as u16);
149 if desired_symmetric_private_node_count
150 < options.current_inventory.symmetric_private_node_count() as u16
151 {
152 return Err(Error::InvalidUpscaleDesiredSymmetricPrivateNodeCount);
153 }
154 debug!(
155 "Using {desired_symmetric_private_node_count} for desired symmetric private node count"
156 );
157
158 let mut infra_run_options = InfraRunOptions::generate_existing(
159 &options.current_inventory.name,
160 &options.current_inventory.environment_details.region,
161 &self.terraform_runner,
162 Some(&options.current_inventory.environment_details),
163 )
164 .await?;
165 infra_run_options.peer_cache_node_vm_count = Some(desired_peer_cache_node_vm_count);
166 infra_run_options.node_vm_count = Some(desired_node_vm_count);
167 infra_run_options.full_cone_private_node_vm_count =
168 Some(desired_full_cone_private_node_vm_count);
169 infra_run_options.symmetric_private_node_vm_count =
170 Some(desired_symmetric_private_node_vm_count);
171 infra_run_options.client_vm_count = Some(desired_client_vm_count);
172
173 if options.plan {
174 self.plan(&infra_run_options)?;
175 return Ok(());
176 }
177
178 self.create_or_update_infra(&infra_run_options)
179 .map_err(|err| {
180 println!("Failed to create infra {err:?}");
181 err
182 })?;
183
184 if options.infra_only {
185 return Ok(());
186 }
187
188 let mut provision_options = ProvisionOptions {
189 ant_version: options.ant_version.clone(),
190 binary_option: options.current_inventory.binary_option.clone(),
191 chunk_size: None,
192 client_env_variables: None,
193 enable_download_verifier: options.enable_download_verifier,
194 enable_performance_verifier: options.enable_performance_verifier,
195 enable_random_verifier: options.enable_random_verifier,
196 enable_telegraf: true,
197 enable_uploaders: true,
198 evm_data_payments_address: options
199 .current_inventory
200 .environment_details
201 .evm_details
202 .data_payments_address
203 .clone(),
204 evm_network: options
205 .current_inventory
206 .environment_details
207 .evm_details
208 .network
209 .clone(),
210 evm_payment_token_address: options
211 .current_inventory
212 .environment_details
213 .evm_details
214 .payment_token_address
215 .clone(),
216 evm_rpc_url: options
217 .current_inventory
218 .environment_details
219 .evm_details
220 .rpc_url
221 .clone(),
222 expected_hash: None,
223 expected_size: None,
224 file_address: None,
225 full_cone_private_node_count: desired_full_cone_private_node_count,
226 funding_wallet_secret_key: options.funding_wallet_secret_key.clone(),
227 gas_amount: options.gas_amount,
228 interval: Some(options.interval),
229 log_format: None,
230 max_archived_log_files: options.max_archived_log_files,
231 max_log_files: options.max_log_files,
232 max_uploads: None,
233 name: options.current_inventory.name.clone(),
234 network_id: options.current_inventory.environment_details.network_id,
235 network_dashboard_branch: None,
236 node_count: desired_node_count,
237 node_env_variables: options.node_env_variables.clone(),
238 output_inventory_dir_path: self
239 .working_directory_path
240 .join("ansible")
241 .join("inventory"),
242 peer_cache_node_count: desired_peer_cache_node_count,
243 public_rpc: options.public_rpc,
244 rewards_address: options
245 .current_inventory
246 .environment_details
247 .rewards_address
248 .clone(),
249 symmetric_private_node_count: desired_symmetric_private_node_count,
250 token_amount: None,
251 upload_size: None,
252 uploaders_count: options.desired_uploaders_count,
253 wallet_secret_keys: None,
254 };
255 let mut node_provision_failed = false;
256
257 let (initial_multiaddr, initial_ip_addr) = if is_bootstrap_deploy {
258 get_multiaddr(&self.ansible_provisioner.ansible_runner, &self.ssh_client).map_err(
259 |err| {
260 println!("Failed to get node multiaddr {err:?}");
261 err
262 },
263 )?
264 } else {
265 get_genesis_multiaddr(&self.ansible_provisioner.ansible_runner, &self.ssh_client)
266 .map_err(|err| {
267 println!("Failed to get genesis multiaddr {err:?}");
268 err
269 })?
270 };
271 let initial_network_contacts_url = get_bootstrap_cache_url(&initial_ip_addr);
272 debug!("Retrieved initial peer {initial_multiaddr} and initial network contacts {initial_network_contacts_url}");
273
274 if !is_bootstrap_deploy {
275 self.wait_for_ssh_availability_on_new_machines(
276 AnsibleInventoryType::PeerCacheNodes,
277 &options.current_inventory,
278 )?;
279 self.ansible_provisioner
280 .print_ansible_run_banner("Provision Peer Cache Nodes");
281 match self.ansible_provisioner.provision_nodes(
282 &provision_options,
283 Some(initial_multiaddr.clone()),
284 Some(initial_network_contacts_url.clone()),
285 NodeType::PeerCache,
286 ) {
287 Ok(()) => {
288 println!("Provisioned Peer Cache nodes");
289 }
290 Err(err) => {
291 log::error!("Failed to provision Peer Cache nodes: {err}");
292 node_provision_failed = true;
293 }
294 }
295 }
296
297 self.wait_for_ssh_availability_on_new_machines(
298 AnsibleInventoryType::Nodes,
299 &options.current_inventory,
300 )?;
301 self.ansible_provisioner
302 .print_ansible_run_banner("Provision Normal Nodes");
303 match self.ansible_provisioner.provision_nodes(
304 &provision_options,
305 Some(initial_multiaddr.clone()),
306 Some(initial_network_contacts_url.clone()),
307 NodeType::Generic,
308 ) {
309 Ok(()) => {
310 println!("Provisioned normal nodes");
311 }
312 Err(err) => {
313 log::error!("Failed to provision normal nodes: {err}");
314 node_provision_failed = true;
315 }
316 }
317
318 let private_node_inventory = PrivateNodeProvisionInventory::new(
319 &self.ansible_provisioner,
320 Some(desired_full_cone_private_node_vm_count),
321 Some(desired_symmetric_private_node_vm_count),
322 )?;
323
324 if private_node_inventory.should_provision_full_cone_private_nodes() {
325 let full_cone_nat_gateway_inventory = self
326 .ansible_provisioner
327 .ansible_runner
328 .get_inventory(AnsibleInventoryType::FullConeNatGateway, true)?;
329
330 let full_cone_nat_gateway_new_vms: Vec<_> = full_cone_nat_gateway_inventory
331 .into_iter()
332 .filter(|item| {
333 !options
334 .current_inventory
335 .full_cone_nat_gateway_vms
336 .contains(item)
337 })
338 .collect();
339
340 for vm in full_cone_nat_gateway_new_vms.iter() {
341 self.ssh_client.wait_for_ssh_availability(
342 &vm.public_ip_addr,
343 &self.cloud_provider.get_ssh_user(),
344 )?;
345 }
346
347 let full_cone_nat_gateway_new_vms = if full_cone_nat_gateway_new_vms.is_empty() {
348 None
349 } else {
350 debug!("Full Cone NAT Gateway new VMs: {full_cone_nat_gateway_new_vms:?}");
351 Some(full_cone_nat_gateway_new_vms)
352 };
353
354 match self.ansible_provisioner.provision_full_cone(
355 &provision_options,
356 Some(initial_multiaddr.clone()),
357 Some(initial_network_contacts_url.clone()),
358 private_node_inventory.clone(),
359 full_cone_nat_gateway_new_vms,
360 ) {
361 Ok(()) => {
362 println!("Provisioned Full Cone nodes and Gateway");
363 }
364 Err(err) => {
365 log::error!("Failed to provision Full Cone nodes and Gateway: {err}");
366 node_provision_failed = true;
367 }
368 }
369 }
370
371 if private_node_inventory.should_provision_symmetric_private_nodes() {
372 self.wait_for_ssh_availability_on_new_machines(
373 AnsibleInventoryType::SymmetricNatGateway,
374 &options.current_inventory,
375 )?;
376 self.ansible_provisioner
377 .print_ansible_run_banner("Provision Symmetric NAT Gateway");
378 self.ansible_provisioner
379 .provision_symmetric_nat_gateway(&provision_options, &private_node_inventory)
380 .map_err(|err| {
381 println!("Failed to provision symmetric NAT gateway {err:?}");
382 err
383 })?;
384
385 self.wait_for_ssh_availability_on_new_machines(
386 AnsibleInventoryType::SymmetricPrivateNodes,
387 &options.current_inventory,
388 )?;
389 self.ansible_provisioner
390 .print_ansible_run_banner("Provision Symmetric Private Nodes");
391 match self.ansible_provisioner.provision_symmetric_private_nodes(
392 &mut provision_options,
393 Some(initial_multiaddr.clone()),
394 Some(initial_network_contacts_url.clone()),
395 &private_node_inventory,
396 ) {
397 Ok(()) => {
398 println!("Provisioned symmetric private nodes");
399 }
400 Err(err) => {
401 log::error!("Failed to provision symmetric private nodes: {err}");
402 node_provision_failed = true;
403 }
404 }
405 }
406
407 let should_provision_uploaders =
408 options.desired_uploaders_count.is_some() || options.desired_client_vm_count.is_some();
409 if should_provision_uploaders {
410 if provision_options.evm_network == EvmNetwork::Anvil {
412 let anvil_node_data =
413 get_anvil_node_data(&self.ansible_provisioner.ansible_runner, &self.ssh_client)
414 .map_err(|err| {
415 println!("Failed to get evm testnet data {err:?}");
416 err
417 })?;
418
419 provision_options.funding_wallet_secret_key =
420 Some(anvil_node_data.deployer_wallet_private_key);
421 }
422
423 self.wait_for_ssh_availability_on_new_machines(
424 AnsibleInventoryType::Clients,
425 &options.current_inventory,
426 )?;
427 let genesis_network_contacts = get_bootstrap_cache_url(&initial_ip_addr);
428 self.ansible_provisioner
429 .print_ansible_run_banner("Provision Clients");
430 self.ansible_provisioner
431 .provision_clients(
432 &provision_options,
433 Some(initial_multiaddr.clone()),
434 Some(genesis_network_contacts.clone()),
435 )
436 .await
437 .map_err(|err| {
438 println!("Failed to provision Clients {err:?}");
439 err
440 })?;
441 }
442
443 if node_provision_failed {
444 println!();
445 println!("{}", "WARNING!".yellow());
446 println!("Some nodes failed to provision without error.");
447 println!("This usually means a small number of nodes failed to start on a few VMs.");
448 println!("However, most of the time the deployment will still be usable.");
449 println!("See the output from Ansible to determine which VMs had failures.");
450 }
451
452 Ok(())
453 }
454
455 pub async fn upscale_clients(&self, options: &UpscaleOptions) -> Result<()> {
456 let is_bootstrap_deploy = matches!(
457 options
458 .current_inventory
459 .environment_details
460 .deployment_type,
461 DeploymentType::Bootstrap
462 );
463
464 if is_bootstrap_deploy {
465 return Err(Error::InvalidClientUpscaleDeploymentType(
466 "bootstrap".to_string(),
467 ));
468 }
469
470 let desired_client_vm_count = options
471 .desired_client_vm_count
472 .unwrap_or(options.current_inventory.client_vms.len() as u16);
473 if desired_client_vm_count < options.current_inventory.client_vms.len() as u16 {
474 return Err(Error::InvalidUpscaleDesiredClientVmCount);
475 }
476 debug!("Using {desired_client_vm_count} for desired Client VM count");
477
478 let mut infra_run_options = InfraRunOptions::generate_existing(
479 &options.current_inventory.name,
480 &options.current_inventory.environment_details.region,
481 &self.terraform_runner,
482 Some(&options.current_inventory.environment_details),
483 )
484 .await?;
485 infra_run_options.client_vm_count = Some(desired_client_vm_count);
486
487 if options.plan {
488 self.plan(&infra_run_options)?;
489 return Ok(());
490 }
491
492 if !options.provision_only {
493 self.create_or_update_infra(&infra_run_options)
494 .map_err(|err| {
495 println!("Failed to create infra {err:?}");
496 err
497 })?;
498 }
499
500 if options.infra_only {
501 return Ok(());
502 }
503
504 let (initial_multiaddr, initial_ip_addr) =
505 get_genesis_multiaddr(&self.ansible_provisioner.ansible_runner, &self.ssh_client)
506 .map_err(|err| {
507 println!("Failed to get genesis multiaddr {err:?}");
508 err
509 })?;
510 let initial_network_contacts_url = get_bootstrap_cache_url(&initial_ip_addr);
511 debug!("Retrieved initial peer {initial_multiaddr} and initial network contacts {initial_network_contacts_url}");
512
513 let provision_options = ProvisionOptions {
514 ant_version: options.ant_version.clone(),
515 binary_option: options.current_inventory.binary_option.clone(),
516 chunk_size: None,
517 client_env_variables: None,
518 enable_download_verifier: options.enable_download_verifier,
519 enable_random_verifier: options.enable_random_verifier,
520 enable_performance_verifier: options.enable_performance_verifier,
521 enable_telegraf: true,
522 enable_uploaders: true,
523 evm_data_payments_address: options
524 .current_inventory
525 .environment_details
526 .evm_details
527 .data_payments_address
528 .clone(),
529 evm_network: options
530 .current_inventory
531 .environment_details
532 .evm_details
533 .network
534 .clone(),
535 evm_payment_token_address: options
536 .current_inventory
537 .environment_details
538 .evm_details
539 .payment_token_address
540 .clone(),
541 evm_rpc_url: options
542 .current_inventory
543 .environment_details
544 .evm_details
545 .rpc_url
546 .clone(),
547 expected_hash: None,
548 expected_size: None,
549 file_address: None,
550 full_cone_private_node_count: 0,
551 funding_wallet_secret_key: options.funding_wallet_secret_key.clone(),
552 gas_amount: options.gas_amount,
553 interval: Some(options.interval),
554 log_format: None,
555 max_archived_log_files: options.max_archived_log_files,
556 max_log_files: options.max_log_files,
557 max_uploads: None,
558 name: options.current_inventory.name.clone(),
559 network_id: options.current_inventory.environment_details.network_id,
560 network_dashboard_branch: None,
561 node_count: 0,
562 node_env_variables: None,
563 output_inventory_dir_path: self
564 .working_directory_path
565 .join("ansible")
566 .join("inventory"),
567 peer_cache_node_count: 0,
568 public_rpc: options.public_rpc,
569 rewards_address: options
570 .current_inventory
571 .environment_details
572 .rewards_address
573 .clone(),
574 symmetric_private_node_count: 0,
575 token_amount: options.token_amount,
576 uploaders_count: options.desired_uploaders_count,
577 upload_size: None,
578 wallet_secret_keys: None,
579 };
580
581 self.wait_for_ssh_availability_on_new_machines(
582 AnsibleInventoryType::Clients,
583 &options.current_inventory,
584 )?;
585 self.ansible_provisioner
586 .print_ansible_run_banner("Provision Clients");
587 self.ansible_provisioner
588 .provision_clients(
589 &provision_options,
590 Some(initial_multiaddr),
591 Some(initial_network_contacts_url),
592 )
593 .await
594 .map_err(|err| {
595 println!("Failed to provision clients {err:?}");
596 err
597 })?;
598
599 Ok(())
600 }
601
602 fn wait_for_ssh_availability_on_new_machines(
603 &self,
604 inventory_type: AnsibleInventoryType,
605 current_inventory: &DeploymentInventory,
606 ) -> Result<()> {
607 let inventory = self
608 .ansible_provisioner
609 .ansible_runner
610 .get_inventory(inventory_type, true)?;
611 let old_set: HashSet<_> = match inventory_type {
612 AnsibleInventoryType::Clients => current_inventory
613 .client_vms
614 .iter()
615 .map(|client_vm| &client_vm.vm)
616 .cloned()
617 .collect(),
618 AnsibleInventoryType::PeerCacheNodes => current_inventory
619 .peer_cache_node_vms
620 .iter()
621 .map(|node_vm| &node_vm.vm)
622 .cloned()
623 .collect(),
624 AnsibleInventoryType::Nodes => current_inventory
625 .node_vms
626 .iter()
627 .map(|node_vm| &node_vm.vm)
628 .cloned()
629 .collect(),
630 AnsibleInventoryType::FullConeNatGateway => current_inventory
631 .full_cone_nat_gateway_vms
632 .iter()
633 .cloned()
634 .collect(),
635 AnsibleInventoryType::SymmetricNatGateway => current_inventory
636 .symmetric_nat_gateway_vms
637 .iter()
638 .cloned()
639 .collect(),
640 AnsibleInventoryType::FullConePrivateNodes => current_inventory
641 .full_cone_private_node_vms
642 .iter()
643 .map(|node_vm| &node_vm.vm)
644 .cloned()
645 .collect(),
646 AnsibleInventoryType::SymmetricPrivateNodes => current_inventory
647 .symmetric_private_node_vms
648 .iter()
649 .map(|node_vm| &node_vm.vm)
650 .cloned()
651 .collect(),
652 it => return Err(Error::UpscaleInventoryTypeNotSupported(it.to_string())),
653 };
654 let new_vms: Vec<_> = inventory
655 .into_iter()
656 .filter(|item| !old_set.contains(item))
657 .collect();
658 for vm in new_vms.iter() {
659 self.ssh_client.wait_for_ssh_availability(
660 &vm.public_ip_addr,
661 &self.cloud_provider.get_ssh_user(),
662 )?;
663 }
664 Ok(())
665 }
666}