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#[derive(Clone, Debug, Default)]
17pub enum FollowComponents {
18 #[default]
19 None,
21 Named(HashSet<String>),
23 All,
25}
26
27impl FollowComponents {
28 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
38pub 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 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
130pub struct ComponentStdioWriter {
133 inner: ComponentStdioWriterInner,
134}
135
136enum ComponentStdioWriterInner {
137 Inherit,
139 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}