Improve write of arrays

opentelemetry-datadog/src/exporter/intern.rs

@@ -1,5 +1,5 @@
 use indexmap::set::IndexSet;
-use opentelemetry::Value;
+use opentelemetry::{StringValue, Value};
 use rmp::encode::{RmpWrite, ValueWriteError};
 use std::hash::Hash;
 
@@ -18,32 +18,131 @@ impl<'a> Hash for InternValue<'a> {
                 Value::I64(x) => x.hash(state),
                 Value::String(x) => x.hash(state),
                 Value::F64(x) => x.to_bits().hash(state),
-                Value::Array(x) => std::ptr::hash(std::ptr::addr_of!(x), state),
+                Value::Array(a) => match a {
+                    opentelemetry::Array::Bool(x) => x.hash(state),
+                    opentelemetry::Array::I64(x) => x.hash(state),
+                    opentelemetry::Array::F64(floats) => {
+                        for f in floats {
+                            f.to_bits().hash(state);
+                        }
+                    }
+                    opentelemetry::Array::String(x) => x.hash(state),
+                },
             },
         }
     }
 }
 
 impl<'a> Eq for InternValue<'a> {}
 
+const BOOLEAN_TRUE: &str = "true";
+const BOOLEAN_FALSE: &str = "false";
+const LEFT_SQUARE_BRACKET: u8 = 0x5b; // '['
+const RIGHT_SQUARE_BRACKET: u8 = 0x5d; // ']'
+const COMMA: u8 = 0x2c; // ','
+const DOUBLE_QUOTE: u8 = 0x22; // '"'
+const EMPTY_ARRAY: &str = "[]";
+
+trait WriteAsLiteral {
+    fn write_to(&self, buffer: &mut Vec<u8>);
+}
+
+impl WriteAsLiteral for bool {
+    fn write_to(&self, buffer: &mut Vec<u8>) {
+        buffer.extend_from_slice(if *self { BOOLEAN_TRUE } else { BOOLEAN_FALSE }.as_bytes());
+    }
+}
+
+impl WriteAsLiteral for i64 {
+    fn write_to(&self, buffer: &mut Vec<u8>) {
+        buffer.extend_from_slice(itoa::Buffer::new().format(*self).as_bytes());
+    }
+}
+
+impl WriteAsLiteral for f64 {
+    fn write_to(&self, buffer: &mut Vec<u8>) {
+        buffer.extend_from_slice(ryu::Buffer::new().format(*self).as_bytes());
+    }
+}
+
+impl WriteAsLiteral for StringValue {
+    fn write_to(&self, buffer: &mut Vec<u8>) {
+        buffer.push(DOUBLE_QUOTE);
+        buffer.extend_from_slice(self.as_str().as_bytes());
+        buffer.push(DOUBLE_QUOTE);
+    }
+}
+
 impl<'a> InternValue<'a> {
     pub(crate) fn write_as_str<W: RmpWrite>(
         &self,
         payload: &mut W,
+        reusable_buffer: &mut Vec<u8>,
     ) -> Result<(), ValueWriteError<W::Error>> {
         match self {
             InternValue::RegularString(x) => rmp::encode::write_str(payload, x),
             InternValue::OpenTelemetryValue(v) => match v {
                 Value::Bool(x) => {
-                    rmp::encode::write_str(payload, if *x { "true" } else { "false" })
+                    rmp::encode::write_str(payload, if *x { BOOLEAN_TRUE } else { BOOLEAN_FALSE })
                 }
                 Value::I64(x) => rmp::encode::write_str(payload, itoa::Buffer::new().format(*x)),
                 Value::F64(x) => rmp::encode::write_str(payload, ryu::Buffer::new().format(*x)),
                 Value::String(x) => rmp::encode::write_str(payload, x.as_ref()),
-                x => rmp::encode::write_str(payload, x.as_str().as_ref()),
+                Value::Array(array) => match array {
+                    opentelemetry::Array::Bool(x) => {
+                        Self::write_generic_array(payload, reusable_buffer, x)
+                    }
+                    opentelemetry::Array::I64(x) => {
+                        Self::write_generic_array(payload, reusable_buffer, x)
+                    }
+                    opentelemetry::Array::F64(x) => {
+                        Self::write_generic_array(payload, reusable_buffer, x)
+                    }
+                    opentelemetry::Array::String(x) => {
+                        Self::write_generic_array(payload, reusable_buffer, x)
+                    }
+                },
             },
         }
     }
+
+    fn write_empty_array<W: RmpWrite>(payload: &mut W) -> Result<(), ValueWriteError<W::Error>> {
+        rmp::encode::write_str(payload, EMPTY_ARRAY)
+    }
+
+    fn write_buffer_as_string<W: RmpWrite>(
+        payload: &mut W,
+        reusable_buffer: &[u8],
+    ) -> Result<(), ValueWriteError<W::Error>> {
+        rmp::encode::write_str_len(payload, reusable_buffer.len() as u32)?;
+        payload
+            .write_bytes(reusable_buffer)
+            .map_err(ValueWriteError::InvalidDataWrite)
+    }
+
+    fn write_generic_array<W: RmpWrite, T: WriteAsLiteral>(
+        payload: &mut W,
+        reusable_buffer: &mut Vec<u8>,
+        array: &Vec<T>,
+    ) -> Result<(), ValueWriteError<W::Error>> {
+        if array.is_empty() {
+            return Self::write_empty_array(payload);
+        }
+
+        reusable_buffer.clear();
+        reusable_buffer.push(LEFT_SQUARE_BRACKET);
+
+        array[0].write_to(reusable_buffer);
+
+        for value in array[1..].iter() {
+            reusable_buffer.push(COMMA);
+            value.write_to(reusable_buffer);
+        }
+
+        reusable_buffer.push(RIGHT_SQUARE_BRACKET);
+
+        Self::write_buffer_as_string(payload, reusable_buffer)
+    }
 }
 
 pub(crate) struct StringInterner<'a> {
