├── .envrc ├── .gitignore ├── Cargo.toml ├── README.md ├── benches └── codec_comparison.rs ├── examples ├── chat_client.rs ├── chat_server.rs └── no-std │ ├── Cargo.toml │ └── src │ └── main.rs ├── flake.lock ├── flake.nix ├── rust-toolchain.toml ├── rustfmt.toml └── src ├── framed_codec.rs ├── length_codec.rs ├── lib.rs └── rkyv_codec.rs /.envrc: -------------------------------------------------------------------------------- 1 | use flake -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | .vscode 3 | Cargo.lock 4 | .direnv -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "rkyv_codec" 3 | version = "0.5.0" 4 | edition = "2024" 5 | license = "MIT OR Apache-2.0" 6 | description = "Some adaptors to stream rkyv Archives over AsyncRead and AsyncWrite" 7 | repository = "https://github.com/zyansheep/rkyv_codec" 8 | 9 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 10 | 11 | [dependencies] 12 | unsigned-varint = { version = "0.8", default-features = false } 13 | rkyv = { version = "0.8", default-features = false, features = [ 14 | "bytecheck", 15 | "alloc", 16 | ] } 17 | pin-project = "1.1" 18 | futures = { version = "0.3", optional = true } 19 | asynchronous-codec = { version = "0.7", optional = true } 20 | bytes = { version = "1.10" } 21 | thiserror = { version = "1.0", package = "thiserror-core", default-features = false } 22 | 23 | [dev-dependencies] 24 | async-std = { version = "1.13.0", features = ["attributes"] } 25 | # For tests & benchmark comparisons 26 | asynchronous-codec = { version = "0.7", features = ["cbor"] } 27 | criterion = { version = "0.6", features = ["async_std"] } 28 | serde = "^1.0" 29 | lazy_static = "1.5.0" 30 | bincode = "^2.0" # For benchmark comparisons 31 | rand = "^0.9" # For generating test data 32 | 33 | # For macros 34 | paste = "1.0" 35 | 36 | # For examples 37 | async-broadcast = "0.7.1" 38 | rustyline-async = "0.4.0" 39 | anyhow = "1.0" 40 | 41 | 42 | [features] 43 | default = ["std"] # enable variable-length length codec 44 | std = [ 45 | "thiserror/std", 46 | "rkyv/std", 47 | "unsigned-varint/std", 48 | "asynchronous-codec", 49 | "unsigned-varint/futures", 50 | "futures", 51 | ] 52 | 53 | [workspace] 54 | members = ["examples/no-std"] 55 | 56 | [[example]] 57 | name = "chat_server" 58 | required-features = ["std"] 59 | 60 | [[example]] 61 | name = "chat_client" 62 | required-features = ["std"] 63 | 64 | [package.metadata."docs.rs"] 65 | all-features = true 66 | 67 | [[bench]] 68 | name = "codec_comparison" 69 | harness = false 70 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # rkyv_codec 2 |

3 | 4 | docs.rs 5 | 6 | 7 | crates.io 8 | 9 |

