perf(native): reduce mem copy in tmq raw_data_t

Author: zitsen

taos-optin/src/raw/mod.rs

@@ -1774,17 +1774,24 @@ impl RawRes {
     }
 
     #[inline]
-    pub(crate) fn tmq_get_raw(&self) -> RawData {
+    pub(crate) fn tmq_get_raw(&self) -> Result<RawData, RawError> {
         let mut meta = raw_data_t {
             raw: std::ptr::null_mut(),
             raw_len: 0,
             raw_type: 0,
         };
         unsafe {
             let _code = (self.c.tmq.as_ref().unwrap().tmq_get_raw)(self.as_ptr(), &mut meta as _);
-            let raw = RawData::from(&meta);
-            (self.c.tmq.as_ref().unwrap().tmq_free_raw)(meta);
-            raw
+            if _code == 0 {
+                return Ok(RawData::new(
+                    meta,
+                    self.c.tmq.as_ref().unwrap().tmq_free_raw,
+                ));
+            } else {
+                let c_str = (self.c.tmq.as_ref().unwrap().tmq_err2str)(tmq_resp_err_t(_code));
+                let err = CStr::from_ptr(c_str).to_string_lossy().into_owned();
+                Err(RawError::new_with_context(_code, err, "tmq_get_raw error"))
+            }
         }
     }
 }

taos-optin/src/tmq/mod.rs