@@ -83,6 +182,20 @@ impl<'a> StringInterner<'a> {
     pub(crate) fn len(&self) -> u32 {
         self.data.len() as u32
     }
+
+    pub(crate) fn write_dictionary<W: RmpWrite>(
+        &self,
+        payload: &mut W,
+    ) -> Result<(), ValueWriteError<W::Error>> {
+        let mut reusable_buffer: Vec<u8> = vec![];
+
+        rmp::encode::write_array_len(payload, self.len())?;
+        for data in self.iter() {
+            data.write_as_str(payload, &mut reusable_buffer)?;
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -108,4 +221,70 @@ mod tests {
         assert_eq!(d_idx, a_idx);
         assert_eq!(e_idx, c_idx);
     }
+
+    #[test]
+    fn test_intern_bool() {
+        let a = Value::Bool(true);
+        let b = Value::Bool(false);
+        let c = "c";
+
+        let mut intern = StringInterner::new();
+        let a_idx = intern.intern_value(&a);
+        let b_idx = intern.intern_value(&b);
+        let c_idx = intern.intern(c);
+        let d_idx = intern.intern_value(&a);
+        let e_idx = intern.intern(c);
+
+        assert_eq!(a_idx, 0);
+        assert_eq!(b_idx, 1);
+        assert_eq!(c_idx, 2);
+        assert_eq!(d_idx, a_idx);
+        assert_eq!(e_idx, c_idx);
+    }
+
+    #[test]
+    fn test_intern_i64() {
+        let a = Value::I64(1234567890);
+        let b = Value::I64(-1234567890);
+        let c = "c";
+        let d = Value::I64(1234567890);
+
+        let mut intern = StringInterner::new();
+        let a_idx = intern.intern_value(&a);
+        let b_idx = intern.intern_value(&b);
+        let c_idx = intern.intern(c);
+        let d_idx = intern.intern_value(&a);
+        let e_idx = intern.intern(c);
+        let f_idx = intern.intern_value(&d);
+
+        assert_eq!(a_idx, 0);
+        assert_eq!(b_idx, 1);
+        assert_eq!(c_idx, 2);
+        assert_eq!(d_idx, a_idx);
+        assert_eq!(e_idx, c_idx);
+        assert_eq!(f_idx, a_idx);
+    }
+
+    #[test]
+    fn test_intern_f64() {
+        let a = Value::F64(123456.7890);
+        let b = Value::F64(-1234567.890);
+        let c = "c";
+        let d = Value::F64(-1234567.890);
+
+        let mut intern = StringInterner::new();
+        let a_idx = intern.intern_value(&a);
+        let b_idx = intern.intern_value(&b);
+        let c_idx = intern.intern(c);
+        let d_idx = intern.intern_value(&a);
+        let e_idx = intern.intern(c);
+        let f_idx = intern.intern_value(&d);
+
+        assert_eq!(a_idx, 0);
+        assert_eq!(b_idx, 1);
+        assert_eq!(c_idx, 2);
+        assert_eq!(d_idx, a_idx);
+        assert_eq!(e_idx, c_idx);
+        assert_eq!(b_idx, f_idx);
+    }
 }

opentelemetry-datadog/src/exporter/model/v05.rs

@@ -92,10 +92,7 @@ where
     let mut payload = Vec::new();
     rmp::encode::write_array_len(&mut payload, 2)?;
 
-    rmp::encode::write_array_len(&mut payload, interner.len())?;
-    for data in interner.iter() {
-        data.write_as_str(&mut payload)?;
-    }
+    interner.write_dictionary(&mut payload)?;
 
     payload.append(&mut encoded_traces);