viceroy_lib/
logging.rs

1use std::{
2    io::{self, Write},
3    pin::Pin,
4    sync::{Arc, Mutex},
5    task::{Context, Poll},
6};
7use tokio::io::AsyncWrite;
8
9/// A named logging endpoint.
10#[derive(Clone)]
11pub struct LogEndpoint {
12    name: Vec<u8>,
13    writer: Arc<Mutex<dyn Write + Send>>,
14}
15
16impl LogEndpoint {
17    /// Allocate a new `LogEndpoint` with the given name, with log messages sent
18    /// to the given writer.
19    pub fn new(name: &[u8], writer: Arc<Mutex<dyn Write + Send>>) -> LogEndpoint {
20        LogEndpoint {
21            name: name.to_owned(),
22            writer,
23        }
24    }
25
26    /// Write a log entry to this endpoint.
27    ///
28    /// Log entries are prefixed with the endpoint name and terminated with a newline.
29    /// Any newlines in the message will be escaped to the string r"\n".
30    ///
31    /// The entry is written atomically to the writer given to [`LogEndpoint::new`].
32    pub fn write_entry(&self, mut msg: &[u8]) -> io::Result<()> {
33        const LOG_ENDPOINT_DELIM: &[u8] = b" :: ";
34
35        // Strip any trailing newlines; we will add a newline at the end,
36        // and escape any interior newlines.
37        if msg.last() == Some(&b'\n') {
38            msg = &msg[..msg.len() - 1];
39        }
40
41        if msg.is_empty() {
42            return Ok(());
43        }
44
45        // Accumulate log entry into a buffer before writing, while escaping newlines
46        let mut to_write =
47            Vec::with_capacity(msg.len() + self.name.len() + LOG_ENDPOINT_DELIM.len() + 1);
48
49        to_write.extend_from_slice(&self.name);
50        to_write.extend_from_slice(LOG_ENDPOINT_DELIM);
51        for &byte in msg {
52            if byte == b'\n' {
53                to_write.extend_from_slice(br"\n");
54            } else {
55                to_write.push(byte);
56            }
57        }
58        to_write.push(b'\n');
59
60        self.writer.lock().unwrap().write_all(&to_write)
61    }
62}
63
64impl Write for LogEndpoint {
65    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
66        self.write_entry(buf)?;
67        Ok(buf.len())
68    }
69
70    fn flush(&mut self) -> io::Result<()> {
71        self.writer.lock().unwrap().flush()
72    }
73}
74
75impl wasmtime_wasi::cli::StdoutStream for LogEndpoint {
76    fn p2_stream(&self) -> Box<dyn wasmtime_wasi::p2::OutputStream> {
77        Box::new(self.clone())
78    }
79
80    fn async_stream(&self) -> Box<dyn AsyncWrite + Send + Sync> {
81        Box::new(self.clone())
82    }
83}
84
85#[wiggle::async_trait]
86impl wasmtime_wasi::p2::Pollable for LogEndpoint {
87    async fn ready(&mut self) {}
88}
89
90impl wasmtime_wasi::cli::IsTerminal for LogEndpoint {
91    fn is_terminal(&self) -> bool {
92        false
93    }
94}
95
96impl wasmtime_wasi::p2::OutputStream for LogEndpoint {
97    fn write(&mut self, bytes: bytes::Bytes) -> wasmtime_wasi::p2::StreamResult<()> {
98        self.write_entry(&bytes)
99            .map_err(|e| wasmtime_wasi::p2::StreamError::LastOperationFailed(anyhow::anyhow!(e)))
100    }
101
102    fn flush(&mut self) -> wasmtime_wasi::p2::StreamResult<()> {
103        <Self as Write>::flush(self)
104            .map_err(|e| wasmtime_wasi::p2::StreamError::LastOperationFailed(anyhow::anyhow!(e)))
105    }
106
107    fn check_write(&mut self) -> wasmtime_wasi::p2::StreamResult<usize> {
108        Ok(1024 * 1024)
109    }
110}
111
112impl AsyncWrite for LogEndpoint {
113    fn poll_write(
114        self: Pin<&mut Self>,
115        _cx: &mut Context<'_>,
116        buf: &[u8],
117    ) -> Poll<Result<usize, std::io::Error>> {
118        self.write_entry(&buf)?;
119        Poll::Ready(Ok(buf.len()))
120    }
121
122    fn poll_flush(
123        mut self: Pin<&mut Self>,
124        _cx: &mut Context<'_>,
125    ) -> Poll<Result<(), std::io::Error>> {
126        Poll::Ready(<Self as Write>::flush(&mut self))
127    }
128
129    fn poll_shutdown(
130        self: Pin<&mut Self>,
131        _cx: &mut Context<'_>,
132    ) -> Poll<Result<(), std::io::Error>> {
133        Poll::Ready(Ok(()))
134    }
135}
viceroy_lib/logging.rs

viceroy_lib/
logging.rs
