Skip to content

Commit a960ca0

Browse files
gaius-qigaius-qi
authored
feat: add TryFrom and From for vortex (#40)
Signed-off-by: Gaius <gaius.qi@gmail.com>
1 parent d2005fa commit a960ca0

File tree

4 files changed

+97
-72
lines changed

4 files changed

+97
-72
lines changed

‎benches/vortex.rs‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ fn vortex_packet_serialization(c: &mut Criterion) {
7373
let value = Bytes::from("a".repeat(32));
7474
let packet = Vortex::new(tag, value).unwrap();
7575
b.iter(|| {
76-
black_box(packet.to_bytes());
76+
black_box(Bytes::from(packet.clone()));
7777
});
7878
});
7979

@@ -82,7 +82,7 @@ fn vortex_packet_serialization(c: &mut Criterion) {
8282
let value = Bytes::from("a".repeat(1024));
8383
let packet = Vortex::new(tag, value).unwrap();
8484
b.iter(|| {
85-
black_box(packet.to_bytes());
85+
black_box(Bytes::from(packet.clone()));
8686
});
8787
});
8888

@@ -91,7 +91,7 @@ fn vortex_packet_serialization(c: &mut Criterion) {
9191
let value = Bytes::from("a".repeat(1024 * 1024));
9292
let packet = Vortex::new(tag, value).unwrap();
9393
b.iter(|| {
94-
black_box(packet.to_bytes());
94+
black_box(Bytes::from(packet.clone()));
9595
});
9696
});
9797

@@ -100,7 +100,7 @@ fn vortex_packet_serialization(c: &mut Criterion) {
100100
let value = Bytes::from("a".repeat(16 * 1024 * 1024));
101101
let packet = Vortex::new(tag, value).unwrap();
102102
b.iter(|| {
103-
black_box(packet.to_bytes());
103+
black_box(Bytes::from(packet.clone()));
104104
});
105105
});
106106

@@ -109,7 +109,7 @@ fn vortex_packet_serialization(c: &mut Criterion) {
109109
let value = Bytes::from("a".repeat(128 * 1024 * 1024));
110110
let packet = Vortex::new(tag, value).unwrap();
111111
b.iter(|| {
112-
black_box(packet.to_bytes());
112+
black_box(Bytes::from(packet.clone()));
113113
});
114114
});
115115

@@ -123,49 +123,49 @@ fn vortex_packet_deserialization(c: &mut Criterion) {
123123
let tag = Tag::PieceContent;
124124
let value = Bytes::from("a".repeat(32));
125125
let packet = Vortex::new(tag, value).unwrap();
126-
let bytes = packet.to_bytes();
126+
let bytes: Bytes = packet.into();
127127
b.iter(|| {
128-
black_box(Vortex::from_bytes(bytes.clone())).unwrap();
128+
black_box(Vortex::try_from(bytes.clone()).unwrap());
129129
});
130130
});
131131

132132
group.bench_function("Deserialize Medium Packet (1 KB)", |b| {
133133
let tag = Tag::PieceContent;
134134
let value = Bytes::from("a".repeat(1024));
135135
let packet = Vortex::new(tag, value).unwrap();
136-
let bytes = packet.to_bytes();
136+
let bytes: Bytes = packet.into();
137137
b.iter(|| {
138-
black_box(Vortex::from_bytes(bytes.clone())).unwrap();
138+
black_box(Vortex::try_from(bytes.clone()).unwrap());
139139
});
140140
});
141141

142142
group.bench_function("Deserialize Large Packet (1 MB)", |b| {
143143
let tag = Tag::PieceContent;
144144
let value = Bytes::from("a".repeat(1024 * 1024));
145145
let packet = Vortex::new(tag, value).unwrap();
146-
let bytes = packet.to_bytes();
146+
let bytes: Bytes = packet.into();
147147
b.iter(|| {
148-
black_box(Vortex::from_bytes(bytes.clone())).unwrap();
148+
black_box(Vortex::try_from(bytes.clone()).unwrap());
149149
});
150150
});
151151

