RUST-1902 Fix benchmarks for client 3.0 API (#1172)

abr-egn · web-flow · commit 6a4b739a9264 · 2024-07-22T15:44:57.000-04:00
diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
@@ -4,31 +4,18 @@ version = "0.1.0"
 authors = ["benjirewis <ben.rewis98@gmail.com>"]
 edition = "2018"
 
-[features]
-default = ["tokio-runtime"]
-tokio-runtime = [
-  "tokio/fs",
-  "tokio/macros",
-  "tokio/rt",
-  "tokio/rt-multi-thread",
-  "tokio-stream",
-  "mongodb/tokio-runtime"
-]
-async-std-runtime = ["async-std", "mongodb/async-std-runtime"]
-
 [dependencies]
-mongodb = { path = "..", default-features = false }
+mongodb = { path = ".." }
 serde_json = "1.0.59"
 once_cell = "1.19.0"
 clap = "2.33.3"
 indicatif = "0.15.0"
 async-trait = "0.1.41"
-tokio = { version = "1.6", features = ["sync"] }
+tokio = { version = "1.6", features = ["sync", "fs"] }
 tokio-util = "0.7"
-tokio-stream = { version = "0.1.6", features = ["io-util"], optional = true }
-# "unstable" feature is needed for `spawn_blocking`, which is only used in task setup
-async-std = { version = "1.9.0", optional = true, features = ["attributes", "unstable"] }
+tokio-stream = { version = "0.1.6", features = ["io-util"] }
 futures = "0.3.8"
 anyhow = "1.0.34"
 serde = "1"
 num_enum = "0.5"
+futures-util = "0.3.30"
diff --git a/benchmarks/src/bench.rs b/benchmarks/src/bench.rs
@@ -149,7 +149,7 @@ pub async fn drop_database(uri: &str, database: &str) -> Result<()> {
 
     let hello = client
         .database("admin")
-        .run_command(doc! { "hello": true }, None)
+        .run_command(doc! { "hello": true })
         .await?;
 
     client.database(&database).drop().await?;
@@ -162,10 +162,10 @@ pub async fn drop_database(uri: &str, database: &str) -> Result<()> {
         for host in options.hosts {
             client
                 .database("admin")
-                .run_command(
-                    doc! { "flushRouterConfig": 1 },
-                    SelectionCriteria::Predicate(Arc::new(move |s| s.address() == &host)),
-                )
+                .run_command(doc! { "flushRouterConfig": 1 })
+                .selection_criteria(SelectionCriteria::Predicate(Arc::new(move |s| {
+                    s.address() == &host
+                })))
                 .await?;
         }
     }
diff --git a/benchmarks/src/bench/find_many.rs b/benchmarks/src/bench/find_many.rs
@@ -3,7 +3,7 @@ use std::{convert::TryInto, path::PathBuf};
 use anyhow::{bail, Result};
 use futures::stream::StreamExt;
 use mongodb::{
-    bson::{Bson, Document, RawDocumentBuf},
+    bson::{doc, Bson, Document, RawDocumentBuf},
     Client,
     Collection,
     Database,
@@ -59,7 +59,7 @@ impl Benchmark for FindManyBenchmark {
 
         let coll = db.collection(&COLL_NAME);
         let docs = vec![doc.clone(); num_iter];
-        coll.insert_many(docs, None).await?;
+        coll.insert_many(docs).await?;
 
         Ok(FindManyBenchmark {
             db,
@@ -74,7 +74,7 @@ impl Benchmark for FindManyBenchmark {
             bench: &FindManyBenchmark,
         ) -> Result<()> {
             let coll = bench.coll.clone_with_type::<T>();
-            let mut cursor = coll.find(None, None).await?;
+            let mut cursor = coll.find(doc! {}).await?;
             while let Some(doc) = cursor.next().await {
                 doc?;
             }
diff --git a/benchmarks/src/bench/find_one.rs b/benchmarks/src/bench/find_one.rs
@@ -50,7 +50,7 @@ impl Benchmark for FindOneBenchmark {
         let coll = db.collection(&COLL_NAME);
         for i in 0..num_iter {
             doc.insert("_id", i as i32);
-            coll.insert_one(doc.clone(), None).await?;
+            coll.insert_one(doc.clone()).await?;
         }
 
         Ok(FindOneBenchmark {
@@ -63,9 +63,7 @@ impl Benchmark for FindOneBenchmark {
 
     async fn do_task(&self) -> Result<()> {
         for i in 0..self.num_iter {
-            self.coll
-                .find_one(Some(doc! { "_id": i as i32 }), None)
-                .await?;
+            self.coll.find_one(doc! { "_id": i as i32 }).await?;
         }
 
         Ok(())
diff --git a/benchmarks/src/bench/gridfs_download.rs b/benchmarks/src/bench/gridfs_download.rs
@@ -1,6 +1,7 @@
 use std::path::PathBuf;
 
 use anyhow::{Context, Result};
+use futures::{AsyncReadExt, AsyncWriteExt};
 use mongodb::{bson::Bson, gridfs::GridFsBucket, Client};
 
 use crate::{
@@ -33,24 +34,25 @@ impl Benchmark for GridFsDownloadBenchmark {
         let bucket = db.gridfs_bucket(None);
 
         let file = open_async_read_compat(&options.path).await?;
-        let file_id = bucket
-            .upload_from_futures_0_3_reader("gridfstest", file, None)
-            .await
-            .context("upload file")?;
+        let mut upload = bucket.open_upload_stream("gridfstest").await?;
+        let file_id = upload.id().clone();
+        futures_util::io::copy(file, &mut upload).await?;
+        upload.close().await?;
 
         Ok(Self {
             uri: options.uri,
             bucket,
-            file_id: file_id.into(),
+            file_id,
         })
     }
 
     async fn do_task(&self) -> Result<()> {
         let mut buf = vec![];
-        self.bucket
-            .download_to_futures_0_3_writer(self.file_id.clone(), &mut buf)
-            .await
-            .context("download file")?;
+        let mut download = self
+            .bucket
+            .open_download_stream(self.file_id.clone())
+            .await?;
+        download.read_to_end(&mut buf).await?;
 
         Ok(())
     }
diff --git a/benchmarks/src/bench/gridfs_multi_download.rs b/benchmarks/src/bench/gridfs_multi_download.rs
@@ -5,7 +5,7 @@ use std::{
 
 use anyhow::{Context, Result};
 use futures::AsyncWriteExt;
-use mongodb::{bson::oid::ObjectId, gridfs::GridFsBucket, Client};
+use mongodb::{bson::Bson, gridfs::GridFsBucket, Client};
 use once_cell::sync::Lazy;
 
 use crate::{
@@ -19,7 +19,7 @@ static DOWNLOAD_PATH: Lazy<PathBuf> =
 pub struct GridFsMultiDownloadBenchmark {
     uri: String,
     bucket: GridFsBucket,
-    ids: Vec<ObjectId>,
+    ids: Vec<Bson>,
 }
 
 pub struct Options {
@@ -46,10 +46,12 @@ impl Benchmark for GridFsMultiDownloadBenchmark {
             let path = entry?.path();
 
             let file = open_async_read_compat(&path).await?;
-            let id = bucket
-                .upload_from_futures_0_3_reader(path.display().to_string(), file, None)
-                .await
-                .context("upload file")?;
+            let mut upload = bucket
+                .open_upload_stream(path.display().to_string())
+                .await?;
+            let id = upload.id().clone();
+            futures_util::io::copy(file, &mut upload).await?;
+            upload.close().await?;
             ids.push(id);
         }
 
@@ -64,7 +66,7 @@ impl Benchmark for GridFsMultiDownloadBenchmark {
 
     async fn before_task(&mut self) -> Result<()> {
         for id in &self.ids {
-            let path = get_filename(id.clone());
+            let path = get_filename(&id);
             if Path::try_exists(&path)? {
                 remove_file(path)?;
             }
@@ -81,15 +83,12 @@ impl Benchmark for GridFsMultiDownloadBenchmark {
             let id = id.clone();
 
             tasks.push(crate::spawn(async move {
-                let download_path = get_filename(id);
+                let download_path = get_filename(&id);
                 let mut file = open_async_write_compat(&download_path)
                     .await
                     .context("open file")?;
-
-                bucket
-                    .download_to_futures_0_3_writer(id.into(), &mut file)
-                    .await
-                    .context("download file")?;
+                let download = bucket.open_download_stream(id).await?;
+                futures_util::io::copy(download, &mut file).await?;
 
                 file.flush().await?;
 
@@ -114,6 +113,6 @@ impl Benchmark for GridFsMultiDownloadBenchmark {
     }
 }
 
-fn get_filename(id: ObjectId) -> PathBuf {
+fn get_filename(id: &Bson) -> PathBuf {
     DOWNLOAD_PATH.join(format!("file{}.txt", id))
 }
diff --git a/benchmarks/src/bench/gridfs_multi_upload.rs b/benchmarks/src/bench/gridfs_multi_upload.rs
@@ -1,6 +1,7 @@
 use std::{fs::read_dir, path::PathBuf};
 
 use anyhow::{Context, Result};
+use futures::AsyncWriteExt;
 use mongodb::{gridfs::GridFsBucket, Client};
 
 use crate::{
@@ -39,11 +40,9 @@ impl Benchmark for GridFsMultiUploadBenchmark {
 
     async fn before_task(&mut self) -> Result<()> {
         self.bucket.drop().await.context("bucket drop")?;
-
-        self.bucket
-            .upload_from_futures_0_3_reader("beforetask", &[11u8][..], None)
-            .await
-            .context("single byte upload")?;
+        let mut upload = self.bucket.open_upload_stream("beforetask").await?;
+        upload.write_all(&[11u8][..]).await?;
+        upload.close().await?;
 
         Ok(())
     }
@@ -57,10 +56,11 @@ impl Benchmark for GridFsMultiUploadBenchmark {
             tasks.push(crate::spawn(async move {
                 let path = entry?.path();
                 let file = open_async_read_compat(&path).await?;
-                bucket
-                    .upload_from_futures_0_3_reader(path.display().to_string(), file, None)
-                    .await
-                    .context("upload file")?;
+                let mut upload = bucket
+                    .open_upload_stream(path.display().to_string())
+                    .await?;
+                futures_util::io::copy(file, &mut upload).await?;
+                upload.close().await?;
 
                 let ok: anyhow::Result<()> = Ok(());
                 ok
diff --git a/benchmarks/src/bench/gridfs_upload.rs b/benchmarks/src/bench/gridfs_upload.rs
@@ -1,7 +1,7 @@
 use std::path::PathBuf;
 
 use anyhow::{Context, Result};
-use futures::AsyncReadExt;
+use futures::{AsyncReadExt, AsyncWriteExt};
 use mongodb::{gridfs::GridFsBucket, Client};
 
 use crate::{
@@ -44,20 +44,17 @@ impl Benchmark for GridFsUploadBenchmark {
 
     async fn before_task(&mut self) -> Result<()> {
         self.bucket.drop().await.context("bucket drop")?;
-
-        self.bucket
-            .upload_from_futures_0_3_reader("beforetask", &[11u8][..], None)
-            .await
-            .context("single byte upload")?;
+        let mut upload = self.bucket.open_upload_stream("beforetask").await?;
+        upload.write_all(&[11u8][..]).await?;
+        upload.close().await?;
 
         Ok(())
     }
 
     async fn do_task(&self) -> Result<()> {
-        self.bucket
-            .upload_from_futures_0_3_reader("gridfstest", &self.bytes[..], None)
-            .await
-            .context("upload bytes")?;
+        let mut upload = self.bucket.open_upload_stream("gridfstest").await?;
+        upload.write_all(&self.bytes[..]).await?;
+        upload.close().await?;
 
         Ok(())
     }
diff --git a/benchmarks/src/bench/insert_many.rs b/benchmarks/src/bench/insert_many.rs
@@ -64,7 +64,7 @@ impl Benchmark for InsertManyBenchmark {
     async fn before_task(&mut self) -> Result<()> {
         self.coll.drop().await?;
         self.db
-            .create_collection(COLL_NAME.as_str(), None)
+            .create_collection(COLL_NAME.as_str())
             .await
             .context("create in before")?;
 
@@ -74,7 +74,7 @@ impl Benchmark for InsertManyBenchmark {
     async fn do_task(&self) -> Result<()> {
         let insertions = vec![&self.doc; self.num_copies];
         self.coll
-            .insert_many(insertions, None)
+            .insert_many(insertions)
             .await
             .context("insert many")?;
 
diff --git a/benchmarks/src/bench/insert_one.rs b/benchmarks/src/bench/insert_one.rs
@@ -64,7 +64,7 @@ impl Benchmark for InsertOneBenchmark {
     async fn before_task(&mut self) -> Result<()> {
         self.coll.drop().await?;
         self.db
-            .create_collection(COLL_NAME.as_str(), None)
+            .create_collection(COLL_NAME.as_str())
             .await
             .context("create collection")?;
 
@@ -74,7 +74,7 @@ impl Benchmark for InsertOneBenchmark {
     async fn do_task(&self) -> Result<()> {
         for _ in 0..self.num_iter {
             self.coll
-                .insert_one(&self.doc, None)
+                .insert_one(&self.doc)
                 .await
                 .context("insert one")?;
         }
diff --git a/benchmarks/src/bench/json_multi_export.rs b/benchmarks/src/bench/json_multi_export.rs
@@ -55,7 +55,7 @@ impl Benchmark for JsonMultiExportBenchmark {
 
                 for mut doc in docs {
                     doc.insert("file", i as i32);
-                    coll.insert_one(doc, None).await?;
+                    coll.insert_one(doc).await?;
                 }
 
                 let ok: anyhow::Result<()> = Ok(());
@@ -89,10 +89,7 @@ impl Benchmark for JsonMultiExportBenchmark {
                 let file_name = path.join(format!("ldjson{:03}.txt", i));
                 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 mut cursor = coll_ref.find(doc! { "file": i as i32 }).await.unwrap();
 
                 while let Some(doc) = cursor.try_next().await.unwrap() {
                     file.write_line(serde_json::to_string(&doc).unwrap().as_str())
diff --git a/benchmarks/src/bench/json_multi_import.rs b/benchmarks/src/bench/json_multi_import.rs
@@ -2,7 +2,7 @@ use std::path::PathBuf;
 
 use anyhow::Result;
 use futures::stream::TryStreamExt;
-use mongodb::{options::InsertManyOptions, Client, Collection, Database};
+use mongodb::{Client, Collection, Database};
 
 use crate::{
     bench::{Benchmark, COLL_NAME, DATABASE_NAME},
@@ -48,7 +48,7 @@ impl Benchmark for JsonMultiImportBenchmark {
 
     async fn before_task(&mut self) -> Result<()> {
         self.coll.drop().await?;
-        self.db.create_collection(COLL_NAME.as_str(), None).await?;
+        self.db.create_collection(COLL_NAME.as_str()).await?;
 
         Ok(())
     }
@@ -74,8 +74,7 @@ impl Benchmark for JsonMultiImportBenchmark {
                     docs.push(serde_json::from_str(&line).unwrap());
                 }
 
-                let opts = Some(InsertManyOptions::builder().ordered(false).build());
-                coll_ref.insert_many(docs, opts).await.unwrap();
+                coll_ref.insert_many(docs).ordered(false).await.unwrap();
             }));
         }
 
diff --git a/benchmarks/src/bench/run_command.rs b/benchmarks/src/bench/run_command.rs
@@ -40,7 +40,7 @@ impl Benchmark for RunCommandBenchmark {
         for _ in 0..self.num_iter {
             let _doc = self
                 .db
-                .run_command(self.cmd.clone(), None)
+                .run_command(self.cmd.clone())
                 .await
                 .context("run command")?;
         }
diff --git a/benchmarks/src/fs.rs b/benchmarks/src/fs.rs
@@ -64,7 +64,7 @@ impl BufReader {
         }
     }
 
-    pub(crate) fn lines(self) -> impl Stream<Item = std::io::Result<String>> {
+    pub(crate) fn lines(self) -> impl tokio_stream::Stream<Item = std::io::Result<String>> {
         tokio_stream::wrappers::LinesStream::new(self.inner.lines())
     }
 }
