Skip to content

Commit

Permalink
Avoid allocation in logger with intermediate buffer
Browse files Browse the repository at this point in the history 
Signed-off-by: Eloi DEMOLIS <[email protected]>
  • Loading branch information
@Wonshtrum @Keksoj
Wonshtrum authored and Keksoj committed Feb 2, 2024
1 parent 38a74b5 commit 79fb602
Showing 1 changed file with 70 additions and 57 deletions.
127 changes: 70 additions & 57 deletions command/src/logging.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@ use std::{
cmp::{self, Ord},
collections::BTreeMap,
env,
fmt::{format, Arguments},
fmt::Arguments,
fs::{File, OpenOptions},
io::{stdout, Stdout, Write},
io::{stdout, Error as IoError, ErrorKind as IoErrorKind, Stdout, Write},
mem::ManuallyDrop,
net::{SocketAddr, TcpStream, ToSocketAddrs, UdpSocket},
path::Path,
Expand All @@ -14,7 +14,7 @@ use std::{

use libc;
use mio::net::UnixDatagram;
use prost::Message;
use prost::{encoding::encoded_len_varint, Message};
use rand::{distributions::Alphanumeric, thread_rng, Rng};
use rusty_ulid::Ulid;
use time::Duration;
Expand Down Expand Up @@ -47,6 +47,7 @@ pub struct InnerLogger {
pub access_backend: Option<LoggerBackend>,
/// how to format the access logs
pub access_log_format: AccessLogFormat,
pub buffer: Vec<u8>,
}

pub struct Logger {
Expand Down Expand Up @@ -81,6 +82,7 @@ impl Default for Logger {
backend: LoggerBackend::Stdout(stdout()),
access_backend: None,
access_log_format: AccessLogFormat::Ascii,
buffer: Vec::with_capacity(4096),
},
tag: "UNINITIALIZED".to_string(),
pid: 0,
Expand Down Expand Up @@ -126,6 +128,26 @@ impl Logger {
}
}

trait LoggerBuffer {
fn fmt<F: FnOnce(&[u8]) -> Result<usize, IoError>>(
&mut self,
args: Arguments,
flush: F,
) -> Result<(), IoError>;
}
impl LoggerBuffer for Vec<u8> {
fn fmt<F: FnOnce(&[u8]) -> Result<usize, IoError>>(
&mut self,
args: Arguments,
flush: F,
) -> Result<(), IoError> {
self.clear();
self.write_fmt(args)?;
flush(self.as_slice())?;
Ok(())
}
}

impl InnerLogger {
pub fn log(&mut self, args: Arguments) {
let io_result = match &mut self.backend {
Expand All @@ -135,12 +157,8 @@ impl InnerLogger {
}
LoggerBackend::Tcp(socket) => socket.write_fmt(args),
LoggerBackend::File(file) => file.write_fmt(args),
//FIXME: should have a buffer to write to instead of allocating a string
LoggerBackend::Unix(socket) => socket.send(format(args).as_bytes()).map(|_| ()),
//FIXME: should have a buffer to write to instead of allocating a string
LoggerBackend::Udp(socket, address) => socket
.send_to(format(args).as_bytes(), *address)
.map(|_| ()),
LoggerBackend::Unix(s) => self.buffer.fmt(args, |bytes| s.send(bytes)),
LoggerBackend::Udp(s, addr) => self.buffer.fmt(args, |bytes| s.send_to(bytes, *addr)),
};

if let Err(e) = io_result {
Expand All @@ -154,54 +172,53 @@ impl InnerLogger {

let io_result = match self.access_log_format {
AccessLogFormat::Binary => {
let bytes = unsafe {
log.into_binary_access_log(now, precise_time, tag)
.encode_to_vec()
};
let binary_log = unsafe { log.into_binary_access_log(now, precise_time, tag) };
let log_length = binary_log.encoded_len();
let total_length = log_length + encoded_len_varint(log_length as u64);
self.buffer.clear();
if self.buffer.capacity() < total_length {
self.buffer.reserve(total_length - self.buffer.capacity());
}

match backend {
LoggerBackend::Stdout(stdout) => {
let _ = stdout.write(&bytes);
return;
if let Err(e) = binary_log.encode_length_delimited(&mut self.buffer) {
Err(IoError::new(IoErrorKind::InvalidData, e))
} else {
let bytes = &self.buffer;
match backend {
LoggerBackend::Stdout(stdout) => {
let _ = stdout.write(bytes);
return;
}
LoggerBackend::Tcp(socket) => socket.write(bytes),
LoggerBackend::File(file) => file.write(bytes),
LoggerBackend::Unix(socket) => socket.send(bytes),
LoggerBackend::Udp(socket, address) => socket.send_to(bytes, *address),
}
LoggerBackend::Tcp(socket) => socket.write(&bytes),
LoggerBackend::File(file) => file.write(&bytes),
LoggerBackend::Unix(socket) => socket.send(&bytes),
LoggerBackend::Udp(socket, address) => socket.send_to(&bytes, *address),
.map(|_| ())
}
.map(|_| ())
}
AccessLogFormat::Ascii => {
match backend {
LoggerBackend::Stdout(stdout) => {
let _ = stdout.write_fmt(format_args!(
"{now} {precise_time} {pid} {tag} {level_tag} {log}"
));
return;
}
LoggerBackend::Tcp(socket) => socket.write_fmt(format_args!(
"{now} {precise_time} {pid} {tag} {level_tag} {log}"
)),
LoggerBackend::File(file) => file.write_fmt(format_args!(
AccessLogFormat::Ascii => match backend {
LoggerBackend::Stdout(stdout) => {
let _ = stdout.write_fmt(format_args!(
"{now} {precise_time} {pid} {tag} {level_tag} {log}"
)),
//FIXME: should have a buffer to write to instead of allocating a string
LoggerBackend::Unix(socket) => socket
.send(
format!("{now} {precise_time} {pid} {tag} {level_tag} {log}")
.as_bytes(),
)
.map(|_| ()),
//FIXME: should have a buffer to write to instead of allocating a string
LoggerBackend::Udp(socket, address) => socket
.send_to(
format!("{now} {precise_time} {pid} {tag} {level_tag} {log}")
.as_bytes(),
*address,
)
.map(|_| ()),
));
return;
}
}
LoggerBackend::Tcp(socket) => socket.write_fmt(format_args!(
"{now} {precise_time} {pid} {tag} {level_tag} {log}"
)),
LoggerBackend::File(file) => file.write_fmt(format_args!(
"{now} {precise_time} {pid} {tag} {level_tag} {log}"
)),
LoggerBackend::Unix(socket) => self.buffer.fmt(
format_args!("{now} {precise_time} {pid} {tag} {level_tag} {log}"),
|bytes| socket.send(bytes),
),
LoggerBackend::Udp(socket, address) => self.buffer.fmt(
format_args!("{now} {precise_time} {pid} {tag} {level_tag} {log}"),
|bytes| socket.send_to(bytes, *address),
),
},
};

if let Err(e) = io_result {
Expand All @@ -217,12 +234,8 @@ impl InnerLogger {
}
LoggerBackend::Tcp(socket) => socket.write_fmt(args),
LoggerBackend::File(file) => file.write_fmt(args),
//FIXME: should have a buffer to write to instead of allocating a string
LoggerBackend::Unix(socket) => socket.send(format(args).as_bytes()).map(|_| ()),
//FIXME: should have a buffer to write to instead of allocating a string
LoggerBackend::Udp(socket, address) => socket
.send_to(format(args).as_bytes(), *address)
.map(|_| ()),
LoggerBackend::Unix(s) => self.buffer.fmt(args, |bytes| s.send(bytes)),
LoggerBackend::Udp(s, addr) => self.buffer.fmt(args, |bytes| s.send_to(bytes, *addr)),
};

if let Err(e) = io_result {
Expand Down

0 comments on commit 79fb602

Please sign in to comment.