152152
group.bench_function("Deserialize XLarge Packet (16 MB)", |b| {
153153
let tag = Tag::PieceContent;
154154
let value = Bytes::from("a".repeat(16 * 1024 * 1024));
155155
let packet = Vortex::new(tag, value).unwrap();
156-
let bytes = packet.to_bytes();
156+
let bytes: Bytes = packet.into();
157157
b.iter(|| {
158-
black_box(Vortex::from_bytes(bytes.clone())).unwrap();
158+
black_box(Vortex::try_from(bytes.clone()).unwrap());
159159
});
160160
});
161161

162162
group.bench_function("Deserialize XXLarge Packet (128 MB)", |b| {
163163
let tag = Tag::PieceContent;
164164
let value = Bytes::from("a".repeat(128 * 1024 * 1024));
165165
let packet = Vortex::new(tag, value).unwrap();
166-
let bytes = packet.to_bytes();
166+
let bytes: Bytes = packet.into();
167167
b.iter(|| {
168-
black_box(Vortex::from_bytes(bytes.clone())).unwrap();
168+
black_box(Vortex::try_from(bytes.clone()).unwrap());
169169
});
170170
});
171171

‎src/lib.rs‎

Lines changed: 51 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -23,15 +23,15 @@ pub mod tlv;
2323

2424
/// HEADER_SIZE is the size of the Vortex packet header including the packet identifier, tag, and
2525
/// length.
26-
const HEADER_SIZE: usize = 6;
26+
pub const HEADER_SIZE: usize = 6;
2727

2828
/// MAX_VALUE_SIZE is the maximum size of the value field (4 GiB).
2929
const MAX_VALUE_SIZE: usize = 4 * 1024 * 1024 * 1024;
3030

