Merge pull request #40 from marshallpierce/deser-varint

Deserialize optimizations

Author: jonhoo

src/lib.rs

@@ -155,6 +155,7 @@ extern crate num;
 use std::borrow::Borrow;
 use std::cmp;
 use std::ops::{Index, IndexMut, AddAssign, SubAssign};
+use num::ToPrimitive;
 
 use iterators::HistogramIterator;
 
@@ -736,9 +737,9 @@ impl<T: Counter> Histogram<T> {
         };
 
         // determine exponent range needed to support the trackable value with no overflow:
-        let len = h.cover(high);
-
-        h.alloc(len as usize);
+        // TODO usize safety
+        let len = h.cover(high).to_usize().unwrap();
+        h.counts.resize(len, T::zero());
         Ok(h)
     }
 
@@ -748,20 +749,15 @@ impl<T: Counter> Histogram<T> {
         let mut h = Self::new_with_bounds(source.lowest_discernible_value,
                                           source.highest_trackable_value,
                                           source.significant_value_digits)
-            .unwrap();
+            .expect("Using another histogram's parameters failed");
 
         // h.start_time = source.start_time;
         // h.end_time = source.end_time;
         h.auto_resize = source.auto_resize;
-        h.alloc(source.len());
+        h.counts.resize(source.len(), T::zero());
         h
     }
 
-    /// Allocate a counts array of the given size.
-    fn alloc(&mut self, len: usize) {
-        self.counts = std::iter::repeat(T::zero()).take(len).collect();
-    }
-
     // ********************************************************************************************
     // Recording samples.
     // ********************************************************************************************

src/serialization/benchmarks.rs

@@ -2,38 +2,64 @@ extern crate rand;
 extern crate test;
 
 use super::v2_serializer::varint_write;
-use super::deserializer::varint_read;
+use super::deserializer::{varint_read, varint_read_slice};
 use std::io::Cursor;
-use self::rand::Rng;
+use self::rand::distributions::range::Range;
+use self::rand::distributions::IndependentSample;
 use self::test::Bencher;
 
 #[bench]
