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.

Sign up for GitHub

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

feat: add basic codec support for binary le && binary unsafe #163

Merged
merged 1 commit into from
May 24, 2023
Merged

feat: add basic codec support for binary le && binary unsafe #163

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 4 additions & 3 deletions pilota-build/src/codegen/thrift/ty.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,12 +349,13 @@ impl ThriftBackend {
let read_el = self.codegen_decode_ty(helper, ty);
format! {
r#"
{{
unsafe {{
let list_ident = {read_list_begin};
let mut val = Vec::with_capacity(list_ident.size);
for _ in 0..list_ident.size {{
val.push({read_el});
for i in 0..list_ident.size {{
*val.get_unchecked_mut(i) = {read_el};
}};
val.set_len(list_ident.size);
{read_list_end};
val
}}
Expand Down
50 changes: 28 additions & 22 deletions pilota-build/test_data/thrift/wrapper_arc.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -616,22 +616,24 @@ pub mod wrapper_arc {
id = Some(protocol.read_faststr()?);
}
Some(2) if field_ident.field_type == ::pilota::thrift::TType::List => {
name2 = Some({
name2 = Some(unsafe {
let list_ident = protocol.read_list_begin()?;
let mut val = Vec::with_capacity(list_ident.size);
for _ in 0..list_ident.size {
val.push({
for i in 0..list_ident.size {
*val.get_unchecked_mut(i) = unsafe {
let list_ident = protocol.read_list_begin()?;
let mut val = Vec::with_capacity(list_ident.size);
for _ in 0..list_ident.size {
val.push(::std::sync::Arc::new(
for i in 0..list_ident.size {
*val.get_unchecked_mut(i) = ::std::sync::Arc::new(
::pilota::thrift::Message::decode(protocol)?,
));
);
}
val.set_len(list_ident.size);
protocol.read_list_end()?;
val
});
};
}
val.set_len(list_ident.size);
protocol.read_list_end()?;
val
});
Expand All @@ -642,14 +644,15 @@ pub mod wrapper_arc {
let mut val =
::std::collections::HashMap::with_capacity(map_ident.size);
for _ in 0..map_ident.size {
val.insert(protocol.read_i32()?, {
val.insert(protocol.read_i32()?, unsafe {
let list_ident = protocol.read_list_begin()?;
let mut val = Vec::with_capacity(list_ident.size);
for _ in 0..list_ident.size {
val.push(::std::sync::Arc::new(
for i in 0..list_ident.size {
*val.get_unchecked_mut(i) = ::std::sync::Arc::new(
::pilota::thrift::Message::decode(protocol)?,
));
);
}
val.set_len(list_ident.size);
protocol.read_list_end()?;
val
});
Expand Down Expand Up @@ -736,25 +739,27 @@ pub mod wrapper_arc {
id = Some(protocol.read_faststr().await?);
}
Some(2) if field_ident.field_type == ::pilota::thrift::TType::List => {
name2 = Some({
name2 = Some(unsafe {
let list_ident = protocol.read_list_begin().await?;
let mut val = Vec::with_capacity(list_ident.size);
for _ in 0..list_ident.size {
val.push({
for i in 0..list_ident.size {
*val.get_unchecked_mut(i) = unsafe {
let list_ident = protocol.read_list_begin().await?;
let mut val = Vec::with_capacity(list_ident.size);
for _ in 0..list_ident.size {
val.push(::std::sync::Arc::new(
for i in 0..list_ident.size {
*val.get_unchecked_mut(i) = ::std::sync::Arc::new(
::pilota::thrift::Message::decode_async(
protocol,
)
.await?,
));
);
}
val.set_len(list_ident.size);
protocol.read_list_end().await?;
val
});
};
}
val.set_len(list_ident.size);
protocol.read_list_end().await?;
val
});
Expand All @@ -765,17 +770,18 @@ pub mod wrapper_arc {
let mut val =
::std::collections::HashMap::with_capacity(map_ident.size);
for _ in 0..map_ident.size {
val.insert(protocol.read_i32().await?, {
val.insert(protocol.read_i32().await?, unsafe {
let list_ident = protocol.read_list_begin().await?;
let mut val = Vec::with_capacity(list_ident.size);
for _ in 0..list_ident.size {
val.push(::std::sync::Arc::new(
for i in 0..list_ident.size {
*val.get_unchecked_mut(i) = ::std::sync::Arc::new(
::pilota::thrift::Message::decode_async(
protocol,
)
.await?,
));
);
}
val.set_len(list_ident.size);
protocol.read_list_end().await?;
val
});
Expand Down
4 changes: 4 additions & 0 deletions pilota/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,3 +44,7 @@ rand = "0.8"
[[bench]]
name = "faststr"
harness = false

[[bench]]
name = "thrift_binary"
harness = false
262 changes: 262 additions & 0 deletions pilota/benches/thrift_binary.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,262 @@
#![allow(clippy::redundant_clone)]

use bytes::BytesMut;
use criterion::{black_box, criterion_group, criterion_main};
use pilota::thrift::{TInputProtocol, TOutputProtocol};
use rand::{self, Rng};

fn binary_bench(c: &mut criterion::Criterion) {
let size = std::env::var("SIZE")
.unwrap_or("10000".to_string())
.parse()
.unwrap();
let mut group = c.benchmark_group("Bench Thrift Binary");
let mut v: Vec<i64> = Vec::with_capacity(size);
for _ in 0..size {
v.push(rand::thread_rng().gen());
}
let mut buf = BytesMut::new();

let mut p = pilota::thrift::binary::TBinaryProtocol::new(&mut buf, true);
for i in &v {
p.write_i64(*i).unwrap();
}
drop(p);
assert_eq!(buf.len(), 8 * size);

let mut buf_le = BytesMut::new();

let mut p = pilota::thrift::binary_le::TBinaryProtocol::new(&mut buf_le, true);
for i in &v {
p.write_i64(*i).unwrap();
}
drop(p);
assert_eq!(buf_le.len(), 8 * size);

let b = buf_le.clone();
let mut v2: Vec<i64> = Vec::with_capacity(size);
let src = b.as_ptr();
let dst = v2.as_mut_ptr();
unsafe {
std::ptr::copy_nonoverlapping(src, dst as *mut u8, size * 8);
v2.set_len(size);
}
assert_eq!(v, v2);

group.bench_function("big endian decode vec i64", |b| {
b.iter(|| {
black_box({
let b = buf.clone();
black_box(read_be(b, size));
});
})
});

group.bench_function("big endian decode vec i64 unsafe", |b| {
b.iter(|| {
black_box({
let b = buf.clone();
black_box(read_be_unsafe(b, size));
});
})
});

group.bench_function("big endian decode vec i64 unsafe vec", |b| {
b.iter(|| {
black_box({
let b = buf.clone();
black_box(read_be_unsafe_vec(b, size));
});
})
});

group.bench_function("big endian decode vec i64 unsafe optimized", |b| {
b.iter(|| {
black_box({
let b = buf.clone();
black_box(read_be_unsafe_optimized(b, size));
});
})
});

group.bench_function("big endian encode vec i64", |b| {
b.iter(|| {
black_box({
let mut b = BytesMut::with_capacity(8 * size);
black_box(write_be(&mut b, &v, size));
});
})
});

group.bench_function("big endian encode vec i64 unsafe", |b| {
b.iter(|| {
black_box({
let mut b = BytesMut::with_capacity(8 * size);
black_box(write_be_unsafe(&mut b, &v, size));
});
})
});

group.bench_function("little endian decode vec i64", |b| {
b.iter(|| {
black_box({
let b = buf_le.clone();
black_box(read_le(b, size));
});
})
});
group.bench_function("little endian decode vec i64 unsafe optimized", |b| {
b.iter(|| {
black_box({
let b = buf_le.clone();
black_box(read_le_unsafe_optimized(b, size));
});
})
});
group.bench_function("little endian decode vec i64 optimized", |b| {
b.iter(|| {
black_box({
let b = buf_le.clone();
black_box(read_le_optimized(b, size));
});
})
});

group.bench_function("alloc vec", |b| {
b.iter(|| {
let mut b = buf_le.clone();
let _p = pilota::thrift::binary_le::TBinaryProtocol::new(&mut b, true);
let _: Vec<i64> = black_box(Vec::with_capacity(size));
})
});

group.finish();
}

#[inline(never)]
fn read_be(mut b: BytesMut, size: usize) -> Vec<i64> {
let mut p = pilota::thrift::binary::TBinaryProtocol::new(&mut b, true);
let mut v = Vec::with_capacity(size);
for _ in 0..size {
v.push(p.read_i64().unwrap());
}
v
}

#[inline(never)]
fn read_be_unsafe(mut b: BytesMut, size: usize) -> Vec<i64> {
unsafe {
let s = std::slice::from_raw_parts_mut(b.as_mut_ptr(), b.len());
let mut p = pilota::thrift::binary_unsafe::TBinaryProtocol::new(&mut b, s, true);
let mut v = Vec::with_capacity(size);
for _ in 0..size {
v.push(p.read_i64().unwrap());
}
v
}
}

#[inline(never)]
fn read_be_unsafe_vec(mut b: BytesMut, size: usize) -> Vec<i64> {
unsafe {
let s = std::slice::from_raw_parts_mut(b.as_mut_ptr(), b.len());
let mut p = pilota::thrift::binary_unsafe::TBinaryProtocol::new(&mut b, s, true);
let mut v = Vec::with_capacity(size);
for i in 0..size {
*v.get_unchecked_mut(i) = p.read_i64().unwrap();
}
v
}
}

#[inline(never)]
fn read_be_unsafe_optimized(b: BytesMut, size: usize) -> Vec<i64> {
unsafe {
let buf: &[u8] = b.as_ref();
assert!(buf.len() >= size * 8);
let mut index = 0;

let mut v = Vec::with_capacity(size);
for i in 0..size {
*v.get_unchecked_mut(i) = i64::from_be_bytes(
buf.get_unchecked(index..index + 8)
.try_into()
.unwrap_unchecked(),
);
index += 8;
}
v.set_len(size);
v
}
}

#[inline(never)]
fn write_be(b: &mut BytesMut, v: &Vec<i64>, size: usize) {
let mut p = pilota::thrift::binary::TBinaryProtocol::new(b, true);
for el in v {
p.write_i64(*el).unwrap();
}
}

#[inline(never)]
fn write_be_unsafe(b: &mut BytesMut, v: &Vec<i64>, size: usize) {
unsafe {
let s = std::slice::from_raw_parts_mut(b.as_mut_ptr(), b.len());
let mut p = pilota::thrift::binary_unsafe::TBinaryProtocol::new(b, s, true);
for el in v {
p.write_i64(*el).unwrap();
}
}
}

#[inline(never)]
fn read_le(mut b: BytesMut, size: usize) -> Vec<i64> {
let mut p = pilota::thrift::binary_le::TBinaryProtocol::new(&mut b, true);

let mut v = Vec::with_capacity(size);
for _ in 0..size {
v.push(p.read_i64().unwrap());
}
v
}

// cargo asm -p pilota --bench thrift_binary --native --full-name --keep-labels
// --simplify --rust
#[inline(never)]
fn read_le_unsafe_optimized(b: BytesMut, size: usize) -> Vec<i64> {
unsafe {
let buf: &[u8] = b.as_ref();
assert!(buf.len() >= size * 8);
let mut index = 0;

let mut v = Vec::with_capacity(size);
for i in 0..size {
*v.get_unchecked_mut(i) = i64::from_le_bytes(
buf.get_unchecked(index..index + 8)
.try_into()
.unwrap_unchecked(),
);
index += 8;
}
v.set_len(size);
v
}
}

#[inline(never)]
fn read_le_optimized(mut b: BytesMut, size: usize) -> Vec<i64> {
let _p = pilota::thrift::binary_le::TBinaryProtocol::new(&mut b, true);
let mut v: Vec<i64> = Vec::with_capacity(size);
let _ = black_box({
let src = b.as_ptr();
let dst = v.as_mut_ptr();
unsafe {
std::ptr::copy_nonoverlapping(src, dst as *mut u8, size * 8);
v.set_len(size);
}
});
v
}

criterion_group!(benches, binary_bench);
criterion_main!(benches);
Loading