1use std::collections::BTreeMap;
4use std::path::PathBuf;
5use std::process::{Command, Stdio};
6use std::time::{Duration, Instant};
7
8use anyhow::{Context, Result, bail, ensure};
9
10use mj_core::config::{AwsAddressSource, Config, ProjectBundle, TargetTemplate, data_dir};
11use mj_core::state::{
12 PodmanWorkspaceLocator, ResourceAllocationKind, SessionRecord, SessionResourceAllocation,
13 TargetLocator, allocation_cpus,
14};
15
16use crate::targets::{
17 self, AwsTemplate, CommandExecutor, CommandOutput, CommandSpec, ContainerTemplate,
18 ImageRefresh, ProjectBundleSpec, ProvisionStage, RepositorySpec, SshTarget,
19};
20
21use super::{Controller, execute_checked};
22
23impl Controller {
24 pub fn session_working_context(
26 &self,
27 session_id: &str,
28 executor: &impl CommandExecutor,
29 ) -> Result<(PathBuf, String)> {
30 let session = self
31 .state
32 .sessions
33 .get(session_id)
34 .context("session is missing")?;
35 let locator = session
36 .target
37 .as_ref()
38 .context("target is still starting")?;
39 let backend = backend_locator(locator, session, &self.config)?;
40 let launch = self.current_worker_launch_config(session_id, &backend)?;
41 let output = executor.execute(&targets::command_on_locator(
42 &backend,
43 session_id,
44 vec![
45 "git".into(),
46 "-C".into(),
47 launch.cwd.to_string_lossy().into_owned(),
48 "rev-parse".into(),
49 "--abbrev-ref".into(),
50 "HEAD".into(),
51 ],
52 "read current session branch",
53 )?)?;
54 let branch = if output.status == 0 {
55 let branch = String::from_utf8(output.stdout).context("decode session branch")?;
56 if branch.trim() == "HEAD" {
57 "detached HEAD".to_owned()
58 } else {
59 branch.trim().to_owned()
60 }
61 } else {
62 let stderr = String::from_utf8_lossy(&output.stderr);
63 if stderr.contains("not a git repository") {
64 "not a git checkout".to_owned()
68 } else {
69 format!(
70 "unavailable: {}",
71 stderr.lines().next().unwrap_or("").trim()
72 )
73 }
74 };
75 Ok((launch.cwd, branch))
76 }
77
78 pub fn session_git_status(
83 &self,
84 session_id: &str,
85 executor: &impl CommandExecutor,
86 ) -> Result<mj_core::local_git::SessionGitStatus> {
87 let (cwd, branch) = self.session_working_context(session_id, executor)?;
88 if branch.starts_with("not a git") || branch.starts_with("unavailable") {
89 return Ok(mj_core::local_git::parse_git_status(
90 cwd, &branch, None, "", "",
91 ));
92 }
93 let session = self
94 .state
95 .sessions
96 .get(session_id)
97 .context("session is missing")?;
98 let locator = session
99 .target
100 .as_ref()
101 .context("target is still starting")?;
102 let backend = backend_locator(locator, session, &self.config)?;
103 let cwd_text = cwd.to_string_lossy().into_owned();
104 let run = |args: &[&str], purpose: &str| -> Result<Option<String>> {
105 let mut command = vec!["git".to_owned(), "-C".to_owned(), cwd_text.clone()];
106 command.extend(args.iter().map(|arg| (*arg).to_owned()));
107 let output = executor.execute(&targets::command_on_locator(
108 &backend, session_id, command, purpose,
109 )?)?;
110 Ok((output.status == 0).then(|| String::from_utf8_lossy(&output.stdout).into_owned()))
111 };
112 let ahead_behind = run(
115 &["rev-list", "--left-right", "--count", "@{upstream}...HEAD"],
116 "count commits against upstream",
117 )?;
118 let numstat = run(
120 &["--no-optional-locks", "diff", "--numstat", "HEAD"],
121 "count changed lines",
122 )?
123 .unwrap_or_default();
124 let porcelain = run(
125 &[
126 "--no-optional-locks",
127 "status",
128 "--porcelain",
129 "--untracked-files=normal",
130 ],
131 "list changed files",
132 )?
133 .unwrap_or_default();
134 Ok(mj_core::local_git::parse_git_status(
135 cwd,
136 &branch,
137 ahead_behind.as_deref(),
138 &numstat,
139 &porcelain,
140 ))
141 }
142
143 pub fn resolve_aws_resource_options(
144 &self,
145 target_id: &str,
146 executor: &impl CommandExecutor,
147 ) -> Result<Vec<SessionResourceAllocation>> {
148 let TargetTemplate::AwsEc2 {
149 aws_profile,
150 region,
151 launch_template,
152 launch_template_version,
153 ..
154 } = self
155 .config
156 .targets
157 .get(target_id)
158 .with_context(|| format!("unknown target template {target_id:?}"))?
159 else {
160 bail!("target {target_id:?} is not an AWS EC2 target");
161 };
162 let profile = aws_profile.as_deref().unwrap_or("default");
163 let launch_key = if launch_template.starts_with("lt-") {
164 "--launch-template-id"
165 } else {
166 "--launch-template-name"
167 };
168 let version = launch_template_version.as_deref().unwrap_or("$Default");
169 let describe_template = CommandSpec::new(
170 "aws",
171 [
172 "--profile",
173 profile,
174 "--region",
175 region,
176 "ec2",
177 "describe-launch-template-versions",
178 launch_key,
179 launch_template,
180 "--versions",
181 version,
182 "--output",
183 "json",
184 ],
185 )
186 .purpose("resolve EC2 launch template instance family");
187 let output = executor.execute(&describe_template)?;
188 if output.status != 0 {
189 bail!(
190 "{} failed with status {}: {}",
191 describe_template.purpose,
192 output.status,
193 String::from_utf8_lossy(&output.stderr).trim()
194 );
195 }
196 let response: serde_json::Value =
197 serde_json::from_slice(&output.stdout).context("parse EC2 launch template response")?;
198 let instance_type = response
199 .pointer("/LaunchTemplateVersions/0/LaunchTemplateData/InstanceType")
200 .and_then(serde_json::Value::as_str)
201 .context("launch template does not specify a concrete instance type")?;
202 let family = instance_type
203 .rsplit_once('.')
204 .map(|(family, _)| family)
205 .context("launch template instance type has no size suffix")?;
206 let filter = format!("Name=instance-type,Values={family}.*");
207 let describe_types = CommandSpec::new(
208 "aws",
209 [
210 "--profile",
211 profile,
212 "--region",
213 region,
214 "ec2",
215 "describe-instance-types",
216 "--filters",
217 &filter,
218 "--output",
219 "json",
220 ],
221 )
222 .purpose("discover EC2 instance sizes");
223 let output = executor.execute(&describe_types)?;
224 if output.status != 0 {
225 bail!(
226 "{} failed with status {}: {}",
227 describe_types.purpose,
228 output.status,
229 String::from_utf8_lossy(&output.stderr).trim()
230 );
231 }
232 let response: serde_json::Value =
233 serde_json::from_slice(&output.stdout).context("parse EC2 instance type response")?;
234 let mut options = response
235 .get("InstanceTypes")
236 .and_then(serde_json::Value::as_array)
237 .context("EC2 instance type response omitted InstanceTypes")?
238 .iter()
239 .filter_map(|entry| {
240 Some(SessionResourceAllocation::AwsEc2 {
241 instance_type: entry.get("InstanceType")?.as_str()?.to_owned(),
242 vcpus: entry.pointer("/VCpuInfo/DefaultVCpus")?.as_u64()?,
243 memory_bytes: entry
244 .pointer("/MemoryInfo/SizeInMiB")?
245 .as_u64()?
246 .checked_mul(1024 * 1024)?,
247 })
248 })
249 .collect::<Vec<_>>();
250 options.sort_by_key(allocation_cpus);
251 if !options.iter().any(|option| allocation_cpus(option) == 8) {
252 bail!("EC2 family {family:?} has no exact 8-vCPU baseline size");
253 }
254 Ok(options)
255 }
256
257 pub fn reconnect_command(&self, session_id: &str) -> Result<CommandSpec> {
258 let session = self
259 .state
260 .sessions
261 .get(session_id)
262 .with_context(|| format!("unknown session {session_id}"))?;
263 session.validate_configuration(&self.config)?;
264 let locator = session.target.as_ref().context("session has no target")?;
265 let backend = backend_locator(locator, session, &self.config)?;
266 targets::reconnect_plan(&backend, session_id)?
267 .commands
268 .into_iter()
269 .next()
270 .context("reconnect plan is empty")
271 }
272
273 fn session_storage_paths(&self, target_id: &str) -> Vec<String> {
275 self.state
276 .sessions
277 .values()
278 .filter(|session| session.target_template_id == target_id && session.state.is_active())
279 .filter_map(|session| {
280 let paths = mj_core::targets::storage::session_storage_paths(
281 session.target.as_ref()?,
282 &session.id,
283 session.project_directory.as_deref(),
284 );
285 Some(paths.all().map(str::to_owned).collect::<Vec<_>>())
286 })
287 .flatten()
288 .collect()
289 }
290
291 pub fn deployment_capacity_targets(&self) -> Vec<targets::DeploymentCapacityTarget> {
292 use targets::{DeploymentCapacityKind, DeploymentCapacityTarget};
293
294 let mut local_ids = Vec::new();
295 type HostProbes = (Vec<String>, Vec<(SshTarget, Vec<String>)>);
297 let mut ssh_hosts: BTreeMap<String, HostProbes> = BTreeMap::new();
298 let mut targets = Vec::new();
299 for (target_id, template) in &self.config.targets {
300 match template {
301 TargetTemplate::LocalBare
302 | TargetTemplate::LocalPodman { .. }
303 | TargetTemplate::LocalDocker { .. }
304 | TargetTemplate::AppleContainer { .. } => {
305 local_ids.push(target_id.clone());
306 }
307 TargetTemplate::SshBare { ssh, .. }
308 | TargetTemplate::SshPodman { ssh, .. }
309 | TargetTemplate::SshDocker { ssh, .. } => {
310 let entry = ssh_hosts.entry(ssh.host.clone()).or_default();
311 entry.0.push(target_id.clone());
312 let connection = SshTarget::from(ssh);
313 let paths = match entry.1.iter_mut().find(|(known, _)| *known == connection) {
314 Some((_, paths)) => paths,
315 None => {
316 entry.1.push((connection, Vec::new()));
317 &mut entry.1.last_mut().expect("just pushed").1
318 }
319 };
320 push_unique(paths, ssh_storage_paths(template));
321 push_unique(paths, self.session_storage_paths(target_id));
324 }
325 TargetTemplate::AwsEc2 { .. } => {
326 let mut probes = Vec::new();
327 let mut probe_error = None;
328 for session in self.state.sessions.values().filter(|session| {
329 session.target_template_id == *target_id
330 && session.state.is_active()
331 && session.target.is_some()
332 }) {
333 let result = backend_locator(
334 session.target.as_ref().expect("filtered target"),
335 session,
336 &self.config,
337 )
338 .and_then(|locator| {
339 targets::aws_allocated_capacity_command(&locator, &session.id)
340 });
341 match result {
342 Ok(command) => probes.push(command),
343 Err(error) => probe_error = Some(format!("{error:#}")),
344 }
345 }
346 targets.push(DeploymentCapacityTarget {
347 id: format!("aws:{target_id}"),
348 host: target_id.clone(),
349 target_ids: vec![target_id.clone()],
350 kind: DeploymentCapacityKind::AwsFleet,
351 local: false,
352 probes,
353 local_storage_paths: Vec::new(),
354 probe_error,
355 });
356 }
357 }
358 }
359 if !local_ids.is_empty() {
360 let mut paths = Vec::new();
363 for target_id in &local_ids {
364 if let Some(template) = self.config.targets.get(target_id) {
365 push_unique(&mut paths, local_storage_paths(template));
366 }
367 push_unique(&mut paths, self.session_storage_paths(target_id));
368 }
369 targets.push(DeploymentCapacityTarget {
370 id: "local".into(),
371 host: "local".into(),
372 target_ids: local_ids,
373 kind: DeploymentCapacityKind::Host,
374 local: true,
375 probes: Vec::new(),
376 local_storage_paths: paths,
377 probe_error: None,
378 });
379 }
380 targets.extend(
381 ssh_hosts
382 .into_iter()
383 .map(
384 |(host, (target_ids, connections))| DeploymentCapacityTarget {
385 id: format!("ssh:{host}"),
386 host,
387 target_ids,
388 kind: DeploymentCapacityKind::Host,
389 local: false,
390 probes: connections
391 .iter()
392 .map(|(ssh, paths)| targets::ssh_host_capacity_command(ssh, paths))
393 .collect(),
394 local_storage_paths: Vec::new(),
395 probe_error: None,
396 },
397 ),
398 );
399 targets.sort_by(|left, right| left.id.cmp(&right.id));
400 targets
401 }
402
403 pub fn test_target(&self, target_id: &str, executor: &impl CommandExecutor) -> Result<()> {
405 verify_target(
406 self.configured_target(target_id)?,
407 executor,
408 TargetCheck::BeforeLaunch,
409 )
410 }
411
412 pub fn check_target_readiness(
415 &self,
416 target_id: &str,
417 executor: &impl CommandExecutor,
418 ) -> Result<()> {
419 preflight_target(
420 self.configured_target(target_id)?,
421 executor,
422 TargetCheck::BeforeLaunch,
423 )
424 }
425
426 fn configured_target(&self, target_id: &str) -> Result<&TargetTemplate> {
427 self.config
428 .targets
429 .get(target_id)
430 .with_context(|| format!("unknown target template {target_id:?}"))
431 }
432}
433
434#[derive(Debug, Clone, Copy, PartialEq, Eq)]
436pub(super) enum TargetCheck {
437 BeforeLaunch,
440 Launch,
443}
444
445impl TargetCheck {
446 const fn then_retry(self) -> &'static str {
449 match self {
450 Self::BeforeLaunch => "",
451 Self::Launch => ", then Retry launch",
452 }
453 }
454}
455
456pub(crate) fn ssh_storage_paths(template: &TargetTemplate) -> Vec<String> {
461 use mj_core::targets::storage::{
462 ContainerStorage, DEFAULT_BUILD_CACHE_DIRECTORY, REMOTE_CACHE_DIRECTORY,
463 REMOTE_PROFILES_DIRECTORY, REMOTE_WORKERS_DIRECTORY, TEMPORARY_DIRECTORY,
464 };
465 let mut paths = vec![
466 REMOTE_CACHE_DIRECTORY.to_owned(),
467 DEFAULT_BUILD_CACHE_DIRECTORY.to_owned(),
468 TEMPORARY_DIRECTORY.to_owned(),
469 ];
470 match template {
471 TargetTemplate::SshBare {
472 workspace_prefix, ..
473 } => {
474 paths.push(REMOTE_WORKERS_DIRECTORY.to_owned());
475 paths.push(REMOTE_PROFILES_DIRECTORY.to_owned());
476 paths.push(workspace_prefix.to_string_lossy().into_owned());
477 }
478 TargetTemplate::SshPodman { .. } => paths.push(ContainerStorage::Podman.path().to_owned()),
479 TargetTemplate::SshDocker { .. } => paths.push(ContainerStorage::Docker.path().to_owned()),
480 _ => {}
481 }
482 paths.extend(build_cache_directory(template));
483 paths
484}
485
486fn build_cache_directory(template: &TargetTemplate) -> Option<String> {
488 template
489 .container()?
490 .build_cache
491 .as_ref()?
492 .directory
493 .as_ref()
494 .map(|directory| directory.to_string_lossy().into_owned())
495}
496
497fn local_storage_paths(template: &TargetTemplate) -> Vec<String> {
499 use mj_core::targets::storage::{
500 ContainerStorage, DEFAULT_BUILD_CACHE_DIRECTORY, TEMPORARY_DIRECTORY, local_home_path,
501 };
502 let mut paths = vec![
503 data_dir().to_string_lossy().into_owned(),
504 TEMPORARY_DIRECTORY.to_owned(),
505 local_home_path(DEFAULT_BUILD_CACHE_DIRECTORY),
506 ];
507 match template {
508 TargetTemplate::LocalPodman { .. } => {
509 paths.push(local_home_path(ContainerStorage::Podman.path()));
510 }
511 TargetTemplate::LocalDocker { .. } => {
512 paths.push(ContainerStorage::Docker.path().to_owned());
513 }
514 _ => {}
515 }
516 paths.extend(build_cache_directory(template));
517 paths
518}
519
520fn push_unique(paths: &mut Vec<String>, more: impl IntoIterator<Item = String>) {
521 for path in more {
522 if !path.is_empty() && !paths.contains(&path) {
523 paths.push(path);
524 }
525 }
526}
527
528pub(super) fn preflight_target(
535 template: &TargetTemplate,
536 executor: &impl CommandExecutor,
537 check: TargetCheck,
538) -> Result<()> {
539 match template {
540 TargetTemplate::SshBare { ssh, .. }
541 | TargetTemplate::SshPodman { ssh, .. }
542 | TargetTemplate::SshDocker { ssh, .. } => {
543 verify_ssh_connectivity(&SshTarget::from(ssh), executor)
544 }
545 _ => verify_target(template, executor, check),
546 }
547}
548
549#[derive(Debug, Clone, Copy, PartialEq, Eq)]
552pub(crate) enum LocalEngineReadiness {
553 Ready,
554 NotInstalled,
556 NotReady,
558}
559
560pub(crate) fn local_engine_readiness(
565 kind: &str,
566 executor: &impl CommandExecutor,
567) -> Option<LocalEngineReadiness> {
568 let result = match kind {
569 "local-podman" => targets::verify_local_podman(executor).map(|_| ()),
570 "local-docker" => targets::verify_local_docker(executor).map(|_| ()),
571 "apple-container" => executor
574 .execute(
575 &CommandSpec::new("container", ["system", "status"])
576 .purpose("preflight Apple container runtime")
577 .stage(ProvisionStage::Provisioning),
578 )
579 .and_then(|output| {
580 ensure!(output.status == 0, "container system status failed");
581 Ok(())
582 }),
583 _ => return None,
584 };
585 Some(match result {
586 Ok(()) => LocalEngineReadiness::Ready,
587 Err(error) if is_missing_command(&error) => LocalEngineReadiness::NotInstalled,
588 Err(_) => LocalEngineReadiness::NotReady,
589 })
590}
591
592fn is_missing_command(error: &anyhow::Error) -> bool {
594 error.chain().any(|cause| {
595 cause
596 .downcast_ref::<std::io::Error>()
597 .is_some_and(|io| io.kind() == std::io::ErrorKind::NotFound)
598 })
599}
600
601pub(super) fn verify_target(
603 template: &TargetTemplate,
604 executor: &impl CommandExecutor,
605 check: TargetCheck,
606) -> Result<()> {
607 let then_retry = check.then_retry();
608 match template {
609 TargetTemplate::LocalPodman { .. } => targets::verify_local_podman(executor)
610 .map(|_| ())
611 .map_err(|error| {
612 anyhow::anyhow!(
613 "local Podman is not ready. Fix the problem below{then_retry}: {error:#}"
614 )
615 }),
616 TargetTemplate::LocalDocker { .. } => targets::verify_local_docker(executor)
617 .map(|_| ())
618 .map_err(
619 |error| match error.downcast_ref::<targets::DockerUnavailable>() {
620 Some(problem) => anyhow::anyhow!(
624 "{problem} {}",
625 match check {
626 TargetCheck::BeforeLaunch => problem.remedy(),
627 TargetCheck::Launch => problem.launch_remedy(),
628 }
629 ),
630 None => anyhow::anyhow!(
631 "local Docker is not ready. Start Docker or fix the problem below{then_retry}: {error:#}"
632 ),
633 },
634 ),
635 TargetTemplate::SshPodman { ssh, .. } => {
636 let ssh = SshTarget::from(ssh);
637 targets::verify_ssh_podman(&ssh, executor)
638 .map(|preflight| {
639 for warning in preflight.warnings {
640 executor.notify_notice(&warning.notice());
641 }
642 })
643 .map_err(|error| {
644 anyhow::anyhow!(
645 "remote Podman is not ready on {}. Fix the problem below{then_retry}: {error:#}",
646 ssh.destination
647 )
648 })
649 }
650 TargetTemplate::SshDocker { ssh, .. } => {
651 let ssh = SshTarget::from(ssh);
652 targets::verify_ssh_docker(&ssh, executor)
653 .map(|_| ())
654 .map_err(|error| {
655 anyhow::anyhow!(
656 "remote Docker preflight failed for {}. Fix the problem below{then_retry}: {error:#}",
657 ssh.destination
658 )
659 })
660 }
661 TargetTemplate::AppleContainer { .. } => {
662 let command = CommandSpec::new("container", ["system", "status"])
663 .purpose("preflight Apple container runtime")
664 .stage(ProvisionStage::Provisioning);
665 let output = executor.execute(&command).map_err(|error| {
666 anyhow::anyhow!(
667 "Apple container is not ready. Fix the problem below{then_retry}: {error}"
668 )
669 })?;
670 if output.status != 0 {
671 bail!(
672 "Apple container is not ready. Start the runtime with `container system start`{then_retry}: container system status exited {}: {}",
673 output.status,
674 [
675 String::from_utf8_lossy(&output.stdout).trim(),
676 String::from_utf8_lossy(&output.stderr).trim(),
677 ]
678 .into_iter()
679 .filter(|message| !message.is_empty())
680 .collect::<Vec<_>>()
681 .join("\n")
682 );
683 }
684 Ok(())
685 }
686 TargetTemplate::SshBare { ssh, .. } => {
687 verify_ssh_connectivity(&SshTarget::from(ssh), executor)
688 }
689 TargetTemplate::AwsEc2 {
690 aws_profile,
691 region,
692 launch_template,
693 launch_template_version,
694 ..
695 } => {
696 let mut identity_args = vec!["sts".into(), "get-caller-identity".into()];
697 if let Some(profile) = aws_profile {
698 identity_args.extend(["--profile".into(), profile.clone()]);
699 }
700 let identity = CommandSpec::new("aws", identity_args)
701 .purpose("verify AWS credentials")
702 .stage(ProvisionStage::Provisioning);
703 let output = executor.execute(&identity)?;
704 ensure!(
705 output.status == 0,
706 "AWS credential test failed with status {}: {}",
707 output.status,
708 String::from_utf8_lossy(&output.stderr).trim()
709 );
710
711 let mut launch_args = vec![
712 "ec2".into(),
713 "describe-launch-template-versions".into(),
714 "--region".into(),
715 region.clone(),
716 if launch_template.starts_with("lt-") { "--launch-template-id" } else { "--launch-template-name" }.into(),
717 launch_template.clone(),
718 "--versions".into(),
719 launch_template_version
720 .clone()
721 .unwrap_or_else(|| "$Default".into()),
722 ];
723 if let Some(profile) = aws_profile {
724 launch_args.extend(["--profile".into(), profile.clone()]);
725 }
726 let launch = CommandSpec::new("aws", launch_args)
727 .purpose("verify AWS launch template")
728 .stage(ProvisionStage::Provisioning);
729 let output = executor.execute(&launch)?;
730 ensure!(
731 output.status == 0,
732 "AWS launch-template test failed with status {}: {}",
733 output.status,
734 String::from_utf8_lossy(&output.stderr).trim()
735 );
736 Ok(())
737 }
738 TargetTemplate::LocalBare => Ok(()),
739 }
740}
741
742fn verify_ssh_connectivity(ssh: &SshTarget, executor: &impl CommandExecutor) -> Result<()> {
743 let output = executor.execute(&targets::ssh_connectivity_probe(ssh))?;
744 ensure!(
745 output.status == 0,
746 "SSH connectivity test failed for {} with status {}: {}",
747 ssh.destination,
748 output.status,
749 String::from_utf8_lossy(&output.stderr).trim()
750 );
751 Ok(())
752}
753
754#[cfg(test)]
755pub(super) fn backend_bundle(
756 bundle: &ProjectBundle,
757 executor: &impl CommandExecutor,
758) -> Result<ProjectBundleSpec> {
759 backend_bundle_with_sources(bundle, None, executor)
760}
761
762pub(super) fn backend_session_bundle(
763 session: &SessionRecord,
764 config: &mj_core::config::Config,
765 executor: &impl CommandExecutor,
766) -> Result<ProjectBundleSpec> {
767 backend_bundle_with_sources(
768 session
769 .project_bundle(config)
770 .context("session bundle is missing")?,
771 session
772 .project
773 .as_ref()
774 .map(|project| &project.network_sources),
775 executor,
776 )
777}
778
779fn backend_bundle_with_sources(
780 bundle: &ProjectBundle,
781 sources: Option<&std::collections::BTreeMap<String, mj_core::remote_git::NetworkGitSource>>,
782 executor: &impl CommandExecutor,
783) -> Result<ProjectBundleSpec> {
784 let primary = bundle.primary().context("bundle primary is missing")?;
785 Ok(ProjectBundleSpec {
786 primary: primary.destination.to_string_lossy().into_owned(),
787 repositories: bundle
788 .repositories
789 .iter()
790 .map(|repository| {
791 let source = match sources.and_then(|sources| sources.get(&repository.id)) {
792 Some(source) => source.clone(),
793 None => mj_core::remote_git::resolve_repository(repository, executor)
794 .with_context(|| format!("repository {:?}", repository.id))?,
795 };
796 Ok(RepositorySpec {
797 url: Some(source.fetch_url),
798 push_urls: source.push_urls,
799 destination: repository.destination.to_string_lossy().into_owned(),
800 git_ref: None,
801 reference: None,
802 })
803 })
804 .collect::<Result<Vec<_>>>()?,
805 })
806}
807
808#[derive(Debug, Clone, Copy, Default)]
812pub(super) struct ContainerOverrides<'a> {
813 pub cpus: Option<&'a str>,
814 pub memory: Option<&'a str>,
815}
816
817impl<'a> ContainerOverrides<'a> {
818 pub(super) fn for_session(session: &'a SessionRecord) -> Self {
819 Self {
820 cpus: session.container_cpus.as_deref(),
821 memory: session.container_memory.as_deref(),
822 }
823 }
824}
825
826pub(super) fn backend_target(
827 template: &TargetTemplate,
828 allocation: Option<&SessionResourceAllocation>,
829 overrides: ContainerOverrides<'_>,
830) -> Result<targets::TargetTemplate> {
831 Ok(match template {
832 TargetTemplate::LocalBare => targets::TargetTemplate::LocalBare,
833 TargetTemplate::LocalPodman { container } => {
834 let mut backend = backend_container(container, allocation, overrides);
835 backend.workspace_storage = (&container.workspace_storage).into();
836 targets::TargetTemplate::LocalPodman(backend)
837 }
838 TargetTemplate::LocalDocker { container } => targets::TargetTemplate::LocalDocker(
839 backend_container(container, allocation, overrides),
840 ),
841 TargetTemplate::AppleContainer { container } => targets::TargetTemplate::AppleContainer(
842 backend_container(container, allocation, overrides),
843 ),
844 TargetTemplate::AwsEc2 {
845 aws_profile,
846 region,
847 launch_template,
848 launch_template_version,
849 ssh_user,
850 identity_file,
851 ssh_args,
852 ..
853 } => targets::TargetTemplate::AwsEc2(AwsTemplate {
854 profile: aws_profile.clone().unwrap_or_else(|| "default".into()),
855 region: region.clone(),
856 launch_template: launch_template.clone(),
857 launch_template_version: launch_template_version.clone(),
858 instance_type: match allocation {
859 Some(SessionResourceAllocation::AwsEc2 { instance_type, .. }) => {
860 Some(instance_type.clone())
861 }
862 _ => None,
863 },
864 ssh: SshTarget {
866 destination: format!("{ssh_user}@pending.invalid"),
867 ssh_args: targets::ssh_args_with_identity(ssh_args, identity_file.as_deref()),
868 },
869 }),
870 TargetTemplate::SshBare {
871 ssh,
872 workspace_prefix,
873 ..
874 } => targets::TargetTemplate::SshBare {
875 ssh: SshTarget::from(ssh),
876 workspace_prefix: workspace_prefix.to_string_lossy().into_owned(),
877 },
878 TargetTemplate::SshPodman { ssh, container, .. } => {
879 let mut backend = backend_container(container, allocation, overrides);
880 backend.workspace_storage = (&container.workspace_storage).into();
881 targets::TargetTemplate::SshPodman {
882 ssh: SshTarget::from(ssh),
883 container: backend,
884 }
885 }
886 TargetTemplate::SshDocker { ssh, container, .. } => targets::TargetTemplate::SshDocker {
887 ssh: SshTarget::from(ssh),
888 container: backend_container(container, allocation, overrides),
889 },
890 })
891}
892
893pub fn image_refresh_plan(config: &Config) -> Vec<ImageRefresh> {
905 let mut plan: Vec<ImageRefresh> = Vec::new();
906 for target in config.targets.values() {
907 let Some((host, container)) = target.image_host() else {
908 continue;
909 };
910 let Some(refresh) = targets::image_refresh(
911 host,
912 &container.image,
913 container.platform.as_deref(),
914 container.pull_policy,
915 ) else {
916 continue;
917 };
918 if let Some(existing) = plan.iter_mut().find(|entry| {
921 entry.host == refresh.host
922 && entry.image == refresh.image
923 && entry.platform == refresh.platform
924 }) {
925 existing.when = existing.when.max(refresh.when);
926 continue;
927 }
928 plan.push(refresh);
929 }
930 plan
931}
932
933pub(crate) fn controller_github_token() -> Option<String> {
934 for name in ["GH_TOKEN", "GITHUB_TOKEN"] {
935 if let Ok(token) = std::env::var(name)
936 && let Some(token) = usable_github_token(&token)
937 {
938 return Some(token.to_owned());
939 }
940 }
941 let output = match Command::new("gh")
942 .args(["auth", "token", "--hostname", "github.com"])
943 .stdin(Stdio::null())
944 .stderr(Stdio::null())
945 .output()
946 {
947 Ok(output) => output,
948 Err(error) => {
949 tracing::debug!(%error, "could not query the GitHub CLI for a token");
950 return None;
951 }
952 };
953 if !output.status.success() {
954 tracing::debug!(status = ?output.status, "GitHub CLI did not return an authenticated token");
955 return None;
956 }
957 let token = match std::str::from_utf8(&output.stdout) {
958 Ok(token) => token,
959 Err(error) => {
960 tracing::debug!(%error, "GitHub CLI returned a non-UTF-8 token");
961 return None;
962 }
963 };
964 let Some(token) = usable_github_token(token) else {
965 tracing::debug!("GitHub CLI returned an empty or invalid token");
966 return None;
967 };
968 Some(token.to_owned())
969}
970
971fn usable_github_token(token: &str) -> Option<&str> {
972 let token = token.trim();
973 (!token.is_empty() && !token.chars().any(char::is_whitespace)).then_some(token)
974}
975
976pub(super) fn configure_github_token_environment(target: &mut targets::TargetTemplate) -> bool {
977 let container = match target {
978 targets::TargetTemplate::LocalPodman(container)
979 | targets::TargetTemplate::LocalDocker(container)
980 | targets::TargetTemplate::AppleContainer(container)
981 | targets::TargetTemplate::SshPodman { container, .. }
982 | targets::TargetTemplate::SshDocker { container, .. } => container,
983 targets::TargetTemplate::LocalBare
984 | targets::TargetTemplate::AwsEc2(_)
985 | targets::TargetTemplate::SshBare { .. } => return false,
986 };
987 container
988 .extra_run_args
989 .extend(["--env".to_owned(), "GH_TOKEN".to_owned()]);
990 true
991}
992
993pub(super) fn use_github_https_urls(bundle: &mut targets::ProjectBundleSpec) {
994 for repository in &mut bundle.repositories {
995 for source in repository
996 .url
997 .iter_mut()
998 .chain(repository.push_urls.iter_mut())
999 {
1000 if let Some(github) = crate::setup::github_repository_from_origin(source) {
1001 *source = format!(
1002 "https://github.com/{}/{}.git",
1003 github.owner, github.repository
1004 );
1005 }
1006 }
1007 }
1008}
1009
1010fn backend_container(
1011 container: &mj_core::config::ContainerTemplate,
1012 allocation: Option<&SessionResourceAllocation>,
1013 overrides: ContainerOverrides<'_>,
1014) -> ContainerTemplate {
1015 let mut extra_run_args = Vec::new();
1016 if let Some(platform) = &container.platform {
1017 extra_run_args.push(format!("--platform={platform}"));
1018 }
1019 let (cpus, memory) = match allocation {
1020 Some(SessionResourceAllocation::Container { cpus, memory_bytes }) => {
1021 (Some(cpus.to_string()), Some(memory_bytes.to_string()))
1022 }
1023 _ => (container.cpus.clone(), container.memory.clone()),
1024 };
1025 let cpus = overrides.cpus.map(str::to_owned).or(cpus);
1027 let memory = overrides.memory.map(str::to_owned).or(memory);
1028 if let Some(cpus) = cpus {
1029 extra_run_args.push(format!("--cpus={cpus}"));
1030 }
1031 if let Some(memory) = memory {
1032 extra_run_args.push(format!("--memory={memory}"));
1033 }
1034 for (key, value) in &container.environment {
1035 extra_run_args.extend(["--env".to_string(), format!("{key}={value}")]);
1036 }
1037 ContainerTemplate {
1038 image: container.image.clone(),
1039 pull_policy: container.pull_policy,
1040 extra_run_args,
1041 workspace_storage: targets::PodmanWorkspaceStorage::ContainerLayer,
1042 build_cache: container.build_cache.clone(),
1043 }
1044}
1045
1046pub(crate) fn validate_resource_allocation(
1047 target_kind: impl Into<ResourceAllocationKind>,
1048 allocation: Option<&SessionResourceAllocation>,
1049) -> Result<()> {
1050 if let Some(allocation) = allocation {
1051 allocation.validate()?;
1052 }
1053 match (target_kind.into(), allocation) {
1054 (_, None)
1055 | (ResourceAllocationKind::Container, Some(SessionResourceAllocation::Container { .. }))
1056 | (ResourceAllocationKind::AwsEc2, Some(SessionResourceAllocation::AwsEc2 { .. })) => {
1057 Ok(())
1058 }
1059 (ResourceAllocationKind::Fixed, Some(_)) => {
1060 bail!(mj_core::state::BARE_TARGET_FIXED_RESOURCES)
1061 }
1062 _ => bail!("resource allocation does not match the selected target kind"),
1063 }
1064}
1065
1066const AWS_SSH_READY_TIMEOUT: Duration = Duration::from_secs(300);
1068
1069const AWS_SSH_READY_RETRY_DELAY: Duration = Duration::from_secs(3);
1070
1071fn wait_for_ssh_ready(
1076 executor: &impl CommandExecutor,
1077 probe: &CommandSpec,
1078 timeout: Duration,
1079 mut now: impl FnMut() -> Instant,
1080 mut sleep: impl FnMut(Duration),
1081) -> Result<()> {
1082 let started = now();
1083 loop {
1084 if executor.cancellation_requested() {
1085 bail!("cancelled while waiting for SSH on the new instance");
1086 }
1087 let failure = match executor.execute(probe) {
1088 Ok(output) if output.status == 0 => return Ok(()),
1089 Ok(output) => String::from_utf8_lossy(&output.stderr).trim().to_string(),
1090 Err(error) => error.to_string(),
1091 };
1092 if now().duration_since(started) >= timeout {
1093 bail!(
1094 "{} timed out after {}s: {}",
1095 probe.purpose,
1096 timeout.as_secs(),
1097 if failure.is_empty() {
1098 "the SSH probe reported no error output"
1099 } else {
1100 failure.as_str()
1101 }
1102 );
1103 }
1104 sleep(AWS_SSH_READY_RETRY_DELAY);
1105 }
1106}
1107
1108pub(super) fn locator_after_provision_named(
1109 canonical: &TargetTemplate,
1110 backend: &targets::TargetTemplate,
1111 session_id: &str,
1112 first_output: Option<&CommandOutput>,
1113 executor: &(impl CommandExecutor + Sync),
1114 name: &str,
1115) -> Result<TargetLocator> {
1116 let generated = name.to_owned();
1117
1118 Ok(match canonical {
1119 TargetTemplate::LocalBare => TargetLocator::LocalBare {
1120 worker_root: data_dir().join("workers").join(session_id),
1121 },
1122 TargetTemplate::LocalPodman { .. } => {
1123 let targets::TargetTemplate::LocalPodman(container) = backend else {
1124 bail!("session locator/template mismatch")
1125 };
1126 TargetLocator::LocalPodman {
1127 borrowed_from: None,
1128 container_id: generated,
1129 workspace_storage: PodmanWorkspaceLocator::from(
1130 targets::podman_workspace_locator_named(container, name)?,
1131 ),
1132 }
1133 }
1134 TargetTemplate::LocalDocker { .. } => TargetLocator::LocalDocker {
1135 borrowed_from: None,
1136 container_id: generated,
1137 },
1138 TargetTemplate::AppleContainer { .. } => TargetLocator::AppleContainer {
1139 borrowed_from: None,
1140 container_id: generated,
1141 },
1142 TargetTemplate::SshBare { ssh, .. } => TargetLocator::SshBare {
1143 host: ssh.host.clone(),
1144 workspace: PathBuf::from(targets::workspace_for(backend, session_id)?),
1145 worker_id: None,
1146 },
1147 TargetTemplate::SshPodman { ssh, .. } => {
1148 let targets::TargetTemplate::SshPodman { container, .. } = backend else {
1149 bail!("session locator/template mismatch")
1150 };
1151 TargetLocator::SshPodman {
1152 borrowed_from: None,
1153 host: ssh.host.clone(),
1154 container_id: generated,
1155 workspace_storage: PodmanWorkspaceLocator::from(
1156 targets::podman_workspace_locator_named(container, name)?,
1157 ),
1158 }
1159 }
1160 TargetTemplate::SshDocker { ssh, .. } => TargetLocator::SshDocker {
1161 borrowed_from: None,
1162 host: ssh.host.clone(),
1163 container_id: generated,
1164 },
1165 TargetTemplate::AwsEc2 { .. } => {
1166 let output = first_output.context("AWS launch produced no output")?;
1167 let json: serde_json::Value = serde_json::from_slice(&output.stdout)
1168 .context("parse aws ec2 run-instances response")?;
1169 let instance_id = json
1170 .pointer("/Instances/0/InstanceId")
1171 .and_then(serde_json::Value::as_str)
1172 .context("AWS response omitted instance ID")?
1173 .to_string();
1174 return ec2_locator_after_launch(canonical, instance_id, executor);
1175 }
1176 })
1177}
1178
1179pub(super) fn ec2_locator_after_launch(
1181 canonical: &TargetTemplate,
1182 instance_id: String,
1183 executor: &(impl CommandExecutor + Sync),
1184) -> Result<TargetLocator> {
1185 let TargetTemplate::AwsEc2 {
1186 aws_profile,
1187 region,
1188 ssh_user,
1189 address_source,
1190 identity_file,
1191 ssh_args,
1192 ..
1193 } = canonical
1194 else {
1195 bail!("EC2 locator requires an EC2 target");
1196 };
1197 let profile = aws_profile.clone().unwrap_or_else(|| "default".into());
1198 execute_checked(
1199 executor,
1200 CommandSpec::new(
1201 "aws",
1202 [
1203 "--profile".into(),
1204 profile.clone(),
1205 "--region".into(),
1206 region.clone(),
1207 "ec2".into(),
1208 "wait".into(),
1209 "instance-running".into(),
1210 "--instance-ids".into(),
1211 instance_id.clone(),
1212 ],
1213 )
1214 .purpose("wait for EC2 session instance to run")
1215 .stage(ProvisionStage::Booting),
1216 )?;
1217 let field = match address_source {
1218 AwsAddressSource::PublicDns => "PublicDnsName",
1219 AwsAddressSource::PublicIp => "PublicIpAddress",
1220 AwsAddressSource::PrivateDns => "PrivateDnsName",
1221 AwsAddressSource::PrivateIp => "PrivateIpAddress",
1222 };
1223 let address = execute_checked(
1224 executor,
1225 CommandSpec::new(
1226 "aws",
1227 [
1228 "--profile".into(),
1229 profile.clone(),
1230 "--region".into(),
1231 region.clone(),
1232 "ec2".into(),
1233 "describe-instances".into(),
1234 "--instance-ids".into(),
1235 instance_id.clone(),
1236 "--query".into(),
1237 format!("Reservations[0].Instances[0].{field}"),
1238 "--output".into(),
1239 "text".into(),
1240 ],
1241 )
1242 .purpose("resolve EC2 session address")
1243 .stage(ProvisionStage::Booting),
1244 )?;
1245 let address = String::from_utf8(address.stdout)
1246 .context("AWS address was not UTF-8")?
1247 .trim()
1248 .to_string();
1249 if address.is_empty() || address == "None" {
1250 bail!("AWS instance {instance_id} has no configured address");
1251 }
1252 let ssh = SshTarget {
1253 destination: format!("{ssh_user}@{address}"),
1254 ssh_args: targets::ssh_args_with_identity(ssh_args, identity_file.as_deref()),
1255 };
1256 wait_for_ssh_ready(
1257 executor,
1258 &crate::targets::ssh_command(&ssh, ["true"])
1259 .purpose("wait for EC2 SSH availability")
1260 .stage(ProvisionStage::Booting),
1261 AWS_SSH_READY_TIMEOUT,
1262 Instant::now,
1263 std::thread::sleep,
1264 )?;
1265 Ok(TargetLocator::AwsEc2 {
1266 instance_id,
1267 address: Some(address),
1268 })
1269}
1270
1271pub(crate) fn backend_locator(
1276 locator: &TargetLocator,
1277 session: &SessionRecord,
1278 config: &Config,
1279) -> Result<targets::TargetLocator> {
1280 let runtime = if session.target_runtime.is_some() || targets::locator_needs_connection(locator)
1281 {
1282 Some(session.target_runtime_settings(config)?)
1283 } else {
1284 None
1285 };
1286 Ok(targets::TargetLocator::try_from(targets::RecordedTarget {
1287 locator,
1288 runtime: runtime.as_deref(),
1289 session_id: &session.id,
1290 })?)
1291}
1292
1293#[cfg(test)]
1294mod tests;