-fn varint_write_rand_1_k(b: &mut Bencher) {
-    do_varint_write_rand(b, 1000)
+fn varint_write_rand(b: &mut Bencher) {
+    do_varint_write_rand(b, Range::new(0, u64::max_value()))
 }
 
 #[bench]
-fn varint_write_rand_1_m(b: &mut Bencher) {
-    do_varint_write_rand(b, 1000_000)
+fn varint_write_rand_1_byte(b: &mut Bencher) {
+    do_varint_write_rand(b, Range::new(0, 128))
 }
 
 #[bench]
-fn varint_read_rand_1_k(b: &mut Bencher) {
-    do_varint_read_rand(b, 1000)
+fn varint_write_rand_9_bytes(b: &mut Bencher) {
+    do_varint_write_rand(b, Range::new(1 << 56, u64::max_value()))
 }
 
 #[bench]
-fn varint_read_rand_1_m(b: &mut Bencher) {
-    do_varint_read_rand(b, 1000_000)
+fn varint_read_rand(b: &mut Bencher) {
+    do_varint_read_rand(b, Range::new(0, u64::max_value()))
 }
 
-fn do_varint_write_rand(b: &mut Bencher, num: usize) {
-    let mut rng = rand::weak_rng();
+#[bench]
+fn varint_read_rand_1_byte(b: &mut Bencher) {
+    do_varint_read_rand(b, Range::new(0, 128))
+}
+
+#[bench]
+fn varint_read_rand_9_byte(b: &mut Bencher) {
+    do_varint_read_rand(b, Range::new(1 << 56, u64::max_value()))
+}
+
+#[bench]
+fn varint_read_slice_rand(b: &mut Bencher) {
+    do_varint_read_slice_rand(b, Range::new(0, u64::max_value()))
+}
+
+#[bench]
+fn varint_read_slice_rand_1_byte(b: &mut Bencher) {
+    do_varint_read_slice_rand(b, Range::new(0, 128))
+}
 
+#[bench]
+fn varint_read_slice_rand_9_byte(b: &mut Bencher) {
+    do_varint_read_slice_rand(b, Range::new(1 << 56, u64::max_value()))
+}
+
+fn do_varint_write_rand(b: &mut Bencher, range: Range<u64>) {
+    let mut rng = rand::weak_rng();
+    let num = 1000_000;
     let mut vec: Vec<u64> = Vec::new();
 
     for _ in 0..num {
-        vec.push(rng.gen());
+        vec.push(range.ind_sample(&mut rng));
     }
 
     let mut buf = [0; 9];
@@ -44,15 +70,15 @@ fn do_varint_write_rand(b: &mut Bencher, num: usize) {
     });
 }
 
-fn do_varint_read_rand(b: &mut Bencher, num: usize) {
+fn do_varint_read_rand(b: &mut Bencher, range: Range<u64>) {
     let mut rng = rand::weak_rng();
-
+    let num = 1000_000;
     let mut vec = Vec::new();
     vec.resize(9 * num, 0);
     let mut bytes_written = 0;
 
     for _ in 0..num {
-        bytes_written += varint_write(rng.gen(), &mut vec[bytes_written..]);
+        bytes_written += varint_write(range.ind_sample(&mut rng), &mut vec[bytes_written..]);
     }
 
     b.iter(|| {
@@ -62,3 +88,26 @@ fn do_varint_read_rand(b: &mut Bencher, num: usize) {
         }
     });
 }
+
+fn do_varint_read_slice_rand(b: &mut Bencher, range: Range<u64>) {
+    let mut rng = rand::weak_rng();
+    let num = 1000_000;
+    let mut vec = Vec::new();
+
+    vec.resize(9 * num, 0);
+    let mut bytes_written = 0;
+
+    for _ in 0..num {
+        bytes_written += varint_write(range.ind_sample(&mut rng), &mut vec[bytes_written..]);
+    }
+
+    b.iter(|| {
+        let mut input_index = 0;
+        // cheat a little bit: this will skip the last couple numbers, but that's why we do a
+        // million numbers. Losing the last few won't be measurable.
+        while input_index < bytes_written - 9 {
+            let (_, bytes_read) = varint_read_slice(&vec[input_index..(input_index + 9)]);
+            input_index += bytes_read;
+        }
+    });
+}

src/serialization/deserializer.rs

@@ -2,6 +2,7 @@ use super::V2_COOKIE;
 use super::super::{Counter, Histogram, RestatState};
 use super::super::num::ToPrimitive;
 use std::io::{self, Cursor, ErrorKind, Read};
+use std::marker::PhantomData;
 use std;
 use super::byteorder::{BigEndian, ReadBytesExt};
 
@@ -54,8 +55,6 @@ impl Deserializer {
     /// bytes already in slice or `Vec` form.
     pub fn deserialize<T: Counter, R: Read>(&mut self, reader: &mut R)
                                             -> Result<Histogram<T>, DeserializeError> {
-        // TODO benchmark minimizing read calls by reading into a fixed-size header buffer
-
         let cookie = reader.read_u32::<BigEndian>()?;
 
         if cookie != V2_COOKIE {
@@ -87,31 +86,30 @@ impl Deserializer {
         let mut payload_slice = &mut self.payload_buf[0..payload_len];
         reader.read_exact(&mut payload_slice)?;
 
-        let mut cursor = Cursor::new(&payload_slice);
-        let mut dest_index: usize = 0;
+        let mut payload_index: usize = 0;
         let mut restat_state = RestatState::new();
-        while cursor.position() < payload_slice.len() as u64 {
-            let num = zig_zag_decode(varint_read(&mut cursor)?);
-
-            if num < 0 {
-                let zero_count = (-num).to_usize()
-                    .ok_or(DeserializeError::UsizeTypeTooSmall)?;
-                // skip the zeros
-                dest_index = dest_index.checked_add(zero_count)
-                    .ok_or(DeserializeError::UsizeTypeTooSmall)?;
-                continue;
-            } else {
-                let count: T = T::from_i64(num)
-                    .ok_or(DeserializeError::UnsuitableCounterType)?;
-
-                h.set_count_at_index(dest_index, count)
-                    .map_err(|_| DeserializeError::EncodedArrayTooLong)?;
-
-                restat_state.on_nonzero_count(dest_index, count);
-
-                dest_index = dest_index.checked_add(1)
-                    .ok_or(DeserializeError::UsizeTypeTooSmall)?;
-            }
+        let mut decode_state = DecodeLoopState::new();
+
+        while payload_index < payload_len.saturating_sub(9) {
+            // Read with fast loop until we are within 9 of the end. Fast loop can't handle EOF,
+            // so bail to slow version for the last few bytes.
+
+            let (zz_num, bytes_read) = varint_read_slice(
+                &payload_slice[payload_index..(payload_index + 9)]);
+            payload_index += bytes_read;
+
+            let count_or_zeros = zig_zag_decode(zz_num);
+
+            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
+        }
+
+        // Now read the leftovers
+        let leftover_slice = &payload_slice[payload_index..];
+        let mut cursor = Cursor::new(&leftover_slice);
+        while cursor.position() < leftover_slice.len() as u64 {
+            let count_or_zeros = zig_zag_decode(varint_read(&mut cursor)?);
+
+            decode_state.on_decoded_num(count_or_zeros, &mut restat_state, &mut h)?;
         }
 
         restat_state.update_histogram(&mut h);
@@ -120,6 +118,62 @@ impl Deserializer {
     }
 }
 
+// Only public for testing.
+/// Read from a slice that must be 9 bytes long or longer. Returns the decoded number and how many
+/// bytes were consumed.
+#[inline]
+pub fn varint_read_slice(slice: &[u8]) -> (u64, usize) {
+    let mut b = slice[0];
+
+    // take low 7 bits
+    let mut value: u64 = low_7_bits(b);
+    if !is_high_bit_set(b) {
+        return (value, 1);
+    }
+    // high bit set, keep reading
+    b = slice[1];
+    value |= low_7_bits(b) << 7;
+    if !is_high_bit_set(b) {
+        return (value, 2);
+    }
+    b = slice[2];
+    value |= low_7_bits(b) << 7 * 2;
+    if !is_high_bit_set(b) {
+        return (value, 3);
+    }
+    b = slice[3];
+    value |= low_7_bits(b) << 7 * 3;
+    if !is_high_bit_set(b) {
+        return (value, 4);
+    }
+    b = slice[4];
+    value |= low_7_bits(b) << 7 * 4;
+    if !is_high_bit_set(b) {
+        return (value, 5);
+    }
+    b = slice[5];
+    value |= low_7_bits(b) << 7 * 5;
+    if !is_high_bit_set(b) {
+        return (value, 6);
+    }
+    b = slice[6];
+    value |= low_7_bits(b) << 7 * 6;
+    if !is_high_bit_set(b) {
+        return (value, 7);
+    }
+    b = slice[7];
+    value |= low_7_bits(b) << 7 * 7;
+    if !is_high_bit_set(b) {
+        return (value, 8);
+    }
+
+    b = slice[8];
+    // special case: use last byte as is
+    value |= (b as u64) << 7 * 8;
+
+    (value, 9)
+}
+
 // Only public for testing.
 /// Read a LEB128-64b9B from the buffer
 pub fn varint_read<R: Read>(reader: &mut R) -> io::Result<u64> {
@@ -174,7 +228,6 @@ fn low_7_bits(b: u8) -> u64 {
 
 #[inline]
 fn is_high_bit_set(b: u8) -> bool {
-    // TODO benchmark leading zeros rather than masking
     (b & 0x80) != 0
 }
 
@@ -183,3 +236,45 @@ fn is_high_bit_set(b: u8) -> bool {
 pub fn zig_zag_decode(encoded: u64) -> i64 {
     ((encoded >> 1) as i64) ^ -((encoded & 1) as i64)
 }
+
+/// We need to perform the same logic in two different decode loops while carrying over a modicum
+/// of state.
+struct DecodeLoopState<T: Counter> {
+    dest_index:usize,
+    phantom: PhantomData<T>
+}
+
+impl <T: Counter> DecodeLoopState<T> {
+
+    fn new() -> DecodeLoopState<T> {
+        DecodeLoopState {
+            dest_index: 0,
+            phantom: PhantomData
+        }
+    }
+
+    #[inline]
+    fn on_decoded_num(&mut self, count_or_zeros: i64, restat_state: &mut RestatState<T>,
+                      h: &mut Histogram<T>) -> Result<(), DeserializeError> {
+        if count_or_zeros < 0 {
+            let zero_count = (-count_or_zeros).to_usize()
+                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
+            // skip the zeros
+            self.dest_index = self.dest_index.checked_add(zero_count)
+                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
+        } else {
+            let count: T = T::from_i64(count_or_zeros)
+                .ok_or(DeserializeError::UnsuitableCounterType)?;
+
+            h.set_count_at_index(self.dest_index, count)
+                .map_err(|_| DeserializeError::EncodedArrayTooLong)?;
+
+            restat_state.on_nonzero_count(self.dest_index, count);
+
+            self.dest_index = self.dest_index.checked_add(1)
+                .ok_or(DeserializeError::UsizeTypeTooSmall)?;
+        }
+
+        Ok(())
+    }
+}