RUST-611 Add back code mistakenly lost in rebase

saghm · saghm · commit 8a7f03b5a5a0 · 2021-01-06T15:33:22.000-05:00
diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
@@ -6,7 +6,7 @@ edition = "2018"
 
 [features]
 default = ["tokio-runtime"]
-tokio-runtime = ["tokio/macros", "tokio/rt-core", "tokio/rt-threaded"]
+tokio-runtime = ["tokio/fs", "tokio/macros", "tokio/rt-core", "tokio/rt-threaded"]
 async-std-runtime = ["async-std"]
 
 [dependencies]
@@ -17,6 +17,7 @@ clap = "2.33.3"
 indicatif = "0.15.0"
 async-trait = "0.1.41"
 tokio = { version = "0.2.23", features = ["sync"] }
-async-std = { version = "1.6.2", optional = true, features = ["attributes"] }
+# "unstable" feature is needed for `spawn_blocking`, which is only used in task setup
+async-std = { version = "=1.6.2", optional = true, features = ["attributes", "unstable"] }
 futures = "0.3.8"
 anyhow = "1.0.34"
diff --git a/benchmarks/README.md b/benchmarks/README.md
@@ -34,7 +34,6 @@ the single-doc benchmarks. By default, all benchmarks are executed. The table be
 
 Note that in order to compare against the other drivers, an inMemory mongod instance should be used.
 
-At this point, BSON and GridFS benchmarks are not implemented because we do not own the Rust BSON library, and GridFS has not been implemented
-in the driver.
+At this point, GridFS benchmarks are not implemented because it has not been implemented in the driver.
 
 Also note that the parallel benchmarks are implemented to mirror the C++ driver's interpretation of the spec.
diff --git a/benchmarks/src/bench/find_many.rs b/benchmarks/src/bench/find_many.rs
@@ -1,11 +1,14 @@
-use std::{convert::TryInto, fs::File, path::PathBuf};
+use std::{convert::TryInto, path::PathBuf};
 
 use anyhow::{bail, Result};
 use futures::stream::StreamExt;
 use mongodb::{bson::Bson, Client, Collection, Database};
 use serde_json::Value;
 
