RUST-1473 Add GridFS Benchmarks (#772)

Author: isabelatkinson

.evergreen/benchmarks.yml

@@ -275,62 +275,32 @@ post:
   - func: "cleanup"
 
 tasks:
-  - name: "benchmark-4.2-standalone"
-    tags: ["4.2", "standalone"]
+  - name: "benchmark-rapid-standalone"
+    tags: ["rapid", "standalone"]
     commands:
       - func: "bootstrap mongo-orchestration"
         vars:
-          MONGODB_VERSION: "4.2"
+          MONGODB_VERSION: "rapid"
           TOPOLOGY: "server"
       - func: "run driver benchmarks"
       - func: "upload benchmark results"
 
-  - name: "benchmark-4.2-replica_set"
-    tags: ["4.2", "replica_set"]
+  - name: "benchmark-rapid-replica_set"
+    tags: ["rapid", "replica_set"]
     commands:
       - func: "bootstrap mongo-orchestration"
         vars:
-          MONGODB_VERSION: "4.2"
+          MONGODB_VERSION: "rapid"
           TOPOLOGY: "replica_set"
       - func: "run driver benchmarks"
       - func: "upload benchmark results"
 
-  - name: "benchmark-4.2-sharded_cluster"
-    tags: ["4.2", "sharded_cluster"]
+  - name: "benchmark-rapid-sharded_cluster"
+    tags: ["rapid", "sharded_cluster"]
     commands:
       - func: "bootstrap mongo-orchestration"
         vars:
-          MONGODB_VERSION: "4.2"
-          TOPOLOGY: "sharded_cluster"
-      - func: "run driver benchmarks"
-      - func: "upload benchmark results"
-
-  - name: "benchmark-5.0-standalone"
-    tags: ["5.0", "standalone"]
-    commands:
-      - func: "bootstrap mongo-orchestration"
-        vars:
-          MONGODB_VERSION: "5.0"
-          TOPOLOGY: "server"
-      - func: "run driver benchmarks"
-      - func: "upload benchmark results"
-
-  - name: "benchmark-5.0-replica_set"
-    tags: ["5.0", "replica_set"]
-    commands:
-      - func: "bootstrap mongo-orchestration"
-        vars:
-          MONGODB_VERSION: "5.0"
-          TOPOLOGY: "replica_set"
-      - func: "run driver benchmarks"
-      - func: "upload benchmark results"
-
-  - name: "benchmark-5.0-sharded_cluster"
-    tags: ["5.0", "sharded_cluster"]
-    commands:
-      - func: "bootstrap mongo-orchestration"
-        vars:
-          MONGODB_VERSION: "5.0"
+          MONGODB_VERSION: "rapid"
           TOPOLOGY: "sharded_cluster"
       - func: "run driver benchmarks"
       - func: "upload benchmark results"
@@ -350,14 +320,10 @@ axes:
   - id: "mongodb-version"
     display_name: MongoDB Version
     values:
-      - id: "5.0"
-        display_name: "5.0"
-        variables:
-          MONGODB_VERSION: "5.0"
-      - id: "4.2"
-        display_name: "4.2"
+      - id: "rapid"
+        display_name: "rapid"
         variables:
-          MONGODB_VERSION: "4.2"
+          MONGODB_VERSION: "rapid"
 
   - id: "topology"
     display_name: Topology
@@ -434,8 +400,7 @@ buildvariants:
     async-runtime: "*"
   display_name: "${os} ${auth-and-tls} with ${async-runtime}"
   tasks:
-     - ".5.0"
-     - ".4.2"
+     - ".rapid"
 -
   matrix_name: "bson benchmarks"
   matrix_spec:

benchmarks/Cargo.toml

@@ -24,6 +24,7 @@ clap = "2.33.3"
 indicatif = "0.15.0"
 async-trait = "0.1.41"
 tokio = { version = "1.6", features = ["sync"] }
+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"] }

benchmarks/README.md

@@ -20,27 +20,31 @@ When running the benchmarks with the `--ids` flag, you can refer to each benchma
 benchmarks by separating each benchmark's id with a comma. For example `cargo run --release -- --ids 1,2,3,4` would run all 
 the single-doc benchmarks. By default, all benchmarks are executed. The table below lists each benchmark's id.
 
-| Benchmark                      | ID |
-|--------------------------------|----|
-| Run command                    | 1  |
-| Find one by ID                 | 2  |
-| Small doc insertOne            | 3  |
-| Large doc insertOne            | 4  |
-| Find many and empty the cursor | 5  |
-| Small doc bulk insert          | 6  |
-| Large doc bulk insert          | 7  |
-| LDJSON multi-file import       | 8  |
-| LDJSON multi-file export       | 9  |
-| BSON flat decode               | 10 |
-| BSON flat encode               | 11 |
-| BSON deeply nested decode      | 12 |
-| BSON deeply nested encode      | 13 |
-| BSON full document decode      | 14 |
-| BSON full document encode      | 15 |  
-| All benchmarks                 | all|
+| Benchmark                                      | ID |
+|----------------------------------------------- |----|
+| Run command                                    | 1  |
+| Find one by ID                                 | 2  |
+| Small doc insertOne                            | 3  |
+| Large doc insertOne                            | 4  |
+| Find many and empty the cursor                 | 5  |
+| Small doc bulk insert                          | 6  |
+| Large doc bulk insert                          | 7  |
+| LDJSON multi-file import                       | 8  |
+| LDJSON multi-file export                       | 9  |
+| BSON flat decode                               | 10 |
+| BSON flat encode                               | 11 |
+| BSON deeply nested decode                      | 12 |
+| BSON deeply nested encode                      | 13 |
+| BSON full document decode                      | 14 |
+| BSON full document encode                      | 15 |
+| Find many and empty the cursor (raw BSON)      | 16 |
+| Find many and empty the cursor (serde structs) | 17 |
+| GridFS download                                | 18 |
+| GridFS upload                                  | 19 |
+| GridFS multi-file download                     | 20 |
+| GridFS multi-file upload                       | 21 |
+| All benchmarks                                 | all|
 
 Note that in order to compare against the other drivers, an inMemory mongod instance should be used.
 
-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.

benchmarks/src/bench/gridfs_download.rs

@@ -0,0 +1,65 @@
+use std::path::PathBuf;
+
+use anyhow::{Context, Result};
+use mongodb::{bson::Bson, gridfs::GridFsBucket, Client};
+
+use crate::{
+    bench::{drop_database, Benchmark, DATABASE_NAME},
+    fs::open_async_read_compat,
+};
+
+pub struct GridFsDownloadBenchmark {
+    uri: String,
+    bucket: GridFsBucket,
+    file_id: Bson,
+}
+
+pub struct Options {
+    pub uri: String,
+    pub path: PathBuf,
+}
+
+#[async_trait::async_trait]
+impl Benchmark for GridFsDownloadBenchmark {
+    type Options = Options;
+
+    async fn setup(options: Self::Options) -> Result<Self> {
+        let client = Client::with_uri_str(&options.uri).await?;
+        let db = client.database(&DATABASE_NAME);
+        drop_database(&options.uri, &DATABASE_NAME)
+            .await
+            .context("drop database setup")?;
+
+        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")?;
+
+        Ok(Self {
+            uri: options.uri,
+            bucket,
+            file_id: file_id.into(),
+        })
+    }
+
+    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")?;
+
+        Ok(())
+    }
+
+    async fn teardown(&self) -> Result<()> {
+        drop_database(&self.uri, &DATABASE_NAME)
+            .await
+            .context("drop database teardown")?;
+
+        Ok(())
+    }
+}

