Skip to content

Commit 4ef542c

Browse files
levkklevkk
authored
Use Bytes for Bind fields (#559)
* Use Bytes for Bind fields * Fix warning * comment
1 parent 682180d commit 4ef542c

File tree

8 files changed

+47
-35
lines changed

8 files changed

+47
-35
lines changed

‎integration/rust/tests/integration/savepoint.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use rust::setup::{admin_sqlx, connections_sqlx};
1+
use rust::setup::connections_sqlx;
22
use sqlx::Executor;
33

44
#[tokio::test]

‎pgdog/src/backend/server.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1041,7 +1041,7 @@ pub mod test {
10411041
"",
10421042
&[Parameter {
10431043
len: 1,
1044-
data: "1".as_bytes().to_vec(),
1044+
data: "1".as_bytes().into(),
10451045
}],
10461046
&[Format::Text],
10471047
);
@@ -1086,7 +1086,7 @@ pub mod test {
10861086
&name,
10871087
&[Parameter {
10881088
len: 1,
1089-
data: "1".as_bytes().to_vec(),
1089+
data: "1".as_bytes().into(),
10901090
}],
10911091
);
10921092

@@ -1172,7 +1172,7 @@ pub mod test {
11721172
"__pgdog_1",
11731173
&[Parameter {
11741174
len: 1,
1175-
data: "1".as_bytes().to_vec(),
1175+
data: "1".as_bytes().into(),
11761176
}],
11771177
)),
11781178
Execute::new().into(),
@@ -1368,7 +1368,7 @@ pub mod test {
13681368
"test_1",
13691369
&[crate::net::bind::Parameter {
13701370
len: 1,
1371-
data: "1".as_bytes().to_vec(),
1371+
data: "1".as_bytes().into(),
13721372
}],
13731373
)
13741374
.into(),
@@ -1429,7 +1429,7 @@ pub mod test {
14291429
"test",
14301430
&[crate::net::bind::Parameter {
14311431
len: 1,
1432-
data: "1".as_bytes().to_vec(),
1432+
data: "1".as_bytes().into(),
14331433
}],
14341434
)
14351435
.into(),
@@ -1494,7 +1494,7 @@ pub mod test {
14941494
"test",
14951495
&[crate::net::bind::Parameter {
14961496
len: 1,
1497-
data: "1".as_bytes().to_vec(),
1497+
data: "1".as_bytes().into(),
14981498
}],
14991499
)
15001500
.into(),
@@ -1927,7 +1927,7 @@ pub mod test {
19271927
"",
19281928
&[Parameter {
19291929
len: 4,
1930-
data: "1234".as_bytes().to_vec(),
1930+
data: "1234".as_bytes().into(),
19311931
}],
19321932
)
19331933
.into(),

‎pgdog/src/frontend/client/test/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -479,7 +479,7 @@ async fn test_transaction_state() {
479479
"test",
480480
&[Parameter {
481481
len: 1,
482-
data: "1".as_bytes().to_vec(),
482+
data: "1".as_bytes().into(),
483483
}],
484484
)
485485
},

