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