dagger_sdk/core/
cli_session.rs

1use std::{fs::canonicalize, path::Path, process::Stdio, sync::Arc};
2
3use eyre::Context;
4use tokio::{io::AsyncBufReadExt, sync::broadcast};
5
6use crate::core::{config::Config, connect_params::ConnectParams};
7
8#[derive(Clone, Debug)]
9pub struct CliSession {
10    inner: Arc<InnerCliSession>,
11}
12
13impl Default for CliSession {
14    fn default() -> Self {
15        Self::new()
16    }
17}
18
19impl CliSession {
20    pub fn new() -> Self {
21        Self {
22            inner: Arc::new(InnerCliSession {}),
23        }
24    }
25
26    pub async fn connect(
27        &self,
28        config: &Config,
29        cli_path: &Path,
30    ) -> eyre::Result<(ConnectParams, DaggerSessionProc)> {
31        self.inner.connect(config, cli_path).await
32    }
33}
34
35pub struct DaggerSessionProc {
36    shutdown: broadcast::Sender<()>,
37
38    inner: tokio::sync::Mutex<tokio::process::Child>,
39}
40
41impl DaggerSessionProc {
42    pub fn subscribe_shutdown(&self) -> broadcast::Receiver<()> {
43        self.shutdown.subscribe()
44    }
45
46    pub async fn shutdown(&self) -> eyre::Result<()> {
47        let mut proc = self.inner.lock().await;
48
49        tracing::trace!("waiting for dagger subprocess to shutdown");
50
51        tracing::trace!("sending shutdown signal");
52        if let Err(e) = self.shutdown.send(()) {
53            tracing::warn!("failed to send shutdown signal: {}", e);
54        }
55
56        tracing::trace!("closing stdin");
57        proc.wait().await.context("failed to shutdown session")?;
58
59        tracing::trace!("dagger subprocess shutdown");
60
61        Ok(())
62    }
63}
64
65impl From<tokio::process::Child> for DaggerSessionProc {
66    fn from(value: tokio::process::Child) -> Self {
67        let (tx, _) = broadcast::channel::<()>(1);
68
69        Self {
70            inner: tokio::sync::Mutex::new(value),
71            shutdown: tx,
72        }
73    }
74}
75
76#[derive(Debug)]
77struct InnerCliSession {}
78
79impl InnerCliSession {
80    pub async fn connect(
81        &self,
82        config: &Config,
83        cli_path: &Path,
84    ) -> eyre::Result<(ConnectParams, DaggerSessionProc)> {
85        let proc = self.start(config, cli_path)?;
86        let params = self.get_conn(proc, config).await?;
87
88        Ok(params)
89    }
90
91    fn start(&self, config: &Config, cli_path: &Path) -> eyre::Result<tokio::process::Child> {
92        let mut args: Vec<String> = vec!["session".into()];
93        if let Some(workspace) = &config.workdir_path {
94            let abs_path = canonicalize(workspace)?;
95            args.extend(["--workdir".into(), abs_path.to_string_lossy().to_string()])
96        }
97        if let Some(config_path) = &config.config_path {
98            let abs_path = canonicalize(config_path)?;
99            args.extend(["--project".into(), abs_path.to_string_lossy().to_string()])
100        }
101
102        args.extend(["--label".into(), "dagger.io/sdk.name:rust".into()]);
103        args.extend([
104            "--label".into(),
105            format!("dagger.io/sdk.version:{}", env!("CARGO_PKG_VERSION")),
106        ]);
107
108        let proc = tokio::process::Command::new(
109            cli_path
110                .to_str()
111                .ok_or(eyre::anyhow!("could not get string from path"))?,
112        )
113        .args(args.as_slice())
114        .stdin(Stdio::piped())
115        .stdout(Stdio::piped())
116        .stderr(Stdio::piped())
117        .spawn()?;
118
119        //TODO: Add retry mechanism
120
121        Ok(proc)
122    }
123
124    async fn get_conn(
125        &self,
126        mut proc: tokio::process::Child,
127        config: &Config,
128    ) -> eyre::Result<(ConnectParams, DaggerSessionProc)> {
129        let stdout = proc
130            .stdout
131            .take()
132            .ok_or(eyre::anyhow!("could not acquire stdout from child process"))?;
133
134        let stderr = proc
135            .stderr
136            .take()
137            .ok_or(eyre::anyhow!("could not acquire stderr from child process"))?;
138
139        let session: DaggerSessionProc = proc.into();
140
141        let (sender, mut receiver) = tokio::sync::mpsc::channel(1);
142        let logger = config.logger.as_ref().map(|p| p.clone());
143        let mut rx = session.subscribe_shutdown();
144
145        tokio::spawn(async move {
146            let mut stdout_bufr = tokio::io::BufReader::new(stdout).lines();
147            loop {
148                tokio::select! {
149                    line = stdout_bufr.next_line() => {
150                        if let Ok(Some(line)) = line {
151                            if let Ok(conn) = serde_json::from_str::<ConnectParams>(&line) {
152                                sender.send(conn).await.unwrap();
153                                continue;
154                            }
155
156                            if let Some(logger) = &logger {
157                                logger.stdout(&line).unwrap();
158                            }
159                        }
160                    },
161                    _ = rx.recv() => {
162                        drop(stdout_bufr);
163                        tracing::trace!("shutting down stdout");
164                        break;
165                    },
166                };
167            }
168
169            tracing::trace!("closing stdout for dagger session");
170        });
171
172        let mut rx = session.subscribe_shutdown();
173        let logger = config.logger.as_ref().map(|p| p.clone());
174        tokio::spawn(async move {
175            let mut stderr_bufr = tokio::io::BufReader::new(stderr).lines();
176            loop {
177                tokio::select! {
178                    line = stderr_bufr.next_line() => {
179                        if let Ok(Some(line)) = line {
180                            if let Some(logger) = &logger {
181                                logger.stderr(&line).unwrap();
182                            }
183                        }
184                    },
185                    _ = rx.recv() => {
186                        drop(stderr_bufr);
187                        tracing::trace!("shutting down stderr");
188                        break;
189                    },
190                };
191            }
192
193            tracing::trace!("closing stderr for dagger session");
194        });
195
196        let conn = receiver.recv().await.ok_or(eyre::anyhow!(
197            "could not receive ok signal from dagger-engine"
198        ))?;
199
200        Ok((conn, session))
201    }
202}
dagger_sdk/core/cli_session.rs

dagger_sdk/core/
cli_session.rs
