aya-log-ebpf: zero copy!

Author: tamird

aya-log-common/src/lib.rs

@@ -2,13 +2,11 @@
 
 use core::{
     net::{IpAddr, Ipv4Addr, Ipv6Addr},
-    num::{NonZeroUsize, TryFromIntError},
+    num::TryFromIntError,
 };
 
 use num_enum::IntoPrimitive;
 
-pub const LOG_FIELDS: usize = 6;
-
 pub type LogValueLength = u16;
 
 #[repr(u8)]
@@ -93,8 +91,8 @@ impl UpperMacFormatter for [u8; 6] {}
 
 #[repr(u8)]
 #[derive(Copy, Clone, Debug, IntoPrimitive)]
-pub enum RecordField {
-    Target = 1,
+pub enum RecordFieldKind {
+    Target,
     Level,
     Module,
     File,
@@ -106,7 +104,7 @@ pub enum RecordField {
 /// programs to userspace.
 #[repr(u8)]
 #[derive(Copy, Clone, Debug, IntoPrimitive)]
-pub enum Argument {
+pub enum ArgumentKind {
     DisplayHint,
 
     I8,
@@ -158,172 +156,207 @@ pub enum DisplayHint {
     UpperMac,
 }
 
-// Must be inlined, else the BPF backend emits:
-//
-// llvm: <unknown>:0:0: in function _ZN14aya_log_common5write17hc9ed05433e23a663E { i64, i64 } (i8, ptr, i64, ptr, i64): only integer returns supported
-#[inline(always)]
-pub(crate) fn write(tag: u8, value: &[u8], buf: &mut [u8]) -> Option<NonZeroUsize> {
-    // TODO(https://github.com/rust-lang/rust-clippy/issues/14112): Remove this allowance when the
-    // lint behaves more sensibly.
-    #[expect(clippy::manual_ok_err)]
-    let wire_len: LogValueLength = match value.len().try_into() {
-        Ok(wire_len) => Some(wire_len),
-        Err(TryFromIntError { .. }) => None,
-    }?;
-    let mut size = 0;
-    for slice in [&[tag][..], &wire_len.to_ne_bytes()[..], value] {
-        let buf = buf.get_mut(size..)?;
-        let buf = buf.get_mut(..slice.len())?;
-        buf.copy_from_slice(slice);
-        size += slice.len();
-    }
-    NonZeroUsize::new(size)
-}
-
-pub trait WriteToBuf {
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize>;
+pub trait Argument {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>);
 }
 
-macro_rules! impl_write_to_buf {
-    ($type:ident, $arg_type:expr) => {
-        impl WriteToBuf for $type {
-            // This need not be inlined because the return value is Option<N> where N is
-            // mem::size_of<$type>, which is a compile-time constant.
-            #[inline(never)]
-            fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-                write($arg_type.into(), &self.to_ne_bytes(), buf)
+macro_rules! impl_argument {
+    ($self:ident, $arg_type:expr) => {
+        impl Argument for $self {
+            fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+                ($arg_type, self.to_ne_bytes())
             }
         }
     };
 }
 
-impl_write_to_buf!(i8, Argument::I8);
-impl_write_to_buf!(i16, Argument::I16);
-impl_write_to_buf!(i32, Argument::I32);
-impl_write_to_buf!(i64, Argument::I64);
-impl_write_to_buf!(isize, Argument::Isize);
+impl_argument!(i8, ArgumentKind::I8);
+impl_argument!(i16, ArgumentKind::I16);
+impl_argument!(i32, ArgumentKind::I32);
+impl_argument!(i64, ArgumentKind::I64);
+impl_argument!(isize, ArgumentKind::Isize);
+
+impl_argument!(u8, ArgumentKind::U8);
+impl_argument!(u16, ArgumentKind::U16);
+impl_argument!(u32, ArgumentKind::U32);
+impl_argument!(u64, ArgumentKind::U64);
+impl_argument!(usize, ArgumentKind::Usize);
 
-impl_write_to_buf!(u8, Argument::U8);
-impl_write_to_buf!(u16, Argument::U16);
-impl_write_to_buf!(u32, Argument::U32);
-impl_write_to_buf!(u64, Argument::U64);
-impl_write_to_buf!(usize, Argument::Usize);
+impl_argument!(f32, ArgumentKind::F32);
+impl_argument!(f64, ArgumentKind::F64);
 
-impl_write_to_buf!(f32, Argument::F32);
-impl_write_to_buf!(f64, Argument::F64);
+enum Either<L, R> {
+    Left(L),
+    Right(R),
+}
 
-impl WriteToBuf for IpAddr {
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
+impl<L, R> AsRef<[u8]> for Either<L, R>
+where
+    L: AsRef<[u8]>,
+    R: AsRef<[u8]>,
+{
+    fn as_ref(&self) -> &[u8] {
         match self {
-            IpAddr::V4(ipv4_addr) => write(Argument::Ipv4Addr.into(), &ipv4_addr.octets(), buf),
-            IpAddr::V6(ipv6_addr) => write(Argument::Ipv6Addr.into(), &ipv6_addr.octets(), buf),
+            Either::Left(l) => l.as_ref(),
+            Either::Right(r) => r.as_ref(),
         }
     }
 }
 
-impl WriteToBuf for Ipv4Addr {
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-        write(Argument::Ipv4Addr.into(), &self.octets(), buf)
+impl Argument for IpAddr {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        match self {
+            IpAddr::V4(ipv4_addr) => {
+                let (kind, value) = ipv4_addr.as_argument();
+                (kind, Either::Left(value))
+            }
+            IpAddr::V6(ipv6_addr) => {
+                let (kind, value) = ipv6_addr.as_argument();
+                (kind, Either::Right(value))
+            }
+        }
     }
 }
 
-impl WriteToBuf for [u8; 4] {
-    // This need not be inlined because the return value is Option<N> where N is 16, which is a
-    // compile-time constant.
-    #[inline(never)]
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-        write(Argument::ArrU8Len4.into(), &self, buf)
+impl Argument for Ipv4Addr {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        (ArgumentKind::Ipv4Addr, self.octets())
     }
 }
 
-impl WriteToBuf for Ipv6Addr {
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-        write(Argument::Ipv6Addr.into(), &self.octets(), buf)
+impl Argument for [u8; 4] {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        (ArgumentKind::ArrU8Len4, self)
     }
 }
 
-impl WriteToBuf for [u8; 16] {
-    // This need not be inlined because the return value is Option<N> where N is 16, which is a
-    // compile-time constant.
-    #[inline(never)]
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-        write(Argument::ArrU8Len16.into(), &self, buf)
+impl Argument for Ipv6Addr {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        (ArgumentKind::Ipv6Addr, self.octets())
     }
 }
 
-impl WriteToBuf for [u16; 8] {
-    // This need not be inlined because the return value is Option<N> where N is 16, which is a
-    // compile-time constant.
-    #[inline(never)]
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-        let bytes = unsafe { core::mem::transmute::<[u16; 8], [u8; 16]>(self) };
-        write(Argument::ArrU16Len8.into(), &bytes, buf)
+impl Argument for [u8; 16] {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        (ArgumentKind::ArrU8Len16, self)
     }
 }
 
-impl WriteToBuf for [u8; 6] {
-    // This need not be inlined because the return value is Option<N> where N is 6, which is a
-    // compile-time constant.
-    #[inline(never)]
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-        write(Argument::ArrU8Len6.into(), &self, buf)
+impl Argument for [u16; 8] {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        let bytes = unsafe { core::mem::transmute::<&[u16; 8], &[u8; 16]>(self) };
+        (ArgumentKind::ArrU16Len8, bytes)
     }
 }
 
-impl WriteToBuf for &[u8] {
-    // Must be inlined, else the BPF backend emits:
-    //
-    // llvm: <unknown>:0:0: in function _ZN63_$LT$$RF$$u5b$u8$u5d$$u20$as$u20$aya_log_common..WriteToBuf$GT$5write17h08f30a45f7b9f09dE { i64, i64 } (ptr, i64, ptr, i64): only integer returns supported
-    #[inline(always)]
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-        write(Argument::Bytes.into(), self, buf)
+impl Argument for [u8; 6] {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        (ArgumentKind::ArrU8Len6, self)
     }
 }
 
-impl WriteToBuf for &str {
-    // Must be inlined, else the BPF backend emits:
-    //
-    // llvm: <unknown>:0:0: in function _ZN54_$LT$$RF$str$u20$as$u20$aya_log_common..WriteToBuf$GT$5write17h7e2d1ccaa758e2b5E { i64, i64 } (ptr, i64, ptr, i64): only integer returns supported
-    #[inline(always)]
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-        write(Argument::Str.into(), self.as_bytes(), buf)
+impl Argument for &[u8] {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        (ArgumentKind::Bytes, self)
     }
 }
 
-impl WriteToBuf for DisplayHint {
-    // This need not be inlined because the return value is Option<N> where N is 1, which is a
-    // compile-time constant.
-    #[inline(never)]
-    fn write(self, buf: &mut [u8]) -> Option<NonZeroUsize> {
-        let v: u8 = self.into();
-        write(Argument::DisplayHint.into(), &v.to_ne_bytes(), buf)
+impl Argument for &str {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        (ArgumentKind::Str, self.as_bytes())
+    }
+}
+
+impl Argument for DisplayHint {
+    fn as_argument(&self) -> (ArgumentKind, impl AsRef<[u8]>) {
+        let v: u8 = (*self).into();
+        (ArgumentKind::DisplayHint, v.to_ne_bytes())
+    }
+}
+
+fn wire_len(value: &[u8]) -> Option<[u8; 2]> {
+    match LogValueLength::try_from(value.len()) {
+        Ok(wire_len) => Some(wire_len.to_ne_bytes()),
+        Err(TryFromIntError { .. }) => None,
+    }
+}
+
+#[doc(hidden)]
+pub struct Field<T>([u8; 1], [u8; 2], T);
+
+impl<T: AsRef<[u8]>> Field<T> {
+    pub fn new(kind: impl Into<u8>, value: T) -> Option<Self> {
+        let wire_len = wire_len(value.as_ref())?;
+        Some(Self([kind.into()], wire_len, value))
+    }
+
+    pub fn with_bytes(&self, op: &mut impl FnMut(&[u8]) -> Option<()>) -> Option<()> {
+        let Self(kind, wire_len, value) = self;
+        op(&kind[..])?;
+        op(&wire_len[..])?;
+        op(value.as_ref())?;
+        Some(())
     }
 }
 
 #[doc(hidden)]
-#[inline(always)] // This function takes too many arguments to not be inlined.
-pub fn write_record_header(
-    buf: &mut [u8],
-    target: &str,
-    level: Level,
-    module: &str,
-    file: &str,
-    line: u32,
-    num_args: usize,
-) -> Option<NonZeroUsize> {
-    let level: u8 = level.into();
-    let mut size = 0;
-    for (tag, value) in [
-        (RecordField::Target, target.as_bytes()),
-        (RecordField::Level, &level.to_ne_bytes()),
-        (RecordField::Module, module.as_bytes()),
-        (RecordField::File, file.as_bytes()),
-        (RecordField::Line, &line.to_ne_bytes()),
-        (RecordField::NumArgs, &num_args.to_ne_bytes()),
-    ] {
-        let buf = buf.get_mut(size..)?;
-        let len = write(tag.into(), value, buf)?;
-        size += len.get();
+pub struct Header<'a> {
+    target: Field<&'a [u8]>,
+    level: Field<[u8; 1]>,
+    module: Field<&'a [u8]>,
+    file: Field<&'a [u8]>,
+    line: Field<[u8; 4]>,
+    num_args: Field<[u8; 4]>,
+}
+
+impl<'a> Header<'a> {
+    pub fn new(
+        target: &'a str,
+        level: Level,
+        module: &'a str,
+        file: &'a str,
+        line: u32,
+        num_args: u32,
+    ) -> Option<Self> {
+        let target = target.as_bytes();
+        let level: u8 = level.into();
+        let level = level.to_ne_bytes();
+        let module = module.as_bytes();
+        let file = file.as_bytes();
+        let line = line.to_ne_bytes();
+        let num_args = num_args.to_ne_bytes();
+        let target = Field::new(RecordFieldKind::Target, target)?;
+        let level = Field::new(RecordFieldKind::Level, level)?;
+        let module = Field::new(RecordFieldKind::Module, module)?;
+        let file = Field::new(RecordFieldKind::File, file)?;
+        let line = Field::new(RecordFieldKind::Line, line)?;
+        let num_args = Field::new(RecordFieldKind::NumArgs, num_args)?;
+
+        Some(Self {
+            target,
+            level,
+            module,
+            file,
+            line,
+            num_args,
+        })
+    }
+
+    pub fn with_bytes(&self, op: &mut impl FnMut(&[u8]) -> Option<()>) -> Option<()> {
+        let Self {
+            target,
+            level,
+            module,
+            file,
+            line,
+            num_args,
+        } = self;
+        target.with_bytes(op)?;
+        level.with_bytes(op)?;
+        module.with_bytes(op)?;
+        file.with_bytes(op)?;
+        line.with_bytes(op)?;
+        num_args.with_bytes(op)?;
+        Some(())
     }
-    NonZeroUsize::new(size)
 }