3131
/// Header represents the Vortex packet header.
32-
#[derive(Debug)]
32+
#[derive(Debug, Clone)]
3333
pub struct Header {
34-
packet_id: u8,
34+
id: u8,
3535
tag: tlv::Tag,
3636
length: usize,
3737
}
@@ -58,7 +58,7 @@ pub struct Header {
5858
/// ```
5959
///
6060
/// For more information, please refer to the [Vortex Protocol](https://github.com/dragonflyoss/vortex/blob/main/docs/README.md).
61-
#[derive(Debug)]
61+
#[derive(Debug, Clone)]
6262
pub enum Vortex {
6363
DownloadPiece(Header, tlv::download_piece::DownloadPiece),
6464
PieceContent(Header, tlv::piece_content::PieceContent),
@@ -73,23 +73,23 @@ impl Vortex {
7373
pub fn new(tag: tlv::Tag, value: Bytes) -> Result<Self> {
7474
let mut rng = thread_rng();
7575
let header = Header {
76-
packet_id: rng.gen(),
76+
id: rng.gen(),
7777
tag,
7878
length: value.len(),
7979
};
8080

8181
(tag, header, value).try_into()
8282
}
8383

84-
/// packet_id returns the packet identifier of the Vortex packet.
84+
/// id returns the packet identifier of the Vortex packet.
8585
#[inline]
86-
pub fn packet_id(&self) -> u8 {
86+
pub fn id(&self) -> u8 {
8787
match self {
88-
Vortex::DownloadPiece(header, _) => header.packet_id,
89-
Vortex::PieceContent(header, _) => header.packet_id,
90-
Vortex::Reserved(header) => header.packet_id,
91-
Vortex::Close(header) => header.packet_id,
92-
Vortex::Error(header, _) => header.packet_id,
88+
Vortex::DownloadPiece(header, _) => header.id,
89+
Vortex::PieceContent(header, _) => header.id,
90+
Vortex::Reserved(header) => header.id,
91+
Vortex::Close(header) => header.id,
92+
Vortex::Error(header, _) => header.id,
9393
}
9494
}
9595

@@ -117,8 +117,25 @@ impl Vortex {
117117
}
118118
}
119119

120-
/// from_bytes creates a Vortex packet from a byte slice.
121-
pub fn from_bytes(bytes: Bytes) -> Result<Self> {
120+
/// header returns a reference to the packet header.
121+
#[inline]
122+
pub fn header(&self) -> &Header {
123+
match self {
124+
Vortex::DownloadPiece(header, _) => header,
125+
Vortex::PieceContent(header, _) => header,
126+
Vortex::Reserved(header) => header,
127+
Vortex::Close(header) => header,
128+
Vortex::Error(header, _) => header,
129+
}
130+
}
131+
}
132+
133+
/// Implement TryFrom<Bytes> for Vortex.
134+
impl TryFrom<Bytes> for Vortex {
135+
type Error = Error;
136+
137+
/// try_from converts a Bytes into a Vortex packet.
138+
fn try_from(bytes: Bytes) -> Result<Self> {
122139
if bytes.len() < HEADER_SIZE {
123140
return Err(Error::InvalidPacket(format!(
124141
"expected min {HEADER_SIZE} bytes, got {}",
@@ -129,7 +146,7 @@ impl Vortex {
129146
let mut bytes = BytesMut::from(bytes);
130147
let header = bytes.split_to(HEADER_SIZE);
131148
let value = bytes.freeze();
132-
let packet_id = header[0];
149+
let id = header[0];
133150
let tag = header[1]
134151
.try_into()
135152
.map_err(|err| Error::InvalidPacket(format!("invalid tag value: {:?}", err)))?;
@@ -144,21 +161,15 @@ impl Vortex {
144161
)));
145162
}
146163

147-
(
148-
tag,
149-
Header {
150-
packet_id,
151-
tag,
152-
length,
153-
},
154-
value,
155-
)
156-
.try_into()
164+
(tag, Header { id, tag, length }, value).try_into()
157165
}
166+
}
158167

159-
/// to_bytes converts the Vortex packet to a byte slice.
160-
pub fn to_bytes(&self) -> Bytes {
161-
let (header, value) = match self {
168+
/// Implement From<PieceContent> for Bytes.
169+
impl From<Vortex> for Bytes {
170+
/// from converts a Vortex packet to Bytes.
171+
fn from(packet: Vortex) -> Self {
172+
let (header, value) = match packet {
162173
Vortex::DownloadPiece(header, download_piece) => {
163174
(header, Into::<Bytes>::into(download_piece.clone()))
164175
}
@@ -171,7 +182,7 @@ impl Vortex {
171182
};
172183

173184
let mut bytes = BytesMut::with_capacity(HEADER_SIZE + value.len());
174-
bytes.put_u8(header.packet_id);
185+
bytes.put_u8(header.id);
175186
bytes.put_u8(header.tag.into());
176187
bytes.put_u32(value.len() as u32);
177188
bytes.extend_from_slice(&value);
@@ -216,7 +227,7 @@ mod tests {
216227
let value = Bytes::from("a".repeat(32) + "-42");
217228
let packet = Vortex::new(tag, value.clone()).expect("Failed to create Vortex packet");
218229

219-
assert_eq!(packet.packet_id(), packet.packet_id());
230+
assert_eq!(packet.id(), packet.id());
220231
assert_eq!(packet.tag(), &tag);
221232
assert_eq!(packet.length(), value.len());
222233
}
@@ -232,23 +243,26 @@ mod tests {
232243
}
233244

234245
#[test]
235-
fn test_from_bytes() {
246+
fn test_vortex_from_bytes() {
236247
let value = b"test data";
237248
let packet = Vortex::new(Tag::PieceContent, Bytes::from_static(value))
238249
.expect("Failed to create packet");
239-
let bytes = packet.to_bytes();
240-
let deserialized = Vortex::from_bytes(bytes).expect("Failed to deserialize packet");
250+
let tag = *packet.tag();
251+
let length = packet.length();
252+
253+
let bytes: Bytes = packet.into();
254+
let deserialized: Vortex = bytes.try_into().expect("Failed to deserialize packet");
241255

242-
assert_eq!(packet.tag(), deserialized.tag());
243-
assert_eq!(packet.length(), deserialized.length());
256+
assert_eq!(tag, deserialized.tag().clone());
257+
assert_eq!(length, deserialized.length());
244258
}
245259

246260
#[test]
247-
fn test_to_bytes() {
261+
fn test_vortex_to_bytes() {
248262
let value = b"test data";
249263
let packet = Vortex::new(Tag::PieceContent, Bytes::from_static(value))
250264
.expect("Failed to create packet");
251-
let bytes = packet.to_bytes();
265+
let bytes: Bytes = packet.into();
252266

253267
assert_eq!(bytes.len(), HEADER_SIZE + value.len());
254268
}

‎tests/integration_tests.rs‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use bytes::{BufMut, BytesMut};
1+
use bytes::{BufMut, Bytes, BytesMut};
22
use std::sync::{Arc, Mutex};
33
use vortex_protocol::error::Error;
44
use vortex_protocol::tlv::Tag;
@@ -21,10 +21,11 @@ impl MockPeer {
2121
self.pieces.lock().unwrap().push(piece);
2222
}
2323

24-
pub fn handle_packet(&self, packet: &Vortex) -> Result<Vortex, Error> {
24+
pub fn handle_packet(&self, packet: Vortex) -> Result<Vortex, Error> {
2525
match packet.tag() {
2626
Tag::DownloadPiece => {
27-
let value = String::from_utf8(packet.to_bytes()[6..].to_vec())?;
27+
let bytes: Bytes = packet.into();
28+
let value = String::from_utf8(bytes[6..].to_vec())?;
2829
let parts: Vec<&str> = value.split('-').collect();
2930
if parts.len() != 2 {
3031
return Err(Error::InvalidPacket(
@@ -66,11 +67,13 @@ fn test_piece_download_flow() {
6667
.unwrap();
6768

6869
// Handle the request.
69-
let response = peer.handle_packet(&request).unwrap();
70+
let response = peer.handle_packet(request).unwrap();
71+
let tag = *response.tag();
72+
let response_bytes: Bytes = response.into();
7073

7174
// Verify the response.
72-
assert_eq!(response.tag(), &Tag::PieceContent);
73-
assert_eq!(&response.to_bytes()[6..], &piece);
75+
assert_eq!(&tag, &Tag::PieceContent);
76+
assert_eq!(&response_bytes[6..], &piece);
7477
}
7578

7679
#[test]
@@ -84,25 +87,25 @@ fn test_error_propagation() {
8487
format!("{}-42", task_id).into_bytes().into(),
8588
)
8689
.unwrap();
87-
let result = peer.handle_packet(&request);
90+
let result = peer.handle_packet(request);
8891
assert!(matches!(result, Err(Error::InvalidPacket(_))));
8992

9093
// Test with unexpected tag.
9194
let request = Vortex::new(Tag::PieceContent, vec![1, 2, 3, 4].into()).unwrap();
92-
let result = peer.handle_packet(&request);
95+
let result = peer.handle_packet(request);
9396
assert!(matches!(result, Err(Error::InvalidPacket(_))));
9497
}
9598

9699
#[test]
97100
fn test_invalid_length() {
98101
// Create a packet with invalid length in the header.
99102
let mut packet_bytes = BytesMut::with_capacity(6);
100-
packet_bytes.put_u8(42); // packet_id
103+
packet_bytes.put_u8(42); // id
101104
packet_bytes.put_u8(Tag::PieceContent.into()); // tag
102105
packet_bytes.put_u32(u32::MAX); // length (too large)
103106

104107
// Attempt to parse the packet with invalid length.
105-
let result = Vortex::from_bytes(packet_bytes.freeze());
108+
let result = Vortex::try_from(packet_bytes.freeze());
106109
assert!(matches!(result, Err(Error::InvalidLength(_))));
107110
}
108111

0 commit comments

Comments
 (0)