Skip to main content

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(component_id, log_path, follow)
75                .with_context(|| format!("Failed to open log file {}", quoted_path(log_path))),
76            None => ComponentStdioWriter::new_inherit(component_id),
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!(
92                        "The following component(s) specified in --follow do not exist in the application:\n{unknown_list}\nThe following components exist:\n{actual_list}"
93                    );
94                    Err(message)
95                }
96            }
97            _ => Ok(()),
98        }
99    }
100
101    fn truncate_log_files(log_dir: &Path) {
102        if let Ok(entries) = log_dir.read_dir() {
103            for entry in entries.flatten() {
104                let path = entry.path();
105                let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
106                    continue;
107                };
108
109                if name.ends_with(&format!("{STDOUT_LOG_FILE_SUFFIX}.txt"))
110                    || name.ends_with(&format!("{STDERR_LOG_FILE_SUFFIX}.txt"))
111                {
112                    _ = std::fs::File::create(path)
113                }
114            }
115        }
116    }
117}
118
119#[async_trait]
120impl<F: RuntimeFactors, U> ExecutorHooks<F, U> for StdioLoggingExecutorHooks {
121    async fn configure_app(
122        &self,
123        configured_app: &spin_factors::ConfiguredApp<F>,
124    ) -> anyhow::Result<()> {
125        self.validate_follows(configured_app.app())?;
126
127        if let Some(dir) = &self.log_dir {
128            // Ensure log dir exists if set
129            std::fs::create_dir_all(dir)
130                .with_context(|| format!("Failed to create log dir {}", quoted_path(dir)))?;
131
132            if self.truncate_log {
133                Self::truncate_log_files(dir);
134            }
135
136            eprintln!("Logging component stdio to {}", quoted_path(dir.join("")))
137        }
138        Ok(())
139    }
140
141    fn prepare_instance(
142        &self,
143        builder: &mut spin_factors_executor::FactorsInstanceBuilder<F, U>,
144    ) -> anyhow::Result<()> {
145        let component_id = builder.app_component().id().to_string();
146        let Some(wasi_builder) = builder.factor_builder::<WasiFactor>() else {
147            return Ok(());
148        };
149        wasi_builder.stdout_pipe(self.component_stdio_writer(
150            &component_id,
151            STDOUT_LOG_FILE_SUFFIX,
152            self.log_dir.as_deref(),
153        )?);
154        wasi_builder.stderr_pipe(self.component_stdio_writer(
155            &component_id,
156            STDERR_LOG_FILE_SUFFIX,
157            self.log_dir.as_deref(),
158        )?);
159        Ok(())
160    }
161}
162
163/// ComponentStdioWriter forwards output to a log file, (optionally) stderr, and (optionally) to a
164/// tracing compatibility layer.
165pub struct ComponentStdioWriter {
166    component_id: String,
167    inner: ComponentStdioWriterInner,
168}
169
170enum ComponentStdioWriterInner {
171    /// Inherit stdout/stderr from the parent process.
172    Inherit,
173    /// Forward stdout/stderr to a file in addition to the inherited stdout/stderr.
174    Forward {
175        sync_file: std::fs::File,
176        async_file: tokio::fs::File,
177        state: ComponentStdioWriterState,
178        follow: bool,
179    },
180}
181
182#[derive(Debug)]
183enum ComponentStdioWriterState {
184    File,
185    Follow(std::ops::Range<usize>),
186}
187
188impl ComponentStdioWriter {
189    fn new_forward(component_id: &str, log_path: &Path, follow: bool) -> anyhow::Result<Self> {
190        let sync_file = std::fs::File::options()
191            .create(true)
192            .append(true)
193            .open(log_path)?;
194
195        let async_file = sync_file
196            .try_clone()
197            .context("could not get async file handle")?
198            .into();
199
200        Ok(Self {
201            component_id: component_id.to_string(),
202            inner: ComponentStdioWriterInner::Forward {
203                sync_file,
204                async_file,
205                state: ComponentStdioWriterState::File,
206                follow,
207            },
208        })
209    }
210
211    fn new_inherit(component_id: &str) -> anyhow::Result<Self> {
212        Ok(Self {
213            component_id: component_id.to_string(),
214            inner: ComponentStdioWriterInner::Inherit,
215        })
216    }
217}
218
219impl AsyncWrite for ComponentStdioWriter {
220    fn poll_write(
221        self: std::pin::Pin<&mut Self>,
222        cx: &mut std::task::Context<'_>,
223        buf: &[u8],
224    ) -> Poll<std::result::Result<usize, std::io::Error>> {
225        let this = self.get_mut();
226
227        loop {
228            match &mut this.inner {
229                ComponentStdioWriterInner::Inherit => {
230                    let written = futures::ready!(
231                        std::pin::Pin::new(&mut tokio::io::stderr()).poll_write(cx, buf)
232                    );
233                    let written = match written {
234                        Ok(w) => w,
235                        Err(e) => return Poll::Ready(Err(e)),
236                    };
237                    return Poll::Ready(Ok(written));
238                }
239                ComponentStdioWriterInner::Forward {
240                    async_file,
241                    state,
242                    follow,
243                    ..
244                } => match &state {
245                    ComponentStdioWriterState::File => {
246                        let written =
247                            futures::ready!(std::pin::Pin::new(async_file).poll_write(cx, buf));
248                        let written = match written {
249                            Ok(w) => w,
250                            Err(e) => return Poll::Ready(Err(e)),
251                        };
252                        if *follow {
253                            *state = ComponentStdioWriterState::Follow(0..written);
254                        } else {
255                            return Poll::Ready(Ok(written));
256                        }
257                    }
258                    ComponentStdioWriterState::Follow(range) => {
259                        let written = futures::ready!(
260                            std::pin::Pin::new(&mut tokio::io::stderr())
261                                .poll_write(cx, &buf[range.clone()])
262                        );
263                        let written = match written {
264                            Ok(w) => w,
265                            Err(e) => return Poll::Ready(Err(e)),
266                        };
267                        if range.start + written >= range.end {
268                            let end = range.end;
269                            *state = ComponentStdioWriterState::File;
270                            return Poll::Ready(Ok(end));
271                        } else {
272                            *state = ComponentStdioWriterState::Follow(
273                                (range.start + written)..range.end,
274                            );
275                        };
276                    }
277                },
278            }
279        }
280    }
281
282    fn poll_flush(
283        self: std::pin::Pin<&mut Self>,
284        cx: &mut std::task::Context<'_>,
285    ) -> Poll<std::result::Result<(), std::io::Error>> {
286        let this = self.get_mut();
287
288        match &mut this.inner {
289            ComponentStdioWriterInner::Inherit => {
290                std::pin::Pin::new(&mut tokio::io::stderr()).poll_flush(cx)
291            }
292            ComponentStdioWriterInner::Forward {
293                async_file, state, ..
294            } => match state {
295                ComponentStdioWriterState::File => std::pin::Pin::new(async_file).poll_flush(cx),
296                ComponentStdioWriterState::Follow(_) => {
297                    std::pin::Pin::new(&mut tokio::io::stderr()).poll_flush(cx)
298                }
299            },
300        }
301    }
302
303    fn poll_shutdown(
304        self: std::pin::Pin<&mut Self>,
305        cx: &mut std::task::Context<'_>,
306    ) -> Poll<std::result::Result<(), std::io::Error>> {
307        let this = self.get_mut();
308
309        match &mut this.inner {
310            ComponentStdioWriterInner::Inherit => {
311                std::pin::Pin::new(&mut tokio::io::stderr()).poll_flush(cx)
312            }
313            ComponentStdioWriterInner::Forward {
314                async_file, state, ..
315            } => match state {
316                ComponentStdioWriterState::File => std::pin::Pin::new(async_file).poll_shutdown(cx),
317                ComponentStdioWriterState::Follow(_) => {
318                    std::pin::Pin::new(&mut tokio::io::stderr()).poll_flush(cx)
319                }
320            },
321        }
322    }
323}
324
325impl std::io::Write for ComponentStdioWriter {
326    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
327        spin_telemetry::logs::handle_app_log(buf, &self.component_id);
328
329        match &mut self.inner {
330            ComponentStdioWriterInner::Inherit => {
331                std::io::stderr().write_all(buf)?;
332                Ok(buf.len())
333            }
334            ComponentStdioWriterInner::Forward {
335                sync_file, follow, ..
336            } => {
337                let written = sync_file.write(buf)?;
338                if *follow {
339                    std::io::stderr().write_all(&buf[..written])?;
340                }
341                Ok(written)
342            }
343        }
344    }
345
346    fn flush(&mut self) -> std::io::Result<()> {
347        match &mut self.inner {
348            ComponentStdioWriterInner::Inherit => std::io::stderr().flush(),
349            ComponentStdioWriterInner::Forward {
350                sync_file, follow, ..
351            } => {
352                sync_file.flush()?;
353                if *follow {
354                    std::io::stderr().flush()?;
355                }
356                Ok(())
357            }
358        }
359    }
360}
361
362fn bullet_list<S: std::fmt::Display>(items: impl IntoIterator<Item = S>) -> String {
363    items
364        .into_iter()
365        .map(|item| format!("  - {item}"))
366        .collect::<Vec<_>>()
367        .join("\n")
368}