spin_trigger/cli/
stdio.rs

1use std::{
2    collections::HashSet,
3    path::{Path, PathBuf},
4    task::Poll,
5};
6
7use anyhow::{Context, Result};
8use spin_common::ui::quoted_path;
9use spin_core::async_trait;
10use spin_factor_wasi::WasiFactor;
11use spin_factors::RuntimeFactors;
12use spin_factors_executor::ExecutorHooks;
13use tokio::io::AsyncWrite;
14
15/// Which components should have their logs followed on stdout/stderr.
16#[derive(Clone, Debug, Default)]
17pub enum FollowComponents {
18    #[default]
19    /// No components should have their logs followed.
20    None,
21    /// Only the specified components should have their logs followed.
22    Named(HashSet<String>),
23    /// All components should have their logs followed.
24    All,
25}
26
27impl FollowComponents {
28    /// Whether a given component should have its logs followed on stdout/stderr.
29    pub fn should_follow(&self, component_id: &str) -> bool {
30        match self {
31            Self::None => false,
32            Self::All => true,
33            Self::Named(ids) => ids.contains(component_id),
34        }
35    }
36}
37
38/// Implements TriggerHooks, writing logs to a log file and (optionally) stderr
39pub struct StdioLoggingExecutorHooks {
40    follow_components: FollowComponents,
41    log_dir: Option<PathBuf>,
42}
43
44impl StdioLoggingExecutorHooks {
45    pub fn new(follow_components: FollowComponents, log_dir: Option<PathBuf>) -> Self {
46        Self {
47            follow_components,
48            log_dir,
49        }
50    }
51
52    fn component_stdio_writer(
53        &self,
54        component_id: &str,
55        log_suffix: &str,
56        log_dir: Option<&Path>,
57    ) -> Result<ComponentStdioWriter> {
58        let sanitized_component_id = sanitize_filename::sanitize(component_id);
59        let log_path = log_dir
60            .map(|log_dir| log_dir.join(format!("{sanitized_component_id}_{log_suffix}.txt",)));
61        let log_path = log_path.as_deref();
62
63        let follow = self.follow_components.should_follow(component_id);
64        match log_path {
65            Some(log_path) => ComponentStdioWriter::new_forward(log_path, follow)
66                .with_context(|| format!("Failed to open log file {}", quoted_path(log_path))),
67            None => ComponentStdioWriter::new_inherit(),
68        }
69    }
70
71    fn validate_follows(&self, app: &spin_app::App) -> anyhow::Result<()> {
72        match &self.follow_components {
73            FollowComponents::Named(names) => {
74                let component_ids: HashSet<_> =
75                    app.components().map(|c| c.id().to_owned()).collect();
76                let unknown_names: Vec<_> = names.difference(&component_ids).collect();
77                if unknown_names.is_empty() {
78                    Ok(())
79                } else {
80                    let unknown_list = bullet_list(&unknown_names);
81                    let actual_list = bullet_list(&component_ids);
82                    let message = anyhow::anyhow!("The following component(s) specified in --follow do not exist in the application:\n{unknown_list}\nThe following components exist:\n{actual_list}");
83                    Err(message)
84                }
85            }
86            _ => Ok(()),
87        }
88    }
89}
90
91#[async_trait]
92impl<F: RuntimeFactors, U> ExecutorHooks<F, U> for StdioLoggingExecutorHooks {
93    async fn configure_app(
94        &self,
95        configured_app: &spin_factors::ConfiguredApp<F>,
96    ) -> anyhow::Result<()> {
97        self.validate_follows(configured_app.app())?;
98        if let Some(dir) = &self.log_dir {
99            // Ensure log dir exists if set
100            std::fs::create_dir_all(dir)
101                .with_context(|| format!("Failed to create log dir {}", quoted_path(dir)))?;
102
103            println!("Logging component stdio to {}", quoted_path(dir.join("")))
104        }
105        Ok(())
106    }
107
108    fn prepare_instance(
109        &self,
110        builder: &mut spin_factors_executor::FactorsInstanceBuilder<F, U>,
111    ) -> anyhow::Result<()> {
112        let component_id = builder.app_component().id().to_string();
113        let Some(wasi_builder) = builder.factor_builder::<WasiFactor>() else {
114            return Ok(());
115        };
116        wasi_builder.stdout_pipe(self.component_stdio_writer(
117            &component_id,
118            "stdout",
119            self.log_dir.as_deref(),
120        )?);
121        wasi_builder.stderr_pipe(self.component_stdio_writer(
122            &component_id,
123            "stderr",
124            self.log_dir.as_deref(),
125        )?);
126        Ok(())
127    }
128}
129
130/// ComponentStdioWriter forwards output to a log file, (optionally) stderr, and (optionally) to a
131/// tracing compatibility layer.
132pub struct ComponentStdioWriter {
133    inner: ComponentStdioWriterInner,
134}
135
136enum ComponentStdioWriterInner {
137    /// Inherit stdout/stderr from the parent process.
138    Inherit,
139    /// Forward stdout/stderr to a file in addition to the inherited stdout/stderr.
140    Forward {
141        sync_file: std::fs::File,
142        async_file: tokio::fs::File,
143        state: ComponentStdioWriterState,
144        follow: bool,
145    },
146}
147
148#[derive(Debug)]
149enum ComponentStdioWriterState {
150    File,
151    Follow(std::ops::Range<usize>),
152}
153
154impl ComponentStdioWriter {
155    fn new_forward(log_path: &Path, follow: bool) -> anyhow::Result<Self> {
156        let sync_file = std::fs::File::options()
157            .create(true)
158            .append(true)
159            .open(log_path)?;
160
161        let async_file = sync_file
162            .try_clone()
163            .context("could not get async file handle")?
164            .into();
165
166        Ok(Self {
167            inner: ComponentStdioWriterInner::Forward {
168                sync_file,
169                async_file,
170                state: ComponentStdioWriterState::File,
171                follow,
172            },
173        })
174    }
175
176    fn new_inherit() -> anyhow::Result<Self> {
177        Ok(Self {
178            inner: ComponentStdioWriterInner::Inherit,
179        })
180    }
181}
182
183impl AsyncWrite for ComponentStdioWriter {
184    fn poll_write(
185        self: std::pin::Pin<&mut Self>,
186        cx: &mut std::task::Context<'_>,
187        buf: &[u8],
188    ) -> Poll<std::result::Result<usize, std::io::Error>> {
189        let this = self.get_mut();
190
191        loop {
192            match &mut this.inner {
193                ComponentStdioWriterInner::Inherit => {
194                    let written = futures::ready!(
195                        std::pin::Pin::new(&mut tokio::io::stderr()).poll_write(cx, buf)
196                    );
197                    let written = match written {
198                        Ok(w) => w,
199                        Err(e) => return Poll::Ready(Err(e)),
200                    };
201                    return Poll::Ready(Ok(written));
202                }
203                ComponentStdioWriterInner::Forward {
204                    async_file,
205                    state,
206                    follow,
207                    ..
208                } => match &state {
209                    ComponentStdioWriterState::File => {
210                        let written =
211                            futures::ready!(std::pin::Pin::new(async_file).poll_write(cx, buf));
212                        let written = match written {
213                            Ok(w) => w,
214                            Err(e) => return Poll::Ready(Err(e)),
215                        };
216                        if *follow {
217                            *state = ComponentStdioWriterState::Follow(0..written);
218                        } else {
219                            return Poll::Ready(Ok(written));
220                        }
221                    }
222                    ComponentStdioWriterState::Follow(range) => {
223                        let written = futures::ready!(std::pin::Pin::new(&mut tokio::io::stderr())
224                            .poll_write(cx, &buf[range.clone()]));
225                        let written = match written {
226                            Ok(w) => w,
227                            Err(e) => return Poll::Ready(Err(e)),
228                        };
229                        if range.start + written >= range.end {
230                            let end = range.end;
231                            *state = ComponentStdioWriterState::File;
232                            return Poll::Ready(Ok(end));
233                        } else {
234                            *state = ComponentStdioWriterState::Follow(
235                                (range.start + written)..range.end,
236                            );
237                        };
238                    }
239                },
240            }
241        }
242    }
243
244    fn poll_flush(
245        self: std::pin::Pin<&mut Self>,
246        cx: &mut std::task::Context<'_>,
247    ) -> Poll<std::result::Result<(), std::io::Error>> {
248        let this = self.get_mut();
249
250        match &mut this.inner {
251            ComponentStdioWriterInner::Inherit => {
252                std::pin::Pin::new(&mut tokio::io::stderr()).poll_flush(cx)
253            }
254            ComponentStdioWriterInner::Forward {
255                async_file, state, ..
256            } => match state {
257                ComponentStdioWriterState::File => std::pin::Pin::new(async_file).poll_flush(cx),
258                ComponentStdioWriterState::Follow(_) => {
259                    std::pin::Pin::new(&mut tokio::io::stderr()).poll_flush(cx)
260                }
261            },
262        }
263    }
264
265    fn poll_shutdown(
266        self: std::pin::Pin<&mut Self>,
267        cx: &mut std::task::Context<'_>,
268    ) -> Poll<std::result::Result<(), std::io::Error>> {
269        let this = self.get_mut();
270
271        match &mut this.inner {
272            ComponentStdioWriterInner::Inherit => {
273                std::pin::Pin::new(&mut tokio::io::stderr()).poll_flush(cx)
274            }
275            ComponentStdioWriterInner::Forward {
276                async_file, state, ..
277            } => match state {
278                ComponentStdioWriterState::File => std::pin::Pin::new(async_file).poll_shutdown(cx),
279                ComponentStdioWriterState::Follow(_) => {
280                    std::pin::Pin::new(&mut tokio::io::stderr()).poll_flush(cx)
281                }
282            },
283        }
284    }
285}
286
287impl std::io::Write for ComponentStdioWriter {
288    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
289        spin_telemetry::logs::handle_app_log(buf);
290
291        match &mut self.inner {
292            ComponentStdioWriterInner::Inherit => {
293                std::io::stderr().write_all(buf)?;
294                Ok(buf.len())
295            }
296            ComponentStdioWriterInner::Forward {
297                sync_file, follow, ..
298            } => {
299                let written = sync_file.write(buf)?;
300                if *follow {
301                    std::io::stderr().write_all(&buf[..written])?;
302                }
303                Ok(written)
304            }
305        }
306    }
307
308    fn flush(&mut self) -> std::io::Result<()> {
309        match &mut self.inner {
310            ComponentStdioWriterInner::Inherit => std::io::stderr().flush(),
311            ComponentStdioWriterInner::Forward {
312                sync_file, follow, ..
313            } => {
314                sync_file.flush()?;
315                if *follow {
316                    std::io::stderr().flush()?;
317                }
318                Ok(())
319            }
320        }
321    }
322}
323
324fn bullet_list<S: std::fmt::Display>(items: impl IntoIterator<Item = S>) -> String {
325    items
326        .into_iter()
327        .map(|item| format!("  - {item}"))
328        .collect::<Vec<_>>()
329        .join("\n")
330}
spin_trigger/cli/stdio.rs

spin_trigger/cli/
stdio.rs