10 | 11 | Simple async codec for [rkyv](https://github.com/rkyv/rkyv). Reuses streaming buffer for maximum speed! 12 | 13 | This crate provides a makeshift adaptor for streaming `&Archived`s from an `AsyncRead` using a reusable external buffer, as well as a `futures::Sink` implementation to serialize `Object`s to an `AsyncWrite`. 14 | It uses multiformat's [unsigned_varint](https://docs.rs/unsigned-varint/latest/unsigned_varint/) for variable-length length encoding by default, but allows for other kinds of length encoding through the `LengthEncoding` trait. 15 | It also supports directly using `bytes::BytesMut` and `#[no_std]` when default features are disabled. 16 | 17 | ## Examples 18 | This crate has three examples: chat_client, chat_server & no-std. Run the first two at the same time to see a proof-of-concept Archive tcp echo server in action. 19 | 20 | To run: 21 | 22 | `cargo run --example chat_client` 23 | 24 | `cargo run --example chat_server` 25 | 26 | Simple usage example (RkyvCodec): 27 | ```rust 28 | use rkyv::{Archive, Serialize, Deserialize, rancor, util::AlignedVec}; 29 | use rkyv_codec::{RkyvCodec, VarintLength}; 30 | use asynchronous_codec::{Framed, Decoder, Encoder}; 31 | use bytes::BytesMut; 32 | 33 | #[derive(Archive, Deserialize, Serialize, Debug, PartialEq, Clone)] 34 | #[rkyv(attr(derive(Debug)))] 35 | struct Test { 36 | int: u8, 37 | string: String, 38 | option: Option>, 39 | } 40 | 41 | let value = Test { 42 | int: 42, 43 | string: "hello world".to_string(), 44 | option: Some(vec![1, 2, 3, 4]), 45 | }; 46 | let mut codec = RkyvCodec::::default(); 47 | let mut buf = BytesMut::new(); 48 | 49 | // Encoding 50 | codec.encode(&value, &mut buf).unwrap(); 51 | 52 | // Decoding 53 | let decoded_value = codec.decode(&mut buf).unwrap().unwrap(); 54 | assert_eq!(value, decoded_value); 55 | ``` 56 | 57 | Zero-copy archive usage example (RkyvWriter/archive_stream): 58 | ```rust 59 | use rkyv::{Archived, util::AlignedVec, Archive, Serialize, Deserialize, rancor}; 60 | use rkyv_codec::{archive_stream, RkyvWriter, VarintLength}; 61 | use futures::SinkExt; 62 | 63 | #[derive(Archive, Deserialize, Serialize, Debug, PartialEq, Clone)] 64 | #[rkyv(derive(Debug), compare(PartialEq))] 65 | struct Test { 66 | int: u8, 67 | string: String, 68 | option: Option>, 69 | } 70 | 71 | let value = Test { 72 | int: 42, 73 | string: "hello world".to_string(), 74 | option: Some(vec![1, 2, 3, 4]), 75 | }; 76 | 77 | // Writing 78 | let writer = Vec::new(); 79 | let mut codec = RkyvWriter::<_, VarintLength>::new(writer); 80 | codec.send(&value).await.unwrap(); 81 | 82 | // Reading 83 | let mut reader = &codec.inner()[..]; 84 | let mut buffer = AlignedVec::new(); // Aligned streaming buffer for re-use 85 | let value_archived: &Archived = archive_stream::<_, Test, VarintLength>(&mut reader, &mut buffer).await.unwrap(); // This returns a reference into the passed buffer 86 | // can deserialize as normal as well (or do *partial* deserialization for blazingly fast speeds!) 87 | let value_deserialized: Test = rkyv::deserialize::<_, rancor::Error>(value_archived).unwrap(); 88 | 89 | assert_eq!(value, *value_archived); // compare to archived version 90 | 91 | assert_eq!(value, value_deserialized); // compare to deserialized version 92 | ``` 93 | 94 | See [`examples/no-std`](examples/no-std/src/main.rs) for an example with no-std support. 95 | 96 | ## Benchmarks 97 | 98 | These are a set of benchmarks, each benchmark represents 50 test objects being either sent or received. (requires nightly) 99 | ``` 100 | test tests::bench_archive_sink_50 ... bench: 6,276.07 ns/iter (+/- 250.20) 101 | test tests::bench_archive_sink_prearchived_50 ... bench: 2,009.84 ns/iter (+/- 175.23) 102 | test tests::bench_archive_stream_50 ... bench: 2,110.66 ns/iter (+/- 150.51) 103 | test tests::bench_futures_cbor_sink_50 ... bench: 8,178.37 ns/iter (+/- 123.13) 104 | test tests::bench_futures_cbor_stream_50 ... bench: 6,606.01 ns/iter (+/- 129.89) 105 | test tests::bench_rkyv_asynchronous_codec_sink_50 ... bench: 4,328.91 ns/iter (+/- 415.17) 106 | test tests::bench_rkyv_asynchronous_codec_stream_50 ... bench: 4,059.72 ns/iter (+/- 243.54) 107 | test tests::bench_rkyv_writer_50 ... bench: 2,228.49 ns/iter (+/- 185.92) 108 | test tests::bench_u64_length_encoding ... bench: 2,494.63 ns/iter (+/- 143.85) 109 | test tests::bench_varint_length_encoding ... bench: 4,036.83 ns/iter (+/- 471.98) 110 | ``` 111 | The fastest real benchmark (full serialization and bytecheck) is using `RkyvWriter` for writing and `archive_stream` for reading. 112 | This is compared to the slowest benchmark: the `asynchronous-codec` library's `CborCodec`. 113 | 114 | Feel free to contribute your own benchmarks! -------------------------------------------------------------------------------- /benches/codec_comparison.rs: -------------------------------------------------------------------------------- 1 | use bincode::{Decode, Encode}; 2 | use bytes::{Buf, BytesMut}; 3 | use criterion::{ 4 | AxisScale, BenchmarkId, Criterion, PlotConfiguration, Throughput, 5 | async_executor::AsyncStdExecutor, criterion_group, criterion_main, 6 | }; 7 | // rkyv imports for the Packet struct 8 | use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize}; 9 | // rkyv_codec imports 10 | use rkyv_codec::{ 11 | LengthCodec, 12 | RkyvCodec, // The one for asynchronous-codec 13 | RkyvCodecError, // The error type 14 | U32Length, 15 | U64Length, 16 | VarintLength, 17 | }; 18 | use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize}; 19 | use std::{hint::black_box, time::Duration}; 20 | 21 | // Define a sample data structure to be serialized/deserialized 22 | #[derive( 23 | Archive, 24 | RkyvDeserialize, 25 | RkyvSerialize, 26 | SerdeDeserialize, 27 | SerdeSerialize, 28 | Debug, 29 | PartialEq, 30 | Encode, 31 | Decode, 32 | Clone, 33 | )] 34 | struct Packet { 35 | id: u32, 36 | data: Vec, 37 | name: String, 38 | } 39 | 40 | fn create_packet(size: usize) -> Packet { 41 | Packet { 42 | id: rand::random(), 43 | data: vec![0u8; size - 50], // Adjust to make overall packet size roughly `size` 44 | name: "benchmark_packet_name_string_long_enough".to_string(), 45 | } 46 | } 47 | 48 | // --- Rkyv Benchmarking --- 49 | // Adjusted to use asynchronous-codec style with BytesMut 50 | async fn run_rkyv_benchmark( 51 | c: &mut Criterion, 52 | group_name: &str, 53 | packet_size: usize, 54 | ) where 55 | for<'a> ::Error: std::fmt::Debug, 56 | RkyvCodecError: From, // Ensure RkyvCodecError can be created from rkyv errors 57 | { 58 | let packet = create_packet(packet_size); 59 | let mut group = c.benchmark_group(group_name); 60 | group.throughput(Throughput::Bytes(packet_size as u64)); 61 | group.plot_config(PlotConfiguration::default().summary_scale(AxisScale::Logarithmic)); 62 | 63 | group.bench_function( 64 | BenchmarkId::new( 65 | format!( 66 | "rkyv_{}", 67 | std::any::type_name::().split("::").last().unwrap_or("") 68 | ), 69 | packet_size, 70 | ), 71 | |b| { 72 | b.to_async(AsyncStdExecutor).iter_batched( 73 | || { 74 | let codec: RkyvCodec = RkyvCodec::default(); 75 | let transport_buf = BytesMut::new(); 76 | (codec, transport_buf, packet.clone()) 77 | }, 78 | |(mut codec, mut transport_buf, p)| async move { 79 | use asynchronous_codec::{Decoder, Encoder}; 80 | // Encode (Send) 81 | codec.encode(p, &mut transport_buf).unwrap(); 82 | 83 | // Decode (Receive) 84 | let _received_packet: Option = 85 | codec.decode(&mut transport_buf).unwrap(); 86 | black_box((codec, transport_buf, _received_packet)); 87 | }, 88 | criterion::BatchSize::SmallInput, 89 | ) 90 | }, 91 | ); 92 | group.finish(); 93 | } 94 | 95 | // --- Bincode Benchmarking --- 96 | async fn run_bincode_benchmark( 97 | c: &mut Criterion, 98 | group_name: &str, 99 | packet_size: usize, 100 | ) where 101 | for<'a> ::Error: std::fmt::Debug, 102 | { 103 | let packet = create_packet(packet_size); 104 | let mut group = c.benchmark_group(group_name); 105 | group.throughput(Throughput::Bytes(packet_size as u64)); 106 | group.plot_config(PlotConfiguration::default().summary_scale(AxisScale::Logarithmic)); 107 | 108 | group.bench_function( 109 | BenchmarkId::new( 110 | format!( 111 | "bincode_{}", 112 | std::any::type_name::().split("::").last().unwrap_or("") 113 | ), 114 | packet_size, 115 | ), 116 | |b| { 117 | b.to_async(AsyncStdExecutor).iter_batched( 118 | || { 119 | let transport = BytesMut::new(); 120 | let mut buf = Vec::new(); 121 | buf.resize(packet_size + 1024, 0); 122 | // Simulate a framed transport for bincode using LengthCodec manually 123 | (transport, packet.clone(), buf) 124 | }, 125 | |(mut transport, p, mut serialize_buf)| async move { 126 | // Encode 127 | let bytes_written = bincode::encode_into_slice( 128 | &p, 129 | &mut serialize_buf, 130 | bincode::config::standard(), 131 | ) 132 | .unwrap(); 133 | let mut length_storage = L::Buffer::default(); // Use the LengthCodec's own buffer type 134 | let encoded_len_bytes = L::encode(serialize_buf.len(), &mut length_storage); 135 | transport.extend_from_slice(encoded_len_bytes); 136 | transport.extend_from_slice(&serialize_buf[..bytes_written]); 137 | 138 | // Decode 139 | // Simulate reading from the 'transport' BytesMut 140 | let len = L::decode_bytes(&mut transport).unwrap(); 141 | 142 | let data = transport.take(len); 143 | let _received_packet: Packet = bincode::decode_from_std_read( 144 | &mut data.reader(), 145 | bincode::config::standard(), 146 | ) 147 | .unwrap(); 148 | 149 | // For black_box, ensure transport is what's left after reading one message 150 | // Or, if we want to black_box the state *before* decode, clone transport earlier. 151 | // Here, transport is modified. If the bench is just one send/recv, this is fine. 152 | black_box(_received_packet); 153 | }, 154 | criterion::BatchSize::SmallInput, 155 | ) 156 | }, 157 | ); 158 | group.finish(); 159 | } 160 | 161 | fn codec_benchmarks(c: &mut Criterion) { 162 | let packet_sizes = [64, 256, 1024, 4096, 16384]; // Example packet sizes 163 | 164 | for &size in &packet_sizes { 165 | // Rkyv 166 | async_std::task::block_on(run_rkyv_benchmark::( 167 | c, 168 | "Codec Comparison (U32Length)", 169 | size, 170 | )); 171 | async_std::task::block_on(run_rkyv_benchmark::( 172 | c, 173 | "Codec Comparison (U64Length)", 174 | size, 175 | )); 176 | async_std::task::block_on(run_rkyv_benchmark::( 177 | c, 178 | "Codec Comparison (VarintLength)", 179 | size, 180 | )); 181 | 182 | // Bincode 183 | async_std::task::block_on(run_bincode_benchmark::( 184 | c, 185 | "Codec Comparison (U32Length)", 186 | size, 187 | )); 188 | async_std::task::block_on(run_bincode_benchmark::( 189 | c, 190 | "Codec Comparison (U64Length)", 191 | size, 192 | )); 193 | async_std::task::block_on(run_bincode_benchmark::( 194 | c, 195 | "Codec Comparison (VarintLength)", 196 | size, 197 | )); 198 | } 199 | } 200 | 201 | criterion_group! { 202 | name = benches; 203 | config = Criterion::default() 204 | .measurement_time(Duration::from_millis(100)) // Adjust measurement time 205 | .warm_up_time(Duration::from_millis(100)) // Adjust warm-up time 206 | .sample_size(300); // Adjust sample size 207 | targets = codec_benchmarks 208 | } 209 | 210 | criterion_main!(benches); 211 | -------------------------------------------------------------------------------- /examples/chat_client.rs: -------------------------------------------------------------------------------- 1 | use std::{fmt, io::Write}; 2 | 3 | use async_std::{io, net::TcpStream}; 4 | use futures::{FutureExt, SinkExt}; 5 | 6 | use rkyv::{Archive, Deserialize, Serialize, rancor, util::AlignedVec}; 7 | 8 | use rkyv_codec::{RkyvWriter, VarintLength, archive_stream}; 9 | 10 | use rustyline_async::{Readline, ReadlineEvent}; 11 | 12 | #[derive(Archive, Deserialize, Serialize, Debug, PartialEq, Clone)] 13 | // This will generate a PartialEq impl between our unarchived and archived types 14 | // To use the safe API, you must use the check_bytes option for the archive 15 | #[rkyv(compare(PartialEq), attr(derive(Debug)))] 16 | struct ChatMessage { 17 | sender: Option, 18 | message: String, 19 | } 20 | impl fmt::Display for ChatMessage { 21 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 22 | write!( 23 | f, 24 | "<{}>: {}", 25 | self.sender.as_deref().unwrap_or("Anonymous"), 26 | self.message 27 | ) 28 | } 29 | } 30 | 31 | #[async_std::main] 32 | async fn main() -> io::Result<()> { 33 | // Connect to server 34 | let tcp_stream = TcpStream::connect("127.0.0.1:8080").await?; 35 | println!("Connected to {}", &tcp_stream.peer_addr()?); 36 | 37 | // Setup outgoing packet stream 38 | let mut packet_sender = RkyvWriter::<_, VarintLength>::new(tcp_stream.clone()); 39 | 40 | // Incoming packet buffer 41 | let mut buffer = AlignedVec::new(); 42 | 43 | let (mut rl, mut writer) = Readline::new("> ".to_owned()).unwrap(); 44 | 45 | // buffer the tcp stream 46 | let mut tcp_stream = futures::io::BufReader::new(tcp_stream); 47 | 48 | loop { 49 | futures::select! { 50 | archive = archive_stream::<_, ChatMessage, VarintLength>(&mut tcp_stream, &mut buffer).fuse() => match archive { 51 | Ok(archive) => { 52 | let message: ChatMessage = rkyv::deserialize::<_, rancor::Error>(archive).unwrap(); 53 | writeln!(writer, "{message}")?; 54 | } 55 | Err(err) => { 56 | writeln!(writer, "error parsing message: {err}")?; 57 | break; 58 | } 59 | }, 60 | line = rl.readline().fuse() => match line { 61 | Ok(ReadlineEvent::Line(line)) => { 62 | let message = ChatMessage { 63 | sender: None, 64 | message: line, 65 | }; 66 | packet_sender.send(&message).await.unwrap(); 67 | } 68 | Ok(ReadlineEvent::Interrupted) => { 69 | writeln!(writer, "CTRL-C, do CTRL-D to exit")?; 70 | }, 71 | Ok(ReadlineEvent::Eof) => { 72 | writeln!(writer, "CTRL-D")?; 73 | break 74 | }, 75 | Err(err) => { write!(writer, "error reading input: {err}")?; break }, 76 | } 77 | } 78 | } 79 | 80 | let _ = rl.flush(); 81 | 82 | Ok(()) 83 | } 84 | -------------------------------------------------------------------------------- /examples/chat_server.rs: -------------------------------------------------------------------------------- 1 | use std::net::SocketAddr; 2 | 3 | use anyhow::Context; 4 | use async_broadcast::{Receiver, Sender}; 5 | use async_std::{ 6 | io, 7 | net::{TcpListener, TcpStream}, 8 | task, 9 | }; 10 | use futures::{SinkExt, StreamExt, prelude::*}; 11 | 12 | use rkyv::{Archive, Deserialize, Serialize, rancor, util::AlignedVec}; 13 | 14 | use rkyv_codec::{RkyvWriter, VarintLength, archive_stream}; 15 | 16 | #[derive(Archive, Deserialize, Serialize, Debug, PartialEq, Clone)] 17 | // This will generate a PartialEq impl between our unarchived and archived types 18 | // To use the safe API, you have to enable the check_bytes option for the archive 19 | #[rkyv(compare(PartialEq), attr(derive(Debug)))] 20 | struct ChatMessage { 21 | sender: Option, 22 | message: String, 23 | } 24 | 25 | // Process a given TcpStream 26 | async fn handle_conn( 27 | stream: TcpStream, 28 | outgoing: Sender, 29 | mut incoming: Receiver, 30 | addr: &SocketAddr, 31 | ) -> anyhow::Result<()> { 32 | println!("[{addr}] Joined Server"); 33 | outgoing 34 | .broadcast(ChatMessage { 35 | sender: Some("Sever".to_owned()), 36 | message: format!("{} Joined the Chat!", addr), 37 | }) 38 | .await?; 39 | 40 | let reader = stream.clone(); 41 | let mut writer = RkyvWriter::<_, VarintLength>::new(stream); 42 | 43 | let mut buffer = AlignedVec::new(); 44 | 45 | let mut reader = futures::io::BufReader::new(reader); 46 | 47 | loop { 48 | futures::select! { 49 | // Read incoming messages 50 | archive = archive_stream::<_, ChatMessage, VarintLength>(&mut reader, &mut buffer).fuse() => match archive { 51 | Ok(archive) => { 52 | let mut msg: ChatMessage = rkyv::deserialize::(archive).unwrap(); 53 | msg.sender = Some(format!("{addr}")); 54 | println!("[{addr}] sent {msg:?}"); 55 | outgoing.broadcast(msg).await?; 56 | } 57 | _ => break, 58 | }, 59 | msg = incoming.next().fuse() => { 60 | writer.send(&msg.context("incoming channel closed")?).await.unwrap() 61 | } 62 | } 63 | } 64 | 65 | println!("[{addr}] Left Server"); 66 | outgoing 67 | .broadcast(ChatMessage { 68 | sender: Some("Sever".to_owned()), 69 | message: format!("{} Left the Chat!", addr), 70 | }) 71 | .await?; 72 | 73 | Ok(()) 74 | } 75 | 76 | #[async_std::main] 77 | async fn main() -> io::Result<()> { 78 | let listener = TcpListener::bind("127.0.0.1:8080").await?; // Bind to local addr 79 | println!("Listening on {}", listener.local_addr()?); 80 | 81 | // Init broadcast channels 82 | let (broadcast_sender, broadcast_receiver) = async_broadcast::broadcast::(20); 83 | 84 | // Listen for incoming connections 85 | let mut incoming = listener.incoming(); 86 | while let Some(stream) = incoming.next().await { 87 | match stream { 88 | Ok(stream) => { 89 | let outgoing = broadcast_sender.clone(); // clone the channels 90 | let incoming = broadcast_receiver.clone(); 91 | 92 | task::spawn(async move { 93 | // spawn a greenthread to handle the connection 94 | let addr = stream.peer_addr().unwrap(); 95 | if let Err(err) = handle_conn(stream, outgoing, incoming, &addr).await { 96 | println!("[{addr}] error: {err}") 97 | } 98 | }); 99 | } 100 | Err(err) => println!("error: {err}"), 101 | }; 102 | } 103 | 104 | Ok(()) 105 | } 106 | -------------------------------------------------------------------------------- /examples/no-std/Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "no-std" 3 | version = "0.1.0" 4 | edition = "2021" 5 | 6 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 7 | 8 | [dependencies] 9 | bytes = "1.10.0" 10 | rkyv = { version = "0.8", default-features = false } 11 | rkyv_codec = { path = "../.." } 12 | -------------------------------------------------------------------------------- /examples/no-std/src/main.rs: -------------------------------------------------------------------------------- 1 | #![no_std] 2 | extern crate alloc; 3 | 4 | use core::mem::MaybeUninit; 5 | 6 | use alloc::vec::Vec; 7 | 8 | use bytes::BytesMut; 9 | use rkyv::{ 10 | api::low, 11 | rancor, 12 | ser::{allocator::SubAllocator, writer::Buffer}, 13 | util::{Align, AlignedVec}, 14 | Archive, Deserialize, Serialize, 15 | }; 16 | use rkyv_codec::VarintLength; 17 | 18 | #[derive(Debug, PartialEq, Archive, Serialize, Deserialize)] 19 | #[rkyv(compare(PartialEq), attr(derive(Debug)))] 20 | struct Test { 21 | bytes: Vec, 22 | number: u64, 23 | } 24 | 25 | fn main() { 26 | // let archived_bytes = rkyv::a 27 | // Or you can customize your serialization for better performance 28 | // and compatibility with #![no_std] environments 29 | let value = Test { 30 | bytes: Vec::from_iter([1, 2, 3]), 31 | number: 42, 32 | }; 33 | 34 | // serialization scratchpad 35 | let mut output = Align([MaybeUninit::::uninit(); 256]); 36 | let mut alloc = [MaybeUninit::::uninit(); 256]; 37 | 38 | let bytes = low::to_bytes_in_with_alloc::<_, _, rancor::Failure>( 39 | &value, 40 | Buffer::from(&mut *output), 41 | SubAllocator::new(&mut alloc), 42 | ) 43 | .unwrap(); 44 | 45 | let mut buffer = BytesMut::with_capacity(1024); 46 | 47 | rkyv_codec::archive_sink_bytes::(&mut buffer, &bytes[..]).unwrap(); 48 | let mut buffer = buffer.freeze(); 49 | 50 | let mut stream_buffer = AlignedVec::new(); 51 | let result = unsafe { 52 | rkyv_codec::archive_stream_bytes_unsafe::( 53 | &mut buffer, 54 | &mut stream_buffer, 55 | ) 56 | } 57 | .unwrap(); 58 | 59 | assert_eq!(&value, result); 60 | } 61 | -------------------------------------------------------------------------------- /flake.lock: -------------------------------------------------------------------------------- 1 | { 2 | "nodes": { 3 | "flake-utils": { 4 | "inputs": { 5 | "systems": "systems" 6 | }, 7 | "locked": { 8 | "lastModified": 1731533236, 9 | "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", 10 | "owner": "numtide", 11 | "repo": "flake-utils", 12 | "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", 13 | "type": "github" 14 | }, 15 | "original": { 16 | "owner": "numtide", 17 | "repo": "flake-utils", 18 | "type": "github" 19 | } 20 | }, 21 | "nixpkgs": { 22 | "locked": { 23 | "lastModified": 1748190013, 24 | "narHash": "sha256-R5HJFflOfsP5FBtk+zE8FpL8uqE7n62jqOsADvVshhE=", 25 | "owner": "NixOS", 26 | "repo": "nixpkgs", 27 | "rev": "62b852f6c6742134ade1abdd2a21685fd617a291", 28 | "type": "github" 29 | }, 30 | "original": { 31 | "owner": "NixOS", 32 | "ref": "nixos-unstable", 33 | "repo": "nixpkgs", 34 | "type": "github" 35 | } 36 | }, 37 | "nixpkgs_2": { 38 | "locked": { 39 | "lastModified": 1744536153, 40 | "narHash": "sha256-awS2zRgF4uTwrOKwwiJcByDzDOdo3Q1rPZbiHQg/N38=", 41 | "owner": "NixOS", 42 | "repo": "nixpkgs", 43 | "rev": "18dd725c29603f582cf1900e0d25f9f1063dbf11", 44 | "type": "github" 45 | }, 46 | "original": { 47 | "owner": "NixOS", 48 | "ref": "nixpkgs-unstable", 49 | "repo": "nixpkgs", 50 | "type": "github" 51 | } 52 | }, 53 | "root": { 54 | "inputs": { 55 | "flake-utils": "flake-utils", 56 | "nixpkgs": "nixpkgs", 57 | "rust-overlay": "rust-overlay" 58 | } 59 | }, 60 | "rust-overlay": { 61 | "inputs": { 62 | "nixpkgs": "nixpkgs_2" 63 | }, 64 | "locked": { 65 | "lastModified": 1748313401, 66 | "narHash": "sha256-x5UuDKP2Ui/TresAngUo9U4Ss9xfOmN8dAXU8OrkZmA=", 67 | "owner": "oxalica", 68 | "repo": "rust-overlay", 69 | "rev": "9c8ea175cf9af29edbcff121512e44092a8f37e4", 70 | "type": "github" 71 | }, 72 | "original": { 73 | "owner": "oxalica", 74 | "repo": "rust-overlay", 75 | "type": "github" 76 | } 77 | }, 78 | "systems": { 79 | "locked": { 80 | "lastModified": 1681028828, 81 | "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", 82 | "owner": "nix-systems", 83 | "repo": "default", 84 | "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", 85 | "type": "github" 86 | }, 87 | "original": { 88 | "owner": "nix-systems", 89 | "repo": "default", 90 | "type": "github" 91 | } 92 | } 93 | }, 94 | "root": "root", 95 | "version": 7 96 | } 97 | -------------------------------------------------------------------------------- /flake.nix: -------------------------------------------------------------------------------- 1 | { 2 | description = "Rust Overlay"; 3 | 4 | inputs = { 5 | nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; 6 | rust-overlay.url = "github:oxalica/rust-overlay"; 7 | flake-utils.url = "github:numtide/flake-utils"; 8 | }; 9 | 10 | outputs = { self, nixpkgs, rust-overlay, flake-utils, ... }: 11 | flake-utils.lib.eachDefaultSystem (system: 12 | let 13 | overlays = [ (import rust-overlay) ]; 14 | pkgs = import nixpkgs { inherit system overlays; }; 15 | in with pkgs; { 16 | devShells.default = mkShell { 17 | buildInputs = [ 18 | openssl 19 | pkg-config 20 | (rust-bin.fromRustupToolchainFile ./rust-toolchain.toml) 21 | ]; 22 | packages = [ 23 | (pkgs.python3.withPackages 24 | (ps: with ps; [ pandas seaborn matplotlib numpy ])) 25 | ]; 26 | }; 27 | }); 28 | } 29 | -------------------------------------------------------------------------------- /rust-toolchain.toml: -------------------------------------------------------------------------------- 1 | [toolchain] 2 | channel = "nightly" 3 | components = [ "rust-src" ] -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | hard_tabs = true 2 | attr_fn_like_width = 90 -------------------------------------------------------------------------------- /src/framed_codec.rs: -------------------------------------------------------------------------------- 1 | use std::marker::PhantomData; 2 | 3 | use asynchronous_codec::{BytesMut, Decoder, Encoder}; 4 | use bytes::Buf; 5 | use rkyv::{ 6 | Archive, Archived, Deserialize, Portable, Serialize, 7 | api::{ 8 | high::{HighDeserializer, HighSerializer, HighValidator}, 9 | serialize_using, 10 | }, 11 | de::Pool, 12 | rancor, 13 | ser::{ 14 | Serializer, 15 | allocator::{Arena, ArenaHandle}, 16 | sharing::Share, 17 | }, 18 | util::AlignedVec, 19 | }; 20 | 21 | use crate::{RkyvCodecError, length_codec::LengthCodec}; 22 | 23 | /// A futures-compatible Codec using the `asynchronous-codec` library. 24 | /// 25 | /// This struct is used to encode and decode rkyv packets. 26 | /// # Example 27 | /// 28 | /// ```rust 29 | /// # use rkyv::{Archive, Serialize, Deserialize, rancor, util::AlignedVec}; 30 | /// # use rkyv_codec::{RkyvCodec, VarintLength}; 31 | /// # use asynchronous_codec::{Framed, Decoder, Encoder}; 32 | /// # use bytes::BytesMut; 33 | /// # 34 | /// #[derive(Archive, Deserialize, Serialize, Debug, PartialEq, Clone)] 35 | /// #[rkyv(attr(derive(Debug)))] 36 | /// struct Test { 37 | /// int: u8, 38 | /// string: String, 39 | /// option: Option>, 40 | /// } 41 | /// 42 | /// let value = Test { 43 | /// int: 42, 44 | /// string: "hello world".to_string(), 45 | /// option: Some(vec![1, 2, 3, 4]), 46 | /// }; 47 | /// 48 | /// let mut codec = RkyvCodec::::default(); 49 | /// let mut buf = BytesMut::new(); 50 | /// 51 | /// // Encoding 52 | /// codec.encode(&value, &mut buf).unwrap(); 53 | /// 54 | /// // Decoding 55 | /// let decoded_value = codec.decode(&mut buf).unwrap().unwrap(); 56 | /// 57 | /// assert_eq!(value, decoded_value); 58 | /// ``` 59 | 60 | pub struct RkyvCodec { 61 | _data: PhantomData, 62 | _length: PhantomData, 63 | encode_buffer: Option, 64 | ser_arena: Arena, 65 | ser_share: Option, 66 | 67 | decode_buffer: AlignedVec, 68 | des_pool: Pool, 69 | } 70 | impl Default for RkyvCodec { 71 | fn default() -> Self { 72 | Self { 73 | _data: PhantomData, 74 | _length: PhantomData::default(), 75 | encode_buffer: Some(AlignedVec::new()), 76 | ser_arena: Arena::new(), 77 | ser_share: Some(Share::new()), 78 | decode_buffer: AlignedVec::new(), 79 | des_pool: Pool::new(), 80 | } 81 | } 82 | } 83 | /// Encoder impl encodes object streams to bytes 84 | impl Encoder for RkyvCodec 85 | where 86 | Packet: Archive + for<'b> Serialize, rancor::Error>>, 87 | { 88 | type Item<'a> = &'a Packet; 89 | type Error = RkyvCodecError; 90 | 91 | fn encode<'a>(&mut self, data: Self::Item<'a>, buf: &mut BytesMut) -> Result<(), Self::Error> { 92 | let mut encode_buffer = self.encode_buffer.take().unwrap(); 93 | let share = self.ser_share.take().unwrap(); 94 | encode_buffer.clear(); 95 | let mut serializer = Serializer::new(encode_buffer, self.ser_arena.acquire(), share); 96 | let _ = serialize_using(data, &mut serializer)?; 97 | 98 | let (encode_buffer, _, share) = serializer.into_raw_parts(); 99 | 100 | let mut length_buffer = L::Buffer::default(); 101 | let length_buffer = L::encode(encode_buffer.len(), &mut length_buffer); 102 | buf.extend_from_slice(length_buffer); 103 | buf.extend_from_slice(&encode_buffer[..]); 104 | 105 | self.encode_buffer = Some(encode_buffer); 106 | self.ser_share = Some(share); 107 | Ok(()) 108 | } 109 | } 110 | 111 | /// Decoder impl parses json objects from bytes 112 | impl Decoder for RkyvCodec 113 | where 114 | Packet: Archive + 'static, 115 | Packet::Archived: Portable 116 | + for<'b> rkyv::bytecheck::CheckBytes> 117 | + Deserialize>, 118 | { 119 | type Item = Packet; 120 | type Error = RkyvCodecError; 121 | 122 | fn decode(&mut self, buf: &mut BytesMut) -> Result, Self::Error> { 123 | if buf.is_empty() { 124 | return Ok(None); 125 | } 126 | self.decode_buffer.clear(); 127 | 128 | let (length, remaining) = L::decode(buf).map_err(RkyvCodecError::ReadLengthError)?; 129 | self.decode_buffer.extend_from_slice(&remaining[0..length]); 130 | let archive: &Archived = rkyv::access::<_, rancor::Error>(&self.decode_buffer)?; 131 | let packet: Packet = 132 | archive.deserialize(&mut HighDeserializer::wrap(&mut self.des_pool))?; 133 | 134 | // NOTE: This is the only place where I use bytes_old :( 135 | let amount_read = length + buf.len() - remaining.len(); 136 | buf.advance(amount_read); 137 | 138 | Ok(Some(packet)) 139 | } 140 | } 141 | -------------------------------------------------------------------------------- /src/length_codec.rs: -------------------------------------------------------------------------------- 1 | use core::{error, fmt}; 2 | 3 | use bytes::Buf; 4 | #[cfg(feature = "std")] 5 | use futures::{AsyncRead, AsyncReadExt}; 6 | 7 | #[cfg(feature = "std")] 8 | #[derive(Debug)] 9 | pub enum WithIOError { 10 | LengthDecodeError(E), 11 | IoError(std::io::Error), 12 | } 13 | 14 | /// Length encoding trait, allows for different kinds of length encoding 15 | pub trait LengthCodec: fmt::Debug + 'static { 16 | type Error: core::error::Error + Send + Sync + 'static; 17 | type Buffer: Default + Sized + Send; 18 | 19 | fn as_slice(buffer: &mut Self::Buffer) -> &mut [u8]; 20 | /// Encode into buffer and return part of buffer that was written to 21 | fn encode(length: usize, buffer: &mut Self::Buffer) -> &[u8]; 22 | /// Decode from buffer, fails if length is formatted incorrectly, returns length and remaining buffer 23 | fn decode(buffer: &[u8]) -> Result<(usize, &[u8]), Self::Error>; 24 | 25 | fn decode_bytes(buf: &mut impl Buf) -> Result; 26 | 27 | #[cfg(feature = "std")] 28 | #[allow(async_fn_in_trait)] 29 | async fn decode_async<'a, W: AsyncRead + Unpin>( 30 | reader: &'a mut W, 31 | ) -> Result>; 32 | } 33 | 34 | #[derive(Debug, PartialEq, Eq, thiserror::Error)] 35 | pub enum VarintLengthError { 36 | #[error("not enough bytes for varint")] 37 | Insufficient, 38 | #[error("too many bytes")] 39 | Overflow, 40 | #[error("not minimal, too many trailing zero bytes")] 41 | NotMinimal, 42 | } 43 | impl From for VarintLengthError { 44 | fn from(error: unsigned_varint::decode::Error) -> Self { 45 | match error { 46 | unsigned_varint::decode::Error::Insufficient => VarintLengthError::Insufficient, 47 | unsigned_varint::decode::Error::Overflow => VarintLengthError::Overflow, 48 | unsigned_varint::decode::Error::NotMinimal => VarintLengthError::NotMinimal, 49 | _ => unreachable!(), 50 | } 51 | } 52 | } 53 | 54 | /// Variable-bit length encoding based using unsigned_varint crate, currently can handle lengths up to 2^63 55 | #[derive(Debug)] 56 | pub struct VarintLength; 57 | impl LengthCodec for VarintLength { 58 | type Error = VarintLengthError; 59 | type Buffer = [u8; 10]; // max varint length possible 60 | 61 | #[inline] 62 | fn as_slice(buffer: &mut Self::Buffer) -> &mut [u8] { 63 | &mut buffer[..] 64 | } 65 | 66 | #[inline] 67 | fn encode(length: usize, buf: &mut Self::Buffer) -> &[u8] { 68 | unsigned_varint::encode::usize(length, buf) 69 | } 70 | 71 | #[inline] 72 | fn decode(buf: &[u8]) -> Result<(usize, &[u8]), Self::Error> { 73 | Ok(unsigned_varint::decode::usize(buf)?) 74 | } 75 | 76 | fn decode_bytes(buf: &mut impl Buf) -> Result { 77 | let mut n = 0; // number to return 78 | let mut i = 0; // what byte we are reading 79 | 80 | #[cfg(target_pointer_width = "32")] 81 | let max_bytes = 4; 82 | #[cfg(target_pointer_width = "64")] 83 | let max_bytes = 9; 84 | 85 | let mut buf = buf.take(max_bytes); 86 | 87 | while let Ok(b) = buf.try_get_u8() { 88 | let k = usize::from(b & 0x7F); // mask out first byte 89 | n |= k << (i * 7); // record that byte 90 | if unsigned_varint::decode::is_last(b) { 91 | if b == 0 && i > 0 { 92 | // If last byte (of a multi-byte varint) is zero, it could have been "more 93 | // minimally" encoded by dropping that trailing zero. 94 | return Err(VarintLengthError::NotMinimal); 95 | } 96 | return Ok(n); 97 | } 98 | if i == max_bytes { 99 | return Err(VarintLengthError::Overflow); 100 | } 101 | i += 1; 102 | } 103 | Err(VarintLengthError::Insufficient) 104 | } 105 | 106 | #[cfg(feature = "std")] 107 | async fn decode_async<'a, R: AsyncRead + Unpin>( 108 | reader: &'a mut R, 109 | ) -> Result> { 110 | let len = unsigned_varint::aio::read_usize(reader) 111 | .await 112 | .map_err(|e| match e { 113 | unsigned_varint::io::ReadError::Io(error) => WithIOError::IoError(error), 114 | unsigned_varint::io::ReadError::Decode(error) => { 115 | WithIOError::LengthDecodeError(error.into()) 116 | } 117 | _ => todo!(), 118 | })?; 119 | Ok(len) 120 | } 121 | } 122 | 123 | /// Error emitted by const-length encodings when there aren't bytes in the passed buffer when decoding 124 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 125 | pub struct NotEnoughBytesError; 126 | 127 | impl fmt::Display for NotEnoughBytesError { 128 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 129 | write!(f, "not enough bytes in buffer to decode length") 130 | } 131 | } 132 | 133 | impl error::Error for NotEnoughBytesError {} 134 | 135 | macro_rules! impl_uint_length_codec { 136 | ($name:ident, $uint_type:ty, $byte_count:expr) => { 137 | /// Big-endian fixed-size integer length encoding. 138 | #[derive(Debug)] 139 | pub struct $name; 140 | 141 | impl LengthCodec for $name { 142 | type Error = NotEnoughBytesError; 143 | type Buffer = [u8; $byte_count]; 144 | 145 | #[inline] 146 | fn as_slice(buffer: &mut Self::Buffer) -> &mut [u8] { 147 | &mut buffer[..] 148 | } 149 | 150 | #[inline] 151 | fn encode(length: usize, buf: &mut Self::Buffer) -> &[u8] { 152 | *buf = <$uint_type>::to_be_bytes(length as $uint_type); 153 | &buf[..] 154 | } 155 | 156 | #[inline] 157 | fn decode(buf: &[u8]) -> Result<(usize, &[u8]), Self::Error> { 158 | if buf.len() < $byte_count { 159 | return Err(NotEnoughBytesError); 160 | } 161 | let (int_bytes, rest) = buf.split_at($byte_count); 162 | let bytes: [u8; $byte_count] = 163 | int_bytes.try_into().map_err(|_| NotEnoughBytesError)?; 164 | Ok((<$uint_type>::from_be_bytes(bytes) as usize, rest)) 165 | } 166 | 167 | #[inline] 168 | fn decode_bytes(buf: &mut impl Buf) -> Result { 169 | let mut buffer = Self::Buffer::default(); 170 | buf.copy_to_slice(&mut buffer); 171 | Self::decode(&buffer).map(|(len, _)| len) 172 | } 173 | 174 | #[cfg(feature = "std")] 175 | async fn decode_async<'a, R: AsyncRead + Unpin>( 176 | reader: &'a mut R, 177 | ) -> Result> { 178 | let mut buffer = Self::Buffer::default(); 179 | reader 180 | .read_exact(Self::as_slice(&mut buffer)) 181 | .await 182 | .map_err(|e| { 183 | if e.kind() == std::io::ErrorKind::UnexpectedEof { 184 | WithIOError::LengthDecodeError(NotEnoughBytesError) 185 | } else { 186 | WithIOError::IoError(e) 187 | } 188 | })?; 189 | Self::decode(&buffer) 190 | .map(|(len, _)| len) 191 | .map_err(WithIOError::LengthDecodeError) 192 | } 193 | } 194 | }; 195 | } 196 | 197 | // Generate U8Length 198 | impl_uint_length_codec!(U8Length, u8, 1); 199 | impl_uint_length_codec!(U16Length, u16, 2); 200 | impl_uint_length_codec!(U32Length, u32, 4); 201 | impl_uint_length_codec!(U64Length, u64, 8); 202 | 203 | #[cfg(feature = "std")] 204 | #[cfg(test)] 205 | mod tests { 206 | use super::*; 207 | use paste::paste; // For concatenating identifiers in macro 208 | 209 | macro_rules! impl_uint_length_codec_tests { 210 | ($test_suffix:ident, $codec_type:ty, $test_value:expr, $encoded_bytes:expr, $byte_count:expr) => { 211 | paste! { 212 | #[test] 213 | fn []() { 214 | let mut buf = <$codec_type as LengthCodec>::Buffer::default(); 215 | let encoded = <$codec_type>::encode($test_value, &mut buf); 216 | assert_eq!(encoded, $encoded_bytes); 217 | let (decoded, rest) = <$codec_type>::decode(encoded).unwrap(); 218 | assert_eq!(decoded, $test_value); 219 | assert!(rest.is_empty()); 220 | 221 | // Test decoding with extra bytes 222 | let mut extended_bytes: Vec = ($encoded_bytes).to_vec(); 223 | extended_bytes.extend_from_slice(&[1, 2, 3]); 224 | let (decoded_partial, rest_partial) = <$codec_type>::decode(&extended_bytes).unwrap(); 225 | assert_eq!(decoded_partial, $test_value); 226 | assert_eq!(rest_partial, &[1, 2, 3]); 227 | 228 | // Test decoding insufficient bytes 229 | if $byte_count > 0 { 230 | let short_bytes = &$encoded_bytes[..$byte_count -1]; 231 | assert_eq!(<$codec_type>::decode(short_bytes), Err(NotEnoughBytesError)); 232 | } 233 | assert_eq!(<$codec_type>::decode(&[]), Err(NotEnoughBytesError)); 234 | } 235 | 236 | #[async_std::test] 237 | async fn []() { 238 | let data = $encoded_bytes; 239 | let mut reader = &data[..]; 240 | let len = <$codec_type>::decode_async(&mut reader).await.unwrap(); 241 | assert_eq!(len, $test_value); 242 | 243 | if $byte_count > 0 { 244 | let data_short = &$encoded_bytes[..$byte_count-1]; 245 | let mut reader_short = &data_short[..]; 246 | let err = <$codec_type>::decode_async(&mut reader_short) 247 | .await 248 | .unwrap_err(); 249 | match err { 250 | WithIOError::LengthDecodeError(NotEnoughBytesError) => {} 251 | _ => panic!("Unexpected error type for short data: {:?}", err), 252 | } 253 | } 254 | 255 | let data_empty: [u8;0] = []; 256 | let mut reader_empty = &data_empty[..]; 257 | let err_empty = <$codec_type>::decode_async(&mut reader_empty).await.unwrap_err(); 258 | match err_empty { 259 | WithIOError::LengthDecodeError(NotEnoughBytesError) => {} 260 | _ => panic!("Unexpected error type for empty data: {:?}", err_empty), 261 | } 262 | } 263 | } 264 | }; 265 | } 266 | 267 | impl_uint_length_codec_tests!(u8, U8Length, 42, &[42], 1); 268 | impl_uint_length_codec_tests!(u16, U16Length, 300, &[1, 44], 2); // 0x012C 269 | impl_uint_length_codec_tests!(u32, U32Length, 70000, &[0, 1, 17, 112], 4); // 0x00011170 270 | impl_uint_length_codec_tests!( 271 | u64, 272 | U64Length, 273 | 1_000_000_000_000, 274 | &[0, 0, 0, 232, 212, 165, 16, 0], 275 | 8 276 | ); // 0x000000E8_D4A51000 277 | } 278 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! This crate provides zero-copy deserialization for packet streaming using `rkyv`. 2 | #![cfg_attr( 3 | feature = "std", 4 | doc = r##" 5 | Simple usage example: 6 | ```rust 7 | # async_std::task::block_on(async { 8 | use rkyv::{Archived, util::AlignedVec, Archive, Serialize, Deserialize, rancor}; 9 | use rkyv_codec::{archive_stream, RkyvWriter, VarintLength}; 10 | use futures::SinkExt; 11 | #[derive(Archive, Deserialize, Serialize, Debug, PartialEq, Clone)] 12 | #[rkyv(derive(Debug), compare(PartialEq))] 13 | struct Test { 14 | int: u8, 15 | string: String, 16 | option: Option>, 17 | } 18 | let value = Test { 19 | int: 42, 20 | string: "hello world".to_string(), 21 | option: Some(vec![1, 2, 3, 4]), 22 | }; 23 | // Writing 24 | let writer = Vec::new(); 25 | let mut codec = RkyvWriter::<_, VarintLength>::new(writer); 26 | codec.send(&value).await.unwrap(); 27 | // Reading 28 | let mut reader = &codec.inner()[..]; 29 | let mut buffer = AlignedVec::new(); // Aligned streaming buffer for re-use 30 | let value_archived: &Archived = archive_stream::<_, Test, VarintLength>(&mut reader, &mut buffer).await.unwrap(); // This returns a reference into the passed buffer 31 | // can deserialize as normal as well (or do *partial* deserialization for blazingly fast speeds!) 32 | let value_deserialized: Test = rkyv::deserialize::<_, rancor::Error>(value_archived).unwrap(); 33 | assert_eq!(value, *value_archived); 34 | assert_eq!(value, value_deserialized); 35 | # }) 36 | ``` 37 | "## 38 | )] 39 | #![cfg_attr(not(feature = "std"), no_std)] 40 | #![cfg_attr(test, feature(test))] 41 | 42 | /// Abstract length encodings for reading and writing streams 43 | mod length_codec; 44 | pub use length_codec::LengthCodec; 45 | pub use length_codec::{U8Length, U16Length, U32Length, U64Length, VarintLength}; 46 | 47 | #[cfg(feature = "std")] 48 | use length_codec::WithIOError; 49 | 50 | // Asynchronous IO (framed read/writing) 51 | #[cfg(feature = "std")] 52 | mod framed_codec; 53 | #[cfg(feature = "std")] 54 | pub use framed_codec::*; 55 | 56 | /// Error type for rkyv_codec 57 | use thiserror::Error; 58 | #[derive(Debug, Error)] 59 | pub enum RkyvCodecError { 60 | #[cfg(feature = "std")] 61 | #[error(transparent)] 62 | IoError(#[from] futures::io::Error), 63 | #[error("packet not correctly archived: {0}")] 64 | CheckArchiveError(#[from] rkyv::rancor::Error), 65 | #[error("Failed to parse length: {0}")] 66 | ReadLengthError(L::Error), 67 | #[error("Premature End of Buffer Error")] 68 | LengthTooLong { requested: usize, available: usize }, 69 | } 70 | 71 | #[cfg(feature = "std")] 72 | impl From> for RkyvCodecError { 73 | fn from(value: WithIOError) -> Self { 74 | match value { 75 | WithIOError::IoError(err) => RkyvCodecError::IoError(err), 76 | WithIOError::LengthDecodeError(err) => RkyvCodecError::ReadLengthError(err), 77 | } 78 | } 79 | } 80 | 81 | #[cfg(feature = "std")] 82 | mod rkyv_codec; 83 | #[cfg(feature = "std")] 84 | pub use rkyv_codec::*; 85 | 86 | use bytes::{Buf, BufMut}; 87 | #[cfg(feature = "std")] 88 | use rkyv::api::high as cur_api; 89 | #[cfg(feature = "std")] 90 | use rkyv::api::high::HighValidator as CurrentValidator; 91 | #[cfg(not(feature = "std"))] 92 | use rkyv::api::low as cur_api; 93 | #[cfg(not(feature = "std"))] 94 | use rkyv::api::low::LowValidator as CurrentValidator; 95 | 96 | use rkyv::{Archive, Archived, bytecheck::CheckBytes, rancor, util::AlignedVec}; 97 | 98 | /// Writes a single `Object` from a `bytes::Bytes` 99 | pub fn archive_sink_bytes( 100 | bytes: &mut impl BufMut, 101 | archived: &[u8], 102 | ) -> Result<(), RkyvCodecError> { 103 | let length_buf = &mut L::Buffer::default(); 104 | let length_buf = L::encode(archived.len(), length_buf); 105 | bytes.put(length_buf); 106 | bytes.put(archived); 107 | Ok(()) 108 | } 109 | /// Reads a single `&Archived` from a `bytes::Bytes` into the passed buffer 110 | /// # Safety 111 | /// This will cause undefined behavior if the bytestream is not the correct format (i.e. not generated through `archive_sink[_bytes]`, `RkyvWriter`, or `RkyvCodec`) with the correct LengthCodec 112 | pub unsafe fn archive_stream_bytes_unsafe<'b, Packet: Archive, L: LengthCodec>( 113 | bytes: &mut impl Buf, 114 | buffer: &'b mut AlignedVec, 115 | ) -> Result<&'b Archived, RkyvCodecError> { 116 | // Read length 117 | let archive_len = L::decode_bytes(bytes).map_err(RkyvCodecError::ReadLengthError)?; 118 | 119 | // reserve length 120 | buffer.reserve_exact(archive_len.saturating_sub(buffer.len())); 121 | 122 | // Safety: we assume that copy_to_slice will only write (not read) and thus even if buffer is uninit, this should be fine. 123 | unsafe { 124 | buffer.set_len(archive_len); 125 | } 126 | 127 | // Read specific amount into aligned buffer 128 | bytes.copy_to_slice(&mut buffer[..archive_len]); 129 | 130 | // reinterpret cast 131 | let archive = unsafe { rkyv::access_unchecked::(buffer) }; 132 | Ok(archive) 133 | } 134 | /// Reads a single `&Archived` from a `bytes::Bytes` into the passed buffer if validation enabled. 135 | pub fn archive_stream_bytes<'b, Packet: Archive, L: LengthCodec>( 136 | bytes: &mut impl Buf, 137 | buffer: &'b mut AlignedVec, 138 | ) -> Result<&'b Archived, RkyvCodecError> 139 | where 140 | ::Archived: for<'a> CheckBytes>, 141 | { 142 | // Read length 143 | let archive_len = L::decode_bytes(bytes).map_err(RkyvCodecError::ReadLengthError)?; 144 | 145 | // reserve length 146 | buffer.reserve_exact(archive_len.saturating_sub(buffer.len())); 147 | 148 | // Safety: we assume that copy_to_slice will only write (not read) and thus even if buffer is uninit, this should be fine. 149 | unsafe { 150 | buffer.set_len(archive_len); 151 | } 152 | 153 | // Read specific amount into aligned buffer 154 | bytes 155 | .try_copy_to_slice(buffer.as_mut_slice()) 156 | .map_err(|e| RkyvCodecError::LengthTooLong { 157 | requested: e.requested, 158 | available: e.available, 159 | })?; 160 | 161 | let archive = cur_api::access::, rancor::Error>(&*buffer)?; 162 | Ok(archive) 163 | } 164 | 165 | #[cfg(test)] 166 | #[cfg(feature = "std")] 167 | mod tests { 168 | extern crate test; 169 | 170 | use async_std::task::block_on; 171 | use asynchronous_codec::{CborCodec, Framed}; 172 | use bytes::BytesMut; 173 | use futures::{AsyncRead, AsyncWrite, SinkExt, StreamExt, io::Cursor}; 174 | use rkyv::{Archive, Archived, Deserialize, Serialize, rancor, to_bytes, util::AlignedVec}; 175 | 176 | use crate::archive_stream; 177 | use crate::{ 178 | RkyvWriter, archive_sink, archive_sink_bytes, archive_stream_bytes, 179 | length_codec::{self, LengthCodec, U64Length, VarintLength}, 180 | }; 181 | 182 | type TestLengthCodec = length_codec::VarintLength; 183 | 184 | #[derive( 185 | Archive, 186 | Deserialize, 187 | Serialize, 188 | Debug, 189 | PartialEq, 190 | Clone, 191 | serde::Serialize, 192 | serde::Deserialize, 193 | )] 194 | // This will generate a PartialEq impl between our unarchived and archived types 195 | // To use the safe API, you have to use the check_bytes option for the archive 196 | #[rkyv(compare(PartialEq), attr(derive(Debug)))] 197 | struct Test { 198 | int: u8, 199 | string: String, 200 | option: Option>, 201 | } 202 | 203 | lazy_static::lazy_static! { 204 | static ref TEST: Test = Test { 205 | int: 42, 206 | string: "hello world".to_string(), 207 | option: Some(vec![1, 2, 3, 4]), 208 | }; 209 | static ref TEST_BYTES: &'static [u8] = { 210 | let vec = rkyv::to_bytes::(&*TEST).unwrap(); 211 | Box::leak(vec.into_boxed_slice()) 212 | }; 213 | static ref TEST_ARCHIVED: &'static Archived = unsafe { rkyv::access_unchecked::>(*TEST_BYTES) }; 214 | } 215 | 216 | #[inline] 217 | async fn gen_amount(writer: &mut W, count: usize) { 218 | for _ in 0..count { 219 | archive_sink::<_, L>(writer, *TEST_BYTES).await.unwrap(); 220 | } 221 | } 222 | #[inline] 223 | async fn consume_amount(mut reader: R, count: usize) { 224 | let mut buffer = AlignedVec::new(); 225 | for _ in 0..count { 226 | let value = archive_stream::<_, Test, L>(&mut reader, &mut buffer) 227 | .await 228 | .unwrap(); 229 | assert_eq!(*TEST, *value); 230 | } 231 | } 232 | #[test] 233 | fn bytes_functions() { 234 | let mut writer = BytesMut::new(); 235 | for _ in 0..50 { 236 | archive_sink_bytes::(&mut writer, *TEST_BYTES).unwrap(); 237 | } 238 | 239 | let mut reader = writer.freeze(); 240 | 241 | let mut buffer = AlignedVec::with_capacity(256); 242 | 243 | for _ in 0..50 { 244 | let data: &Archived = 245 | archive_stream_bytes::(&mut reader, &mut buffer).unwrap(); 246 | assert_eq!(*TEST, *data); 247 | } 248 | } 249 | 250 | #[async_std::test] 251 | async fn functions() { 252 | let mut writer = Vec::new(); 253 | archive_sink::<_, TestLengthCodec>(&mut writer, *TEST_BYTES) 254 | .await 255 | .unwrap(); 256 | 257 | let mut reader = &writer[..]; 258 | 259 | let mut buffer = AlignedVec::with_capacity(256); 260 | let data: &Archived = 261 | archive_stream::<_, Test, TestLengthCodec>(&mut reader, &mut buffer) 262 | .await 263 | .unwrap(); 264 | 265 | let value_sent: Test = rkyv::deserialize::<_, rancor::Error>(data).unwrap(); 266 | 267 | assert_eq!(*TEST, value_sent); 268 | } 269 | 270 | #[async_std::test] 271 | /// Tests the edgecase where we hit EOF while reading the length because the whole structure fits within the 10 byte max varint length. 272 | async fn functions_varint_edge_case() { 273 | #[derive(Archive, Serialize, Deserialize, Debug, PartialEq, Clone)] 274 | #[rkyv(compare(PartialEq), attr(derive(Debug)))] 275 | struct SmallTest { 276 | int: u8, 277 | } 278 | const SMALL_TEST: SmallTest = SmallTest { int: 1 }; 279 | let mut writer = Vec::new(); 280 | let mut sink = RkyvWriter::<_, TestLengthCodec>::new(&mut writer); 281 | for _ in 0..2 { 282 | sink.send(&SMALL_TEST).await.unwrap(); 283 | } 284 | 285 | let mut reader = &writer[..]; 286 | println!("reader: {reader:?}"); 287 | 288 | let mut buffer = AlignedVec::with_capacity(256); 289 | let data: &Archived = 290 | archive_stream::<_, SmallTest, VarintLength>(&mut reader, &mut buffer) 291 | .await 292 | .unwrap(); 293 | 294 | let value_sent: SmallTest = rkyv::deserialize::<_, rancor::Error>(data).unwrap(); 295 | 296 | assert_eq!(SMALL_TEST, value_sent); 297 | 298 | println!("reader: {reader:?}"); 299 | 300 | let data: &Archived = 301 | archive_stream::<_, SmallTest, VarintLength>(&mut reader, &mut buffer) 302 | .await 303 | .unwrap(); 304 | 305 | let value_sent: SmallTest = rkyv::deserialize::<_, rancor::Error>(data).unwrap(); 306 | 307 | assert_eq!(SMALL_TEST, value_sent); 308 | } 309 | 310 | #[async_std::test] 311 | async fn rkyv_writer() { 312 | let mut writer = Vec::new(); 313 | let mut sink = RkyvWriter::<_, TestLengthCodec>::new(&mut writer); 314 | sink.send(&*TEST).await.unwrap(); 315 | 316 | let mut reader = &writer[..]; 317 | 318 | let mut buffer = AlignedVec::with_capacity(256); 319 | let data: &Archived = 320 | archive_stream::<_, Test, TestLengthCodec>(&mut reader, &mut buffer) 321 | .await 322 | .unwrap(); 323 | 324 | assert_eq!(*TEST, *data); 325 | } 326 | 327 | #[async_std::test] 328 | 329 | async fn futures_ser_de() { 330 | let codec = crate::framed_codec::RkyvCodec::::default(); 331 | let mut buffer = vec![0u8; 256]; 332 | let mut framed = asynchronous_codec::Framed::new(Cursor::new(&mut buffer), codec); 333 | framed.send(&*TEST).await.unwrap(); 334 | 335 | let codec = framed.into_parts().codec; 336 | 337 | let mut framed = asynchronous_codec::Framed::new(Cursor::new(&mut buffer), codec); 338 | let received_value = framed.next().await.unwrap().unwrap(); 339 | 340 | assert_eq!(*TEST, received_value); 341 | } 342 | #[async_std::test] 343 | async fn futures_cbor_ser_de() { 344 | let codec = CborCodec::::new(); 345 | 346 | let mut buffer = vec![0u8; 256]; 347 | let mut framed = Framed::new(Cursor::new(&mut buffer), codec); 348 | 349 | framed.send(TEST.clone()).await.unwrap(); 350 | 351 | let codec = framed.into_parts().codec; 352 | 353 | let mut framed = asynchronous_codec::Framed::new(Cursor::new(&mut buffer), codec); 354 | let received_value = framed.next().await.unwrap().unwrap(); 355 | 356 | assert_eq!(*TEST, received_value); 357 | } 358 | 359 | use test::Bencher; 360 | 361 | #[bench] 362 | fn bench_varint_length_encoding(b: &mut Bencher) { 363 | let mut buffer = Vec::with_capacity(1024); 364 | b.iter(|| { 365 | block_on(async { 366 | buffer.clear(); 367 | gen_amount::<_, VarintLength>(&mut buffer, 50).await; 368 | consume_amount::<_, VarintLength>(&mut &buffer[..], 50).await; 369 | }) 370 | }) 371 | } 372 | #[bench] 373 | fn bench_u64_length_encoding(b: &mut Bencher) { 374 | let mut buffer = Vec::with_capacity(1024); 375 | 376 | b.iter(|| { 377 | block_on(async { 378 | buffer.clear(); 379 | gen_amount::<_, U64Length>(&mut buffer, 50).await; 380 | consume_amount::<_, U64Length>(&mut &buffer[..], 50).await; 381 | }) 382 | }) 383 | } 384 | #[bench] 385 | fn bench_archive_sink_prearchived_50(b: &mut Bencher) { 386 | let mut buffer = Vec::with_capacity(1024); 387 | b.iter(|| { 388 | block_on(async { 389 | buffer.clear(); 390 | for _ in 0..50 { 391 | archive_sink::<_, TestLengthCodec>(&mut buffer, &*TEST_BYTES) 392 | .await 393 | .unwrap() 394 | } 395 | }) 396 | }); 397 | block_on(consume_amount::<_, TestLengthCodec>(&mut &buffer[..], 50)); 398 | } 399 | #[bench] 400 | fn bench_archive_sink_50(b: &mut Bencher) { 401 | let mut buffer = Vec::with_capacity(1024); 402 | b.iter(|| { 403 | block_on(async { 404 | buffer.clear(); 405 | 406 | for _ in 0..50 { 407 | let bytes = to_bytes::(&*TEST).unwrap(); // This makes it very slow 408 | archive_sink::<_, TestLengthCodec>(&mut buffer, &bytes) 409 | .await 410 | .unwrap(); 411 | } 412 | }) 413 | }); 414 | block_on(consume_amount::<_, TestLengthCodec>(&mut &buffer[..], 50)) 415 | } 416 | #[bench] 417 | fn bench_rkyv_writer_50(b: &mut Bencher) { 418 | let mut buffer = Vec::with_capacity(1024); 419 | b.iter(|| { 420 | block_on(async { 421 | buffer.clear(); 422 | let mut sink = RkyvWriter::<_, TestLengthCodec>::new(&mut buffer); 423 | for _ in 0..50 { 424 | sink.send(&*TEST).await.unwrap(); 425 | } 426 | }) 427 | }); 428 | block_on(consume_amount::<_, TestLengthCodec>(&mut &buffer[..], 50)) 429 | } 430 | #[bench] 431 | fn bench_archive_stream_50(b: &mut Bencher) { 432 | let mut buffer = Vec::with_capacity(1024); 433 | 434 | block_on(gen_amount::<_, TestLengthCodec>(&mut buffer, 50)); 435 | 436 | b.iter(|| { 437 | block_on(async { 438 | consume_amount::<_, TestLengthCodec>(&mut &buffer[..], 50).await; 439 | }) 440 | }); 441 | } 442 | 443 | #[bench] 444 | fn bench_rkyv_asynchronous_codec_sink_50(b: &mut Bencher) { 445 | let mut buffer = Vec::with_capacity(1024); 446 | b.iter(|| { 447 | block_on(async { 448 | buffer.clear(); 449 | let codec = crate::framed_codec::RkyvCodec::::default(); 450 | let mut framed = asynchronous_codec::Framed::new(Cursor::new(&mut buffer), codec); 451 | for _ in 0..50 { 452 | framed.send(&TEST).await.unwrap(); 453 | } 454 | }) 455 | }); 456 | block_on(consume_amount::<_, TestLengthCodec>(&mut &buffer[..], 50)); 457 | } 458 | 459 | #[bench] 460 | fn bench_rkyv_asynchronous_codec_stream_50(b: &mut Bencher) { 461 | use futures::TryStreamExt; 462 | let mut buffer = Vec::with_capacity(1024); 463 | 464 | block_on(gen_amount::<_, TestLengthCodec>(&mut buffer, 50)); 465 | 466 | let codec = crate::framed_codec::RkyvCodec::::default(); 467 | let mut framed = asynchronous_codec::Framed::new(Cursor::new(&mut buffer), codec); 468 | b.iter(|| { 469 | block_on(async { 470 | framed.set_position(0); 471 | while let Some(_value) = framed.try_next().await.unwrap() {} 472 | }) 473 | }); 474 | } 475 | 476 | #[bench] 477 | fn bench_futures_cbor_sink_50(b: &mut Bencher) { 478 | let mut buffer = vec![0u8; 256]; 479 | 480 | b.iter(|| { 481 | block_on(async { 482 | buffer.clear(); 483 | let codec = CborCodec::::new(); 484 | let mut framed = Framed::new(Cursor::new(&mut buffer), codec); 485 | framed.set_position(0); 486 | for _ in 0..50 { 487 | framed.send(TEST.clone()).await.unwrap(); 488 | } 489 | }) 490 | }); 491 | } 492 | #[bench] 493 | fn bench_futures_cbor_stream_50(b: &mut Bencher) { 494 | let codec = CborCodec::::new(); 495 | 496 | let mut buffer = vec![0u8; 256]; 497 | let mut framed = Framed::new(Cursor::new(&mut buffer), codec); 498 | 499 | block_on(async { 500 | for _ in 0..50 { 501 | framed.send(TEST.clone()).await.unwrap(); 502 | } 503 | }); 504 | 505 | let codec = framed.into_parts().codec; 506 | 507 | let mut framed = asynchronous_codec::Framed::new(Cursor::new(&mut buffer), codec); 508 | 509 | b.iter(|| { 510 | block_on(async { 511 | framed.set_position(0); 512 | while let Some(value) = framed.next().await { 513 | test::black_box(value.unwrap()); 514 | } 515 | }) 516 | }); 517 | } 518 | } 519 | -------------------------------------------------------------------------------- /src/rkyv_codec.rs: -------------------------------------------------------------------------------- 1 | use futures::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, Sink, ready}; 2 | use std::{ 3 | ops::Range, 4 | pin::Pin, 5 | task::{Context, Poll}, 6 | }; 7 | 8 | use pin_project::pin_project; 9 | 10 | use rkyv::{ 11 | Archive, Archived, Portable, Serialize, 12 | api::{ 13 | high::{HighSerializer, HighValidator}, 14 | serialize_using, 15 | }, 16 | rancor, 17 | ser::{ 18 | Serializer, 19 | allocator::{Arena, ArenaHandle}, 20 | sharing::Share, 21 | }, 22 | util::AlignedVec, 23 | }; 24 | 25 | use crate::{RkyvCodecError, length_codec::LengthCodec}; 26 | 27 | /// Rewrites a single buffer representing an Archive to an `AsyncWrite` 28 | pub async fn archive_sink<'b, Inner: AsyncWrite + Unpin, L: LengthCodec>( 29 | inner: &mut Inner, 30 | archived: &[u8], 31 | ) -> Result<(), RkyvCodecError> { 32 | let length_buf = &mut L::Buffer::default(); 33 | let length_buf = L::encode(archived.len(), length_buf); 34 | inner.write_all(length_buf).await?; 35 | inner.write_all(archived).await?; 36 | Ok(()) 37 | } 38 | /// Reads a single `&Archived` from an `AsyncRead` without checking for correct byte formatting 39 | /// # Safety 40 | /// This may cause undefined behavior if the bytestream is not a valid archive (i.e. not generated through `archive_sink[_bytes]`, or `RkyvWriter`) 41 | /// 42 | /// As an optimisation, this function may pass uninitialized bytes to the reader for the reader to read into. Make sure the particular reader in question is implemented correctly and does not read from its passed buffer in the poll_read() function without first writing to it. 43 | /// # Warning 44 | /// Passed buffer is reallocated so it may fit the size of the packet being written. This may allow for DOS attacks if remote sends too large a length encoding 45 | /// # Errors 46 | /// Will return an error if there are not enough bytes to read to read the length of the packet, or read the packet itself. Will also return an error if the length encoding format is invalid. 47 | pub async unsafe fn archive_stream_unsafe< 48 | 'b, 49 | Inner: AsyncRead + Unpin, 50 | Packet: Archive + Portable + 'b, 51 | L: LengthCodec, 52 | >( 53 | inner: &mut Inner, 54 | buffer: &'b mut AlignedVec, 55 | ) -> Result<&'b Archived, RkyvCodecError> { 56 | buffer.clear(); 57 | 58 | // parse archive length 59 | let archive_len = L::decode_async(inner).await?; 60 | 61 | // If not enough capacity in buffer to fit `archive_len`, reserve more. 62 | if buffer.capacity() < archive_len { 63 | buffer.reserve(archive_len - buffer.capacity()) 64 | } 65 | 66 | // Safety: Caller should make sure that reader does not read from this potentially uninitialized buffer passed to poll_read() 67 | unsafe { buffer.set_len(archive_len) } 68 | 69 | // read exactly amount specified by archive_len into buffer 70 | inner.read_exact(buffer).await?; 71 | 72 | // Safety: Caller should make sure that reader does not produce invalid packets. 73 | unsafe { Ok(rkyv::access_unchecked(buffer)) } 74 | } 75 | 76 | /// Reads a single `&Archived` from an `AsyncRead` using the passed buffer. 77 | /// 78 | /// Until streaming iterators (and streaming futures) are implemented in rust, this currently the fastest method I could come up with that requires no recurring heap allocations. 79 | /// 80 | /// Requires rkyv "validation" feature 81 | /// # Safety 82 | /// As an optimisation, this function may pass uninitialized bytes to the reader for the reader to read into. Make sure the particular reader in question is implemented correctly and does not read from its passed buffer in the poll_read() function without first writing to it. 83 | /// # Warning 84 | /// Passed buffer is reallocated so it may fit the size of the packet being written. This may allow for DOS attacks if remote sends too large a length encoding 85 | /// # Errors 86 | /// Will return an error if there are not enough bytes to read to read the length of the packet, or read the packet itself. Will also return an error if the length encoding format is invalid or the packet archive itself is invalid. 87 | pub async fn archive_stream<'b, Inner: AsyncRead + Unpin, Packet, L: LengthCodec>( 88 | inner: &mut Inner, 89 | buffer: &'b mut AlignedVec, 90 | ) -> Result<&'b Archived, RkyvCodecError> 91 | where 92 | Packet: rkyv::Archive + 'b, 93 | Packet::Archived: for<'a> rkyv::bytecheck::CheckBytes>, 94 | { 95 | buffer.clear(); 96 | 97 | let archive_len = L::decode_async(inner).await?; 98 | 99 | // If not enough capacity in buffer to fit `archive_len`, reserve more. 100 | if buffer.capacity() < archive_len { 101 | buffer.reserve(archive_len - buffer.capacity()) 102 | } 103 | 104 | // Safety: Caller should make sure that reader does not read from this potentially uninitialized buffer passed to poll_read() 105 | unsafe { buffer.set_len(archive_len) } 106 | 107 | inner.read_exact(buffer).await?; 108 | 109 | let archive = rkyv::access::(buffer)?; 110 | 111 | Ok(archive) 112 | } 113 | 114 | /// Wraps an `AsyncWrite` and implements `Sink` to serialize `Archive` objects. 115 | #[pin_project] 116 | pub struct RkyvWriter { 117 | #[pin] 118 | writer: Writer, 119 | length_buffer: L::Buffer, 120 | len_state: Range, // How much of the length buffer has been written 121 | buf_state: usize, // Whether or not the aligned buf is being written and if so, how much so far 122 | buffer: Option, 123 | arena: Arena, 124 | share: Option, 125 | } 126 | 127 | // Safety: Arena is Send and Share is Send, if Writer is Send RkyvWriter should be Send. 128 | unsafe impl Send for RkyvWriter {} 129 | 130 | impl RkyvWriter { 131 | pub fn new(writer: Writer) -> Self { 132 | Self { 133 | writer, 134 | length_buffer: L::Buffer::default(), 135 | len_state: Default::default(), 136 | buf_state: 0, 137 | buffer: Some(AlignedVec::new()), 138 | arena: Arena::new(), 139 | share: Some(Share::new()), 140 | } 141 | } 142 | pub fn inner(self) -> Writer { 143 | self.writer 144 | } 145 | } 146 | 147 | impl Sink<&Packet> 148 | for RkyvWriter 149 | where 150 | Packet: Archive + for<'b> Serialize, rancor::Error>>, 151 | { 152 | type Error = RkyvCodecError; 153 | 154 | fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 155 | self.project() 156 | .writer 157 | .poll_flush(cx) 158 | .map_err(RkyvCodecError::IoError) 159 | } 160 | 161 | fn start_send(self: Pin<&mut Self>, item: &Packet) -> Result<(), Self::Error> { 162 | let this = self.project(); 163 | let buffer_len = { 164 | // Serializer 165 | let mut buffer = this.buffer.take().unwrap(); 166 | buffer.clear(); 167 | let share = this.share.take().unwrap(); 168 | let mut serializer = Serializer::new(buffer, this.arena.acquire(), share); 169 | // serialize 170 | let _ = serialize_using(item, &mut serializer)?; 171 | 172 | let (buffer, _, share) = serializer.into_raw_parts(); 173 | let buffer_len = buffer.len(); 174 | *this.buffer = Some(buffer); 175 | *this.share = Some(share); 176 | buffer_len 177 | }; 178 | 179 | *this.len_state = 0..L::encode(buffer_len, this.length_buffer).len(); 180 | *this.buf_state = 0; 181 | 182 | Ok(()) 183 | } 184 | 185 | fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 186 | let mut this = self.project(); 187 | 188 | // keep writing length buffer for as long as is required 189 | let len_state = this.len_state; 190 | if len_state.start <= len_state.end { 191 | let length_buffer = L::as_slice(this.length_buffer); 192 | let length_buffer = &mut length_buffer[len_state.clone()]; 193 | 194 | let written = ready!(Pin::new(&mut this.writer).poll_write(cx, length_buffer)?); 195 | len_state.start += written; 196 | } 197 | let buffer = this.buffer.take().unwrap(); 198 | 199 | while *this.buf_state < buffer.len() { 200 | let buffer_left = &buffer[*this.buf_state..buffer.len()]; 201 | let bytes_written = ready!(Pin::new(&mut this.writer).poll_write(cx, buffer_left))?; 202 | if bytes_written == 0 { 203 | return Poll::Ready(Err(RkyvCodecError::LengthTooLong { 204 | requested: buffer.capacity(), 205 | available: buffer.len(), 206 | })); 207 | } 208 | *this.buf_state += bytes_written; 209 | } 210 | 211 | *this.buffer = Some(buffer); 212 | 213 | ready!(this.writer.poll_flush(cx)?); 214 | Poll::Ready(Ok(())) 215 | } 216 | 217 | fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 218 | self.project() 219 | .writer 220 | .poll_close(cx) 221 | .map_err(RkyvCodecError::IoError) 222 | } 223 | } 224 | --------------------------------------------------------------------------------