‎pgdog/src/frontend/router/parser/insert.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ mod test {
130130
use crate::config::ShardedTable;
131131
use crate::net::bind::Parameter;
132132
use crate::net::Format;
133+
use bytes::Bytes;
133134

134135
use super::super::Value;
135136
use super::*;
@@ -243,7 +244,7 @@ mod test {
243244
"",
244245
&[Parameter {
245246
len: 1,
246-
data: "3".as_bytes().to_vec(),
247+
data: "3".as_bytes().into(),
247248
}],
248249
);
249250

@@ -254,7 +255,7 @@ mod test {
254255
"",
255256
&[Parameter {
256257
len: 8,
257-
data: 234_i64.to_be_bytes().to_vec(),
258+
data: Bytes::copy_from_slice(&234_i64.to_be_bytes()),
258259
}],
259260
&[Format::Binary],
260261
);

‎pgdog/src/frontend/router/parser/query/explain.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ mod tests {
5252
use crate::frontend::{ClientRequest, PreparedStatements, RouterContext};
5353
use crate::net::messages::{Bind, Parameter, Parse, Query};
5454
use crate::net::Parameters;
55+
use bytes::Bytes;
5556
use std::sync::Once;
5657

5758
fn enable_expanded_explain() {
@@ -88,7 +89,7 @@ mod tests {
8889
.iter()
8990
.map(|v| Parameter {
9091
len: v.len() as i32,
91-
data: v.to_vec(),
92+
data: Bytes::copy_from_slice(v),
9293
})
9394
.collect::<Vec<_>>();
9495

‎pgdog/src/frontend/router/parser/query/test.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use crate::{
88
Close, Format, Sync,
99
},
1010
};
11+
use bytes::Bytes;
1112

1213
use super::{super::Shard, *};
1314
use crate::backend::Cluster;
@@ -99,7 +100,7 @@ macro_rules! parse {
99100
.into_iter()
100101
.map(|p| Parameter {
101102
len: p.len() as i32,
102-
data: p.to_vec(),
103+
data: Bytes::copy_from_slice(&p),
103104
})
104105
.collect::<Vec<_>>();
105106
let bind = Bind::new_params_codes($name, &params, $codes);

‎pgdog/src/frontend/router/sharding/test/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ async fn test_binary_encoding() {
124124
"",
125125
&[Parameter {
126126
len: 5,
127-
data: "test1".as_bytes().to_vec(),
127+
data: "test1".as_bytes().into(),
128128
}],
129129
&[Format::Binary],
130130
&[1],

‎pgdog/src/net/messages/bind.rs‎

Lines changed: 30 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use super::prelude::*;
77
use super::Error;
88
use super::FromDataType;
99
use super::Vector;
10+
use bytes::BytesMut;
1011

1112
use std::fmt::Debug;
1213
use std::str::from_utf8;
@@ -34,7 +35,7 @@ pub struct Parameter {
3435
/// Parameter data length.
3536
pub len: i32,
3637
/// Parameter data.
37-
pub data: Vec<u8>,
38+
pub data: Bytes,
3839
}
3940

4041
impl Debug for Parameter {
@@ -59,14 +60,14 @@ impl Parameter {
5960
pub fn new_null() -> Self {
6061
Self {
6162
len: -1,
62-
data: vec![],
63+
data: Bytes::new(),
6364
}
6465
}
6566

6667
pub fn new(data: &[u8]) -> Self {
6768
Self {
6869
len: data.len() as i32,
69-
data: data.to_vec(),
70+
data: Bytes::copy_from_slice(data),
7071
}
7172
}
7273
}
@@ -124,8 +125,8 @@ pub struct Bind {
124125
codes: Vec<Format>,
125126
/// Parameters.
126127
params: Vec<Parameter>,
127-
/// Results format.
128-
results: Vec<i16>,
128+
/// Results format (raw bytes, 2 bytes per i16).
129+
results: Bytes,
129130
/// Original payload.
130131
original: Option<Bytes>,
131132
}
@@ -137,7 +138,7 @@ impl Default for Bind {
137138
statement: Bytes::from("\0"),
138139
codes: vec![],
139140
params: vec![],
140-
results: vec![],
141+
results: Bytes::new(),
141142
original: None,
142143
}
143144
}
@@ -149,7 +150,7 @@ impl Bind {
149150
+ self.statement.len()
150151
+ self.codes.len() * std::mem::size_of::<i16>() + 2 // num codes
151152
+ self.params.iter().map(|p| p.len()).sum::<usize>() + 2 // num params
152-
+ self.results.len() * std::mem::size_of::<i16>() + 2 // num results
153+
+ self.results.len() + 2 // num results (results already stores raw bytes)
153154
+ 4 // len
154155
+ 1 // code
155156
}
@@ -237,7 +238,11 @@ impl Bind {
237238
results: &[i16],
238239
) -> Self {
239240
let mut me = Self::new_params_codes(name, params, codes);
240-
me.results = results.to_vec();
241+
let mut buf = BytesMut::with_capacity(results.len() * 2);
242+
for result in results {
243+
buf.put_i16(*result);
244+
}
245+
me.results = buf.freeze();
241246

242247
me
243248
}
@@ -274,17 +279,19 @@ impl FromBytes for Bind {
274279
.map(|_| {
275280
let len = bytes.get_i32();
276281
let data = if len >= 0 {
277-
let mut data = Vec::with_capacity(len as usize);
278-
(0..len).for_each(|_| data.push(bytes.get_u8()));
279-
data
282+
bytes.split_to(len as usize)
280283
} else {
281-
vec![]
284+
Bytes::new()
282285
};
283286
Parameter { len, data }
284287
})
285288
.collect();
286289
let num_results = bytes.get_i16();
287-
let results = (0..num_results).map(|_| bytes.get_i16()).collect();
290+
let results = if num_results > 0 {
291+
bytes.split_to((num_results * 2) as usize)
292+
} else {
293+
Bytes::new()
294+
};
288295

289296
Ok(Self {
290297
portal,
@@ -321,10 +328,8 @@ impl ToBytes for Bind {
321328
payload.put_i32(param.len);
322329
payload.put(&param.data[..]);
323330
}
324-
payload.put_i16(self.results.len() as i16);
325-
for result in &self.results {
326-
payload.put_i16(*result);
327-
}
331+
payload.put_i16((self.results.len() / 2) as i16);
332+
payload.put(self.results.clone());
328333
Ok(payload.freeze())
329334
}
330335
}
@@ -358,14 +363,18 @@ mod test {
358363
params: vec![
359364
Parameter {
360365
len: 2,
361-
data: vec![0, 1],
366+
data: Bytes::copy_from_slice(&[0, 1]),
362367
},
363368
Parameter {
364369
len: 4,
365-
data: "test".as_bytes().to_vec(),
370+
data: Bytes::from("test"),
366371
},
367372
],
368-
results: vec![0],
373+
results: {
374+
let mut buf = BytesMut::with_capacity(2);
375+
buf.put_i16(0);
376+
buf.freeze()
377+
},
369378
};
370379
let bytes = bind.to_bytes().unwrap();
371380
let mut original = Bind::from_bytes(bytes.clone()).unwrap();
@@ -400,7 +409,7 @@ mod test {
400409
statement: "test\0".into(),
401410
codes: vec![Format::Binary],
402411
params: vec![Parameter {
403-
data: jsonb.as_bytes().to_vec(),
412+
data: Bytes::copy_from_slice(jsonb.as_bytes()),
404413
len: jsonb.len() as i32,
405414
}],
406415
..Default::default()

0 commit comments

Comments
 (0)