benchmarks/src/bench/gridfs_multi_download.rs

@@ -0,0 +1,121 @@
+use std::{
+    fs::{create_dir, read_dir, remove_file},
+    path::{Path, PathBuf},
+};
+
+use anyhow::{Context, Result};
+use futures::AsyncWriteExt;
+use lazy_static::lazy_static;
+use mongodb::{bson::oid::ObjectId, gridfs::GridFsBucket, Client};
+
+use crate::{
+    bench::{drop_database, Benchmark, DATABASE_NAME},
+    fs::{open_async_read_compat, open_async_write_compat},
+};
+
+lazy_static! {
+    static ref DOWNLOAD_PATH: PathBuf =
+        Path::new(env!("CARGO_MANIFEST_DIR")).join("gridfs_multi_download");
+}
+
+pub struct GridFsMultiDownloadBenchmark {
+    uri: String,
+    bucket: GridFsBucket,
+    ids: Vec<ObjectId>,
+}
+
+pub struct Options {
+    pub uri: String,
+    pub path: PathBuf,
+}
+
+#[async_trait::async_trait]
+impl Benchmark for GridFsMultiDownloadBenchmark {
+    type Options = Options;
+
+    async fn setup(options: Self::Options) -> Result<Self> {
+        let client = Client::with_uri_str(&options.uri).await?;
+        let db = client.database(&DATABASE_NAME);
+        drop_database(&options.uri, &DATABASE_NAME)
+            .await
+            .context("drop database setup")?;
+
+        let bucket = db.gridfs_bucket(None);
+        bucket.drop().await.context("drop bucket setup")?;
+
+        let mut ids = vec![];
+        for entry in read_dir(options.path)? {
+            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")?;
+            ids.push(id);
+        }
+
+        create_dir(DOWNLOAD_PATH.as_path())?;
+
+        Ok(Self {
+            uri: options.uri,
+            bucket,
+            ids,
+        })
+    }
+
+    async fn before_task(&mut self) -> Result<()> {
+        for id in &self.ids {
+            let path = get_filename(id.clone());
+            if Path::try_exists(&path)? {
+                remove_file(path)?;
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn do_task(&self) -> Result<()> {
+        let mut tasks = vec![];
+
+        for id in &self.ids {
+            let bucket = self.bucket.clone();
+            let id = id.clone();
+
+            tasks.push(crate::spawn(async move {
+                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")?;
+
+                file.flush().await?;
+
+                let ok: anyhow::Result<()> = Ok(());
+                ok
+            }));
+        }
+
+        for task in tasks {
+            task.await?;
+        }
+
+        Ok(())
+    }
+
+    async fn teardown(&self) -> Result<()> {
+        drop_database(&self.uri, &DATABASE_NAME)
+            .await
+            .context("drop database teardown")?;
+
+        Ok(())
+    }
+}
+
+fn get_filename(id: ObjectId) -> PathBuf {
+    DOWNLOAD_PATH.join(format!("file{}.txt", id))
+}