-use crate::bench::{Benchmark, COLL_NAME, DATABASE_NAME};
+use crate::{
+    bench::{Benchmark, COLL_NAME, DATABASE_NAME},
+    fs::read_to_string,
+};
 
 pub struct FindManyBenchmark {
     db: Database,
@@ -30,9 +33,9 @@ impl Benchmark for FindManyBenchmark {
 
         let num_iter = options.num_iter;
 
-        let mut file = spawn_blocking_and_await!(File::open(options.path))?;
+        let mut file = read_to_string(&options.path).await?;
 
-        let json: Value = spawn_blocking_and_await!(serde_json::from_reader(&mut file))?;
+        let json: Value = serde_json::from_str(&mut file)?;
         let doc = match json.try_into()? {
             Bson::Document(doc) => doc,
             _ => bail!("invalid json test file"),
diff --git a/benchmarks/src/bench/find_one.rs b/benchmarks/src/bench/find_one.rs
@@ -1,4 +1,4 @@
-use std::{convert::TryInto, fs::File, path::PathBuf};
+use std::{convert::TryInto, path::PathBuf};
 
 use anyhow::{bail, Result};
 use mongodb::{
@@ -9,7 +9,10 @@ use mongodb::{
 };
 use serde_json::Value;
 
-use crate::bench::{Benchmark, COLL_NAME, DATABASE_NAME};
+use crate::{
+    bench::{Benchmark, COLL_NAME, DATABASE_NAME},
+    fs::read_to_string,
+};
 
 pub struct FindOneBenchmark {
     db: Database,
@@ -35,9 +38,9 @@ impl Benchmark for FindOneBenchmark {
 
         let num_iter = options.num_iter;
 
-        let mut file = spawn_blocking_and_await!(File::open(options.path))?;
+        let mut file = read_to_string(&options.path).await?;
 
-        let json: Value = spawn_blocking_and_await!(serde_json::from_reader(&mut file))?;
+        let json: Value = serde_json::from_str(&mut file)?;
         let mut doc = match json.try_into()? {
             Bson::Document(doc) => doc,
             _ => bail!("invalid json test file"),
diff --git a/benchmarks/src/bench/insert_many.rs b/benchmarks/src/bench/insert_many.rs
@@ -36,8 +36,13 @@ impl Benchmark for InsertManyBenchmark {
 
         let num_copies = options.num_copies;
 
+        // This benchmark uses a file that's quite large, and unfortunately `serde_json` has no
+        // async version of `from_reader`, so rather than read the whole file into memory at once,
+        // we use the runtime's `spawn_blocking` functionality to do this efficiently.
+        //
+        // Note that the setup is _not_ measured as part of the benchmark runtime, so even if
+        // `spawn_blocking` turned out not to be super efficient, it wouldn't be a big deal.
         let mut file = spawn_blocking_and_await!(File::open(options.path))?;
-
         let json: Value = spawn_blocking_and_await!(serde_json::from_reader(&mut file))?;
 
         let coll = db.collection(&COLL_NAME);
diff --git a/benchmarks/src/bench/insert_one.rs b/benchmarks/src/bench/insert_one.rs
@@ -36,8 +36,13 @@ impl Benchmark for InsertOneBenchmark {
 
         let num_iter = options.num_iter;
 
+        // This benchmark uses a file that's quite large, and unfortunately `serde_json` has no
+        // async version of `from_reader`, so rather than read the whole file into memory at once,
+        // we use the runtime's `spawn_blocking` functionality to do this efficiently.
+        //
+        // Note that the setup is _not_ measured as part of the benchmark runtime, so even if
+        // `spawn_blocking` turned out not to be super efficient, it wouldn't be a big deal.
         let mut file = spawn_blocking_and_await!(File::open(options.path))?;
-
         let json: Value = spawn_blocking_and_await!(serde_json::from_reader(&mut file))?;
 
         let coll = db.collection(&COLL_NAME);
diff --git a/benchmarks/src/bench/json_multi_export.rs b/benchmarks/src/bench/json_multi_export.rs
@@ -1,14 +1,13 @@
-use std::{
-    fs::{File, OpenOptions},
-    io::Write,
-    path::PathBuf,
-};
+use std::path::PathBuf;
 
 use anyhow::Result;
 use futures::stream::{FuturesUnordered, StreamExt, TryStreamExt};
 use mongodb::{bson::doc, Client, Collection, Database};
 
-use crate::bench::{parse_json_file_to_documents, Benchmark, COLL_NAME, DATABASE_NAME};
+use crate::{
+    bench::{parse_json_file_to_documents, Benchmark, COLL_NAME, DATABASE_NAME},
+    fs::File,
+};
 
 const TOTAL_FILES: usize = 100;
 
@@ -42,9 +41,9 @@ impl Benchmark for JsonMultiExportBenchmark {
 
             tasks.push(async move {
                 let json_file_name = path.join(format!("ldjson{:03}.txt", i));
-                let file = spawn_blocking_and_await!(File::open(&json_file_name))?;
+                let file = File::open_read(&json_file_name).await?;
 
-                let docs = spawn_blocking_and_await!(parse_json_file_to_documents(file))?;
+                let docs = parse_json_file_to_documents(file).await?;
 
                 for mut doc in docs {
                     doc.insert("file", i as i32);
@@ -57,7 +56,6 @@ impl Benchmark for JsonMultiExportBenchmark {
         }
 
         while let Some(result) = tasks.next().await {
-            println!("done!");
             result?;
         }
 
@@ -77,32 +75,16 @@ impl Benchmark for JsonMultiExportBenchmark {
                 // lot of work for little gain since we `unwrap()` in
                 // main.rs anyway.
                 let file_name = path.join(format!("ldjson{:03}.txt", i));
-                let mut file = OpenOptions::new()
-                    .create(true)
-                    .write(true)
-                    .open(&file_name)
-                    .unwrap();
+                let mut file = File::open_write(&file_name).await.unwrap();
 
                 let mut cursor = coll_ref
                     .find(Some(doc! { "file": i as i32 }), None)
                     .await
                     .unwrap();
 
-                let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
-
-                let send_future = spawn!(async move {
-                    while let Some(doc) = cursor.try_next().await.unwrap() {
-                        sender.send(doc.to_string()).unwrap();
-                    }
-                });
-
-                let rec_future = spawn_blocking_and_await!(async move {
-                    while let Some(s) = receiver.next().await {
-                        writeln!(file, "{}", s).unwrap();
-                    }
-                });
-
-                futures::future::join(send_future, rec_future).await
+                while let Some(doc) = cursor.try_next().await.unwrap() {
+                    file.write_line(&doc.to_string()).await.unwrap();
+                }
             });
         }
 
diff --git a/benchmarks/src/bench/json_multi_import.rs b/benchmarks/src/bench/json_multi_import.rs
@@ -1,10 +1,13 @@
-use std::{fs::File, path::PathBuf};
+use std::path::PathBuf;
 
 use anyhow::Result;
 use futures::stream::{FuturesUnordered, StreamExt};
 use mongodb::{options::InsertManyOptions, Client, Collection, Database};
 
-use crate::bench::{parse_json_file_to_documents, Benchmark, COLL_NAME, DATABASE_NAME};
+use crate::{
+    bench::{parse_json_file_to_documents, Benchmark, COLL_NAME, DATABASE_NAME},
+    fs::File,
+};
 
 const TOTAL_FILES: usize = 100;
 
@@ -59,10 +62,9 @@ impl Benchmark for JsonMultiImportBenchmark {
                 let mut docs = Vec::new();
 
                 let json_file_name = path.join(format!("ldjson{:03}.txt", i));
-                let file: File = spawn_blocking_and_await!(File::open(&json_file_name)).unwrap();
+                let file = File::open_read(&json_file_name).await.unwrap();
 
-                let mut new_docs =
-                    spawn_blocking_and_await!(parse_json_file_to_documents(file)).unwrap();
+                let mut new_docs = parse_json_file_to_documents(file).await.unwrap();
 
                 docs.append(&mut new_docs);
 
diff --git a/benchmarks/src/bench/mod.rs b/benchmarks/src/bench/mod.rs
@@ -8,17 +8,18 @@ pub mod run_command;
 
 use std::{
     convert::TryInto,
-    fs::File,
-    io::{BufRead, BufReader},
     time::{Duration, Instant},
 };
 
 use anyhow::{bail, Result};
+use futures::stream::TryStreamExt;
 use indicatif::{ProgressBar, ProgressStyle};
 use lazy_static::lazy_static;
 use mongodb::bson::{Bson, Document};
 use serde_json::Value;
 
+use crate::fs::{BufReader, File};
+
 lazy_static! {
     static ref DATABASE_NAME: String = option_env!("DATABASE_NAME")
         .unwrap_or("perftest")
@@ -61,11 +62,13 @@ pub trait Benchmark: Sized {
     async fn teardown(&self) -> Result<()>;
 }
 
-pub fn parse_json_file_to_documents(file: File) -> Result<Vec<Document>> {
+pub(crate) async fn parse_json_file_to_documents(file: File) -> Result<Vec<Document>> {
     let mut docs: Vec<Document> = Vec::new();
 
-    for line in BufReader::new(file).lines() {
-        let json: Value = serde_json::from_str(&line?)?;
+    let mut lines = BufReader::new(file).lines();
+
+    while let Some(line) = lines.try_next().await? {
+        let json: Value = serde_json::from_str(&line)?;
 
         docs.push(match json.try_into()? {
             Bson::Document(doc) => doc,
diff --git a/benchmarks/src/error.rs b/benchmarks/src/error.rs
diff --git a/benchmarks/src/fs.rs b/benchmarks/src/fs.rs
@@ -0,0 +1,64 @@
+use std::path::Path;
+
+use anyhow::Result;
+#[cfg(feature = "async-std-runtime")]
+use async_std::{
+    fs::{self, OpenOptions},
+    io::{
+        self,
+        prelude::{BufReadExt, WriteExt},
+    },
+};
+use futures::stream::Stream;
+#[cfg(feature = "tokio-runtime")]
+use tokio::{
+    fs::{self, OpenOptions},
+    io::{self, AsyncBufReadExt, AsyncWriteExt},
+};
+
+pub(crate) async fn read_to_string(path: &Path) -> Result<String> {
+    let s = fs::read_to_string(path).await?;
+    Ok(s)
+}
+
+pub(crate) struct File {
+    inner: fs::File,
+}
+
+impl File {
+    pub(crate) async fn open_write(name: &Path) -> Result<Self> {
+        let inner = OpenOptions::new()
+            .create(true)
+            .write(true)
+            .open(name)
+            .await?;
+
+        Ok(Self { inner })
+    }
+
+    pub(crate) async fn open_read(name: &Path) -> Result<Self> {
+        let inner = fs::File::open(name).await?;
+
+        Ok(Self { inner })
+    }
+
+    pub(crate) async fn write_line(&mut self, s: &str) -> Result<()> {
+        Ok(self.inner.write_all(format!("{}\n", s).as_bytes()).await?)
+    }
+}
+
+pub(crate) struct BufReader {
+    inner: io::BufReader<fs::File>,
+}
+
+impl BufReader {
+    pub(crate) fn new(file: File) -> Self {
+        Self {
+            inner: io::BufReader::new(file.inner),
+        }
+    }
+
+    pub(crate) fn lines(self) -> impl Stream<Item = std::io::Result<String>> {
+        self.inner.lines()
+    }
+}
diff --git a/benchmarks/src/main.rs b/benchmarks/src/main.rs
@@ -1,18 +1,3 @@
-#[macro_use]
-macro_rules! spawn {
-    ($task:expr) => {{
-        #[cfg(feature = "tokio-runtime")]
-        {
-            tokio::task::spawn($task)
-        }
-
-        #[cfg(feature = "async-std-runtime")]
-        {
-            async_std::task::spawn($task)
-        }
-    }};
-}
-
 #[macro_use]
 macro_rules! spawn_blocking_and_await {
     ($blocking_call:expr) => {{
@@ -31,6 +16,7 @@ macro_rules! spawn_blocking_and_await {
 }
 
 mod bench;
+mod fs;
 
 use std::{
     path::{Path, PathBuf},
