RUST-871 Serialize directly to BSON bytes in insert operations #406
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,10 +1,15 @@ | ||
| use std::{ | ||
| convert::{TryFrom, TryInto}, | ||
| io::{Read, Write}, | ||
| time::Duration, | ||
| }; | ||
|
|
||
| use bson::spec::ElementType; | ||
| use serde::{de::Error as SerdeDeError, ser, Deserialize, Deserializer, Serialize, Serializer}; | ||
|
|
||
| use crate::{ | ||
| bson::{doc, Bson, Document}, | ||
| error::{Error, ErrorKind, Result}, | ||
| runtime::{SyncLittleEndianRead, SyncLittleEndianWrite}, | ||
| }; | ||
|
|
||
|
|
@@ -164,128 +169,30 @@ where | |
| .ok_or_else(|| D::Error::custom(format!("could not deserialize u64 from {:?}", bson))) | ||
| } | ||
|
|
||
|
There was a problem hiding this comment. These were no longer needed as we now just use the length of the serialized vec. |
||
| /// The size in bytes of the provided document's entry in a BSON array at the given index. | ||
| pub(crate) fn array_entry_size_bytes(index: usize, doc_len: usize) -> u64 { | ||
| // | ||
| // * type (1 byte) | ||
| // * number of decimal digits in key | ||
| // * null terminator for the key (1 byte) | ||
| // * size of value | ||
|
|
||
| 1 + num_decimal_digits(index) + 1 + doc_len as u64 | ||
| } | ||
|
|
||
| /// The number of digits in `n` in base 10. | ||
| /// Useful for calculating the size of an array entry in BSON. | ||
| fn num_decimal_digits(mut n: usize) -> u64 { | ||
|
There was a problem hiding this comment. the old implementation of this method was incorrect actually so I rewrote it. |
||
| let mut digits = 0; | ||
|
|
||
| loop { | ||
| n /= 10; | ||
| digits += 1; | ||
|
|
||
| if n == 0 { | ||
| return digits; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// Read a document's raw BSON bytes from the provided reader. | ||
|
|
@@ -300,63 +207,161 @@ pub(crate) fn read_document_bytes<R: Read>(mut reader: R) -> Result<Vec<u8>> { | |
| Ok(bytes) | ||
| } | ||
|
|
||
| /// Get the value for the provided key from a buffer containing a BSON document. | ||
| /// If the key is not present, None will be returned. | ||
| /// If the BSON is not properly formatted, an internal error would be returned. | ||
| /// | ||
| /// TODO: RUST-924 replace this with raw document API usage. | ||
| pub(crate) fn raw_get(doc: &[u8], key: &str) -> Result<Option<Bson>> { | ||
|
There was a problem hiding this comment. Unfortunately this function needs to live in the driver for now as we need to be able to comb through the serialized command for |
||
| fn read_i32(reader: &mut std::io::Cursor<&[u8]>) -> Result<i32> { | ||
| reader.read_i32().map_err(deserialize_error) | ||
| } | ||
|
|
||
| fn read_u8(reader: &mut std::io::Cursor<&[u8]>) -> Result<u8> { | ||
| reader.read_u8().map_err(deserialize_error) | ||
| } | ||
|
|
||
| fn deserialize_error<T: std::error::Error>(_e: T) -> Error { | ||
| deserialize_error_no_arg() | ||
| } | ||
|
|
||
| fn deserialize_error_no_arg() -> Error { | ||
| Error::from(ErrorKind::Internal { | ||
| message: "failed to read from serialized document".to_string(), | ||
| }) | ||
| } | ||
|
There was a problem hiding this comment. The documents queried here should always be valid as they're coming straight out of our serializer, so any error encountered here is mapped to an internal error. |
||
|
|
||
| let mut reader = std::io::Cursor::new(doc); | ||
| let len: u64 = read_i32(&mut reader)? | ||
| .try_into() | ||
| .map_err(deserialize_error)?; | ||
|
|
||
| while reader.position() < len { | ||
| let element_start: usize = reader.position().try_into().map_err(deserialize_error)?; | ||
|
|
||
| // read the element type | ||
| let tag = read_u8(&mut reader)?; | ||
|
|
||
| // check if we reached the end of the document | ||
| if tag == 0 && reader.position() == len { | ||
| return Ok(None); | ||
| } | ||
|
|
||
| let element_type = ElementType::from(tag).ok_or_else(deserialize_error_no_arg)?; | ||
|
|
||
| // walk through the document until a null byte is encountered | ||
| while read_u8(&mut reader)? != 0 { | ||
| if reader.position() >= len { | ||
| return Err(deserialize_error_no_arg()); | ||
| } | ||
| } | ||
|
|
||
| // parse the key | ||
| let string_end: usize = reader | ||
| .position() | ||
| .checked_sub(1) // back from null byte | ||
| .and_then(|u| usize::try_from(u).ok()) | ||
| .ok_or_else(deserialize_error_no_arg)?; | ||
| let slice = &reader.get_ref()[(element_start + 1)..string_end]; | ||
| let k = std::str::from_utf8(slice).map_err(deserialize_error)?; | ||
|
|
||
| // move to the end of the element | ||
| let skip_len = match element_type { | ||
| ElementType::Array | ||
| | ElementType::EmbeddedDocument | ||
| | ElementType::JavaScriptCodeWithScope => { | ||
| let l = read_i32(&mut reader)?; | ||
| // length includes the 4 bytes for the length, so subtrack them out | ||
| l.checked_sub(4).ok_or_else(deserialize_error_no_arg)? | ||
| } | ||
| ElementType::Binary => read_i32(&mut reader)? | ||
| .checked_add(1) // add one for subtype | ||
| .ok_or_else(deserialize_error_no_arg)?, | ||
| ElementType::Int32 => 4, | ||
| ElementType::Int64 => 8, | ||
| ElementType::String | ElementType::Symbol | ElementType::JavaScriptCode => { | ||
| read_i32(&mut reader)? | ||
| } | ||
| ElementType::Boolean => 1, | ||
| ElementType::Double => 8, | ||
| ElementType::Timestamp => 8, | ||
| ElementType::Decimal128 => 16, | ||
| ElementType::MinKey | ||
| | ElementType::MaxKey | ||
| | ElementType::Null | ||
| | ElementType::Undefined => 0, | ||
| ElementType::DateTime => 8, | ||
| ElementType::ObjectId => 12, | ||
| ElementType::DbPointer => read_i32(&mut reader)? | ||
| .checked_add(12) // add 12 for objectid | ||
| .ok_or_else(deserialize_error_no_arg)?, | ||
| ElementType::RegularExpression => { | ||
| // read two cstr's | ||
| for _i in 0..2 { | ||
| while read_u8(&mut reader)? != 0 { | ||
| if reader.position() >= len { | ||
| return Err(deserialize_error_no_arg()); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| 0 // don't need to skip anymore since we already read the whole value | ||
| } | ||
| }; | ||
| let skip_len: u64 = skip_len.try_into().map_err(deserialize_error)?; | ||
| reader.set_position( | ||
| reader | ||
| .position() | ||
| .checked_add(skip_len) | ||
| .ok_or_else(deserialize_error_no_arg)?, | ||
| ); | ||
|
|
||
| if k == key { | ||
| // if this is the element we're looking for, extract it. | ||
| let element_end: usize = reader.position().try_into().map_err(deserialize_error)?; | ||
| let element_slice = &reader.get_ref()[element_start..element_end]; | ||
| let element_length: i32 = element_slice.len().try_into().map_err(deserialize_error)?; | ||
|
|
||
| // create a new temporary document which just has the element we want and grab the value | ||
| let mut temp_doc = Vec::new(); | ||
|
|
||
| // write the document length | ||
| let temp_len: i32 = element_length | ||
| .checked_add(4 + 1) | ||
| .ok_or_else(deserialize_error_no_arg)?; | ||
| temp_doc | ||
| .write_all(&temp_len.to_le_bytes()) | ||
| .map_err(deserialize_error)?; | ||
|
|
||
| // add in the element | ||
| temp_doc.extend(element_slice); | ||
|
|
||
| // write the null byte | ||
| temp_doc.push(0); | ||
|
|
||
| let d = Document::from_reader(temp_doc.as_slice()).map_err(deserialize_error)?; | ||
| return Ok(Some( | ||
| d.get("_id").ok_or_else(deserialize_error_no_arg)?.clone(), | ||
| )); | ||
| } | ||
| } | ||
|
|
||
| // read all bytes but didn't reach null byte | ||
| Err(deserialize_error_no_arg()) | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod test { | ||
| use crate::bson_util::num_decimal_digits; | ||
|
|
||
| #[cfg_attr(feature = "tokio-runtime", tokio::test)] | ||
| #[cfg_attr(feature = "async-std-runtime", async_std::test)] | ||
| async fn num_digits() { | ||
| assert_eq!(num_decimal_digits(0), 1); | ||
| assert_eq!(num_decimal_digits(1), 1); | ||
| assert_eq!(num_decimal_digits(10), 2); | ||
| assert_eq!(num_decimal_digits(15), 2); | ||
| assert_eq!(num_decimal_digits(100), 3); | ||
| assert_eq!(num_decimal_digits(125), 3); | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Since the benchmarks are single threaded, this ends up giving a nice boost by avoiding the tokio overhead of managing a thread pool.