Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimize VarInt decoding #1117

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 38 additions & 5 deletions quinn-proto/src/varint.rs
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,11 @@ impl Codec for VarInt {
return Err(UnexpectedEnd);
}
let mut buf = [0; 8];
buf[0] = r.get_u8();
// This does not make use of `r.get_u8()` to avoid the additional bounds
// check. We already performed this one earlier on.
buf[0] = r.chunk()[0];
r.advance(1);

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this specific change actually help? I wouldn't expect this to have any difference in bounds checking semantics.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's an assert! in the get_u8, not just a rust array bounds check. I don't think those are removed. And since even without the "le" change this proved to be faster, I guess practical experience shows that it helps.

let tag = buf[0] >> 6;
buf[0] &= 0b0011_1111;
let x = match tag {
Expand All @@ -159,21 +163,35 @@ impl Codec for VarInt {
if r.remaining() < 1 {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we do this without unsafe if we use <&[u8; 1]>::try_from() to extract an array out of the bounds check?

Along these lines:

diff --git a/quinn-proto/src/varint.rs b/quinn-proto/src/varint.rs
index 7f57bc3c..8f4b9b7c 100644
--- a/quinn-proto/src/varint.rs
+++ b/quinn-proto/src/varint.rs
@@ -1,4 +1,4 @@
-use std::{convert::TryInto, fmt};
+use std::{convert::{TryInto, TryFrom}, fmt};
 
 use bytes::{Buf, BufMut};
 use thiserror::Error;
@@ -146,38 +146,37 @@ pub struct VarIntBoundsExceeded;
 
 impl Codec for VarInt {
     fn decode<B: Buf>(r: &mut B) -> coding::Result<Self> {
-        if !r.has_remaining() {
-            return Err(UnexpectedEnd);
-        }
         let mut buf = [0; 8];
-        buf[0] = r.get_u8();
+        let input = r.chunk();
+        buf[0] = match input.get(0) {
+            Some(&b) => b,
+            None => return Err(UnexpectedEnd),
+        };
+
         let tag = buf[0] >> 6;
         buf[0] &= 0b0011_1111;
-        let x = match tag {
-            0b00 => u64::from(buf[0]),
+        let (x, len) = match tag {
+            0b00 => (u64::from(buf[0]), 1),
             0b01 => {
-                if r.remaining() < 1 {
-                    return Err(UnexpectedEnd);
-                }
-                r.copy_to_slice(&mut buf[1..2]);
-                u64::from(u16::from_be_bytes(buf[..2].try_into().unwrap()))
+                let arr = [0, buf[0]];
+                let read = <&[u8; 1]>::try_from(&input[1..]).map_err(|_| UnexpectedEnd)?;
+                let arr = [read[0], buf[0]];
+                (u64::from(u16::from_le_bytes(arr)), 2)
             }
             0b10 => {
-                if r.remaining() < 3 {
-                    return Err(UnexpectedEnd);
-                }
-                r.copy_to_slice(&mut buf[1..4]);
-                u64::from(u32::from_be_bytes(buf[..4].try_into().unwrap()))
+                let read = <&[u8; 3]>::try_from(&input[1..]).map_err(|_| UnexpectedEnd)?;
+                let arr = [read[2], read[1], read[0], buf[0]];
+                (u64::from(u32::from_le_bytes(arr)), 4)
             }
             0b11 => {
-                if r.remaining() < 7 {
-                    return Err(UnexpectedEnd);
-                }
-                r.copy_to_slice(&mut buf[1..8]);
-                u64::from_be_bytes(buf)
+                let read = <&[u8; 7]>::try_from(&input[1..]).map_err(|_| UnexpectedEnd)?;
+                let arr = [read[6], read[5], read[4], read[3], read[2], read[1], read[0], buf[0]];
+                (u64::from(u64::from_le_bytes(arr)), 8)
             }
             _ => unreachable!(),
         };
+
+        r.advance(len);
         Ok(VarInt(x))
     }

return Err(UnexpectedEnd);
}
r.copy_to_slice(&mut buf[1..2]);
u64::from(u16::from_be_bytes(buf[..2].try_into().unwrap()))

// A little-endian buffer is used here since that's our most
// likely target platform and it performs slightly better than
// the be conversion. For bigger numbers we don't bother, because
// it makes the byte copy more complicated.
let mut b = [0, buf[0]];
b[0] = r.chunk()[0];
r.advance(1);

u64::from(u16::from_le_bytes(b))
}
0b10 => {
if r.remaining() < 3 {
return Err(UnexpectedEnd);
}
r.copy_to_slice(&mut buf[1..4]);

// We performed the bounds check upfront
unsafe { copy_unchecked(r, &mut buf[1..4]) };

u64::from(u32::from_be_bytes(buf[..4].try_into().unwrap()))
}
0b11 => {
if r.remaining() < 7 {
return Err(UnexpectedEnd);
}
r.copy_to_slice(&mut buf[1..8]);

// We performed the bounds check upfront
unsafe { copy_unchecked(r, &mut buf[1..8]) };

u64::from_be_bytes(buf)
}
_ => unreachable!(),
Expand All @@ -196,3 +214,18 @@ impl Codec for VarInt {
}
}
}

unsafe fn copy_unchecked<B: Buf>(src: &mut B, dst: &mut [u8]) {
let mut off = 0;
let len = dst.len();

while off < dst.len() {
let chunk = src.chunk();
let to_copy = chunk.len().min(len - off);

std::ptr::copy_nonoverlapping(chunk.as_ptr(), dst[off..].as_mut_ptr(), to_copy);

off += to_copy;
src.advance(to_copy);
}
}