@@ -317,7 +317,8 @@ impl AsyncOnSync for Meta {}
 
 impl IsMeta for Meta {
     fn as_raw_meta(&self) -> RawResult<RawMeta> {
-        Ok(unsafe { std::mem::transmute(self.raw.clone()) })
+        self.res.tmq_get_raw()
+        // Ok(unsafe { std::mem::transmute(self.raw.clone()) })
     }
 
     fn as_json_meta(&self) -> RawResult<taos_query::common::JsonMeta> {
@@ -328,13 +329,13 @@ impl IsMeta for Meta {
 }
 impl Meta {
     fn new(res: RawRes) -> Self {
-        let raw = res.tmq_get_raw();
+        let raw = res.tmq_get_raw().expect("get raw message error");
         Self { res, raw }
     }
 
-    pub fn to_raw(&self) -> raw_data_t {
-        self.raw.as_raw_data_t()
-    }
+    // pub fn to_raw(&self) -> raw_data_t {
+    //     self.raw.as_raw_data_t()
+    // }
 
     pub fn to_json(&self) -> serde_json::Value {
         serde_json::from_slice(self.res.tmq_get_json_meta().as_bytes())
@@ -359,7 +360,7 @@ impl Data {
 #[async_trait::async_trait]
 impl IsAsyncData for Data {
     async fn as_raw_data(&self) -> RawResult<taos_query::common::RawData> {
-        Ok(self.raw.tmq_get_raw())
+        self.raw.tmq_get_raw()
     }
 
     async fn fetch_raw_block(&self) -> RawResult<Option<RawBlock>> {
@@ -369,7 +370,7 @@ impl IsAsyncData for Data {
 
 impl IsData for Data {
     fn as_raw_data(&self) -> RawResult<taos_query::common::RawData> {
-        Ok(self.raw.tmq_get_raw())
+        self.raw.tmq_get_raw()
     }
 
     fn fetch_raw_block(&self) -> RawResult<Option<RawBlock>> {

taos-query/src/common/raw/data.rs

@@ -11,7 +11,8 @@ const RAW_PTR_OFFSET: usize = std::mem::size_of::<u32>() + std::mem::size_of::<u
 ///
 /// It can be copy/cloned, but should not use it outbound away a offset lifetime.
 #[repr(C)]
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
+#[allow(non_camel_case_types)]
 pub struct raw_data_t {
     pub raw: *const c_void,
     pub raw_len: u32,
@@ -20,98 +21,141 @@ pub struct raw_data_t {
 
 unsafe impl Send for raw_data_t {}
 
-impl raw_data_t {
-    pub fn to_bytes(&self) -> Bytes {
-        let cap = // raw data len
-            self.raw_len as usize +
-            // self.raw_len
-            std::mem::size_of::<u32>() +
-            // self.raw_type
-            std::mem::size_of::<u16>();
-        let mut data = Vec::with_capacity(cap);
-
-        // first 4 bytes: raw_len
-        data.extend(self.raw_len.to_le_bytes());
-
-        // next 2 bytes: raw_type
-        data.extend(self.raw_type.to_le_bytes());
-
-        unsafe {
-            let ptr = data.as_mut_ptr().add(RAW_PTR_OFFSET);
-            std::ptr::copy_nonoverlapping(self.raw, ptr as _, self.raw_len as _);
-            data.set_len(cap);
-        }
-        Bytes::from(data)
-    }
+// impl raw_data_t {
+//     pub fn to_bytes(&self) -> Bytes {
+//         let cap = // raw data len
+//             self.raw_len as usize +
+//             // self.raw_len
+//             std::mem::size_of::<u32>() +
+//             // self.raw_type
+//             std::mem::size_of::<u16>();
+//         let mut data = Vec::with_capacity(cap);
+
+//         // first 4 bytes: raw_len
+//         data.extend(self.raw_len.to_le_bytes());
+
+//         // next 2 bytes: raw_type
+//         data.extend(self.raw_type.to_le_bytes());
+
+//         unsafe {
+//             let ptr = data.as_mut_ptr().add(RAW_PTR_OFFSET);
+//             std::ptr::copy_nonoverlapping(self.raw, ptr as _, self.raw_len as _);
+//             data.set_len(cap);
+//         }
+//         Bytes::from(data)
+//     }
+// }
+
+/// TMQ message raw data container.
+///
+/// It's a wrapper for raw data from native library, and will be auto free when drop.
+#[derive(Debug)]
+pub struct RawData {
+    free: unsafe extern "C" fn(raw: raw_data_t) -> i32,
+    raw: raw_data_t,
 }
-
-#[derive(Debug, Clone)]
-pub struct RawData(Bytes);
-
 unsafe impl Send for RawData {}
 unsafe impl Sync for RawData {}
 
-impl From<&raw_data_t> for RawData {
-    fn from(raw: &raw_data_t) -> Self {
-        RawData(raw.to_bytes())
+impl Drop for RawData {
+    /// Use native free function to free raw_data_t
+    fn drop(&mut self) {
+        unsafe {
+            (self.free)(self.raw);
+        }
     }
 }
 
-impl<T: Into<Bytes>> From<T> for RawData {
-    fn from(bytes: T) -> Self {
-        RawData(bytes.into())
+impl RawData {
+    pub fn new(raw: raw_data_t, free: unsafe extern "C" fn(raw: raw_data_t) -> i32) -> Self {
+        RawData { free, raw }
     }
 }
 
+// #[derive(Debug, Clone)]
+// pub struct RawData(Bytes);
+
+// unsafe impl Send for RawData {}
+// unsafe impl Sync for RawData {}
+
+// // impl From<&raw_data_t> for RawData {
+// //     fn from(raw: &raw_data_t) -> Self {
+// //         RawData(raw.to_bytes())
+// //     }
+// // }
+
+// impl<T: Into<Bytes>> From<T> for RawData {
+//     fn from(bytes: T) -> Self {
+//         RawData(bytes.into())
+//     }
+// }
+
 impl RawData {
-    pub fn new(raw: Bytes) -> Self {
-        raw.into()
-    }
-    pub fn raw(&self) -> *const c_void {
-        unsafe { self.0.as_ptr().add(RAW_PTR_OFFSET) as _ }
+    pub fn raw_ptr(&self) -> *const c_void {
+        self.raw.raw
     }
     pub fn raw_len(&self) -> u32 {
-        unsafe { *(self.0.as_ptr() as *const u32) }
+        self.raw.raw_len
     }
     pub fn raw_type(&self) -> u16 {
-        unsafe { *(self.0.as_ptr().add(std::mem::size_of::<u32>()) as *const u16) }
+        self.raw.raw_type
+    }
+    pub fn raw_slice(&self) -> &[u8] {
+        unsafe { std::slice::from_raw_parts(self.raw.raw as *const u8, self.raw.raw_len as _) }
     }
 
     pub fn as_raw_data_t(&self) -> raw_data_t {
-        raw_data_t {
-            raw: self.raw(),
-            raw_len: self.raw_len(),
-            raw_type: self.raw_type(),
-        }
+        self.raw
     }
 
     pub fn as_bytes(&self) -> Cow<Bytes> {
-        Cow::Borrowed(&self.0)
+        todo!("replace as_bytes")
     }
 }
 
+extern "C" fn _rust_free_raw(raw: raw_data_t) -> i32 {
+    unsafe {
+        let ptr = raw.raw as *mut u8;
+        let len = raw.raw_len as usize;
+        std::alloc::dealloc(
+            ptr,
+            std::alloc::Layout::from_size_align(len, 1).expect("Invalid layout"),
+        );
+    }
+    0
+}
 impl Inlinable for RawData {
     fn read_inlined<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
-        let mut data = Vec::new();
-
         let len = reader.read_u32()?;
-        data.extend(len.to_le_bytes());
-
         let meta_type = reader.read_u16()?;
-        data.extend(meta_type.to_le_bytes());
 
-        data.resize(data.len() + len as usize, 0);
+        let layout = std::alloc::Layout::from_size_align(len as _, 1).map_err(|_| {
+            std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid raw data length")
+        })?;
+        let ptr = unsafe { std::alloc::alloc(layout) };
+        let buf = unsafe { std::slice::from_raw_parts_mut(ptr, len as _) };
 
-        let buf = &mut data[RAW_PTR_OFFSET..];
+        reader.read_exact(buf).inspect_err(|_| unsafe {
+            // free memory if read failed
+            std::alloc::dealloc(ptr, layout);
+        })?;
 
-        reader.read_exact(buf)?;
-        Ok(data.into())
+        let raw = raw_data_t {
+            raw: ptr as _,
+            raw_len: len,
+            raw_type: meta_type,
+        };
+
+        let message = RawData::new(raw, _rust_free_raw);
+
+        Ok(message)
     }
 
     fn write_inlined<W: std::io::Write>(&self, wtr: &mut W) -> std::io::Result<usize> {
-        let bytes = self.as_bytes();
-        wtr.write_all(&bytes)?;
-        Ok(bytes.len())
+        wtr.write_all(self.raw_len().to_le_bytes().as_ref())?;
+        wtr.write_all(self.raw_type().to_le_bytes().as_ref())?;
+        wtr.write_all(self.raw_slice())?;
+        Ok(self.raw_len() as usize + 6)
     }
 }
 
@@ -121,29 +165,31 @@ impl crate::util::AsyncInlinable for RawData {
         reader: &mut R,
     ) -> std::io::Result<Self> {
         use tokio::io::*;
-        let mut data = Vec::new();
-
         let len = reader.read_u32_le().await?;
-        data.extend(len.to_le_bytes());
-
         let meta_type = reader.read_u16_le().await?;
-        data.extend(meta_type.to_le_bytes());
-
-        data.resize(data.len() + len as usize, 0);
-
-        let buf = &mut data[RAW_PTR_OFFSET..];
-
-        reader.read_exact(buf).await?;
-        Ok(data.into())
+        let mut vec: Vec<u8> = Vec::with_capacity(len as usize);
+        reader.read_exact(&mut vec).await?;
+
+        let ptr = Box::into_raw(vec.into_boxed_slice());
+        let raw = raw_data_t {
+            raw: ptr as _,
+            raw_len: len,
+            raw_type: meta_type,
+        };
+
+        let message = RawData::new(raw, _rust_free_raw);
+        Ok(message)
     }
 
     async fn write_inlined<W: tokio::io::AsyncWrite + Send + Unpin>(
         &self,
         wtr: &mut W,
     ) -> std::io::Result<usize> {
         use tokio::io::*;
-        let bytes = self.as_bytes();
-        wtr.write_all(&bytes).await?;
-        Ok(bytes.len())
+        wtr.write_all(self.raw_len().to_le_bytes().as_ref()).await?;
+        wtr.write_all(self.raw_type().to_le_bytes().as_ref())
+            .await?;
+        wtr.write_all(self.raw_slice()).await?;
+        Ok(self.raw_len() as usize + 6)
     }
 }