updated

Verticaleap · Verticaleap · commit 0e29dee114b4 · 2021-05-03T14:40:42.000-05:00
diff --git a/src/db.rs b/src/db.rs
@@ -1,6 +1,9 @@
+use std::sync::Arc;
+use chrono::offset::Utc;
+use tokio::sync::Mutex;
 use mongodb::bson::{doc, document::Document};
 //use mongodb::{options::ClientOptions, options::FindOptions, Client, Collection};
-use mongodb::{options::ClientOptions, options::FindOptions, Client, Cursor};
+use mongodb::{options::ClientOptions, options::FindOptions, options::InsertManyOptions, Client, Cursor};
 //use serde::{Deserialize, Serialize};
 use futures::StreamExt;
 //use clap::ArgMatches;
@@ -19,19 +22,32 @@ type BoxResult<T> = std::result::Result<T, Box<dyn error::Error + Send + Sync>>;
 impl DB {
     pub async fn init(url: &str, db: &str) -> BoxResult<Self> {
         let mut client_options = ClientOptions::parse(url).await?;
-        client_options.app_name = Some("json-bucket".to_string());
+        client_options.app_name = Some("mongodb-stream-rs".to_string());
         Ok(Self {
             client: Client::with_options(client_options)?,
             db: db.to_owned(),
             counter: Counter::new()
         })
     }
 
-    pub async fn find(&mut self, collection: &str, query: Document) -> BoxResult<(Cursor, f64)> {
+    pub async fn find(&mut self, collection: &str, query: Document, bulk_size: Option<u64>) -> BoxResult<(Cursor, f64)> {
         // Log which collection this is going into
         log::debug!("Reading {}.{}", self.db, collection);
 
+        let batch_size = match bulk_size {
+            Some(bulk_size) => {
+                if bulk_size > 64000 {
+                    log::info!("Setting mongo cursor batch_size to 64000");
+                    Some(64000u32)
+                } else {
+                    Some(bulk_size as u32)
+                }
+            },
+            None => None
+        };
+
         let find_options = FindOptions::builder()
+            .batch_size(batch_size)
             .sort(doc! { "_id": -1 })
             .build();
 
@@ -48,27 +64,93 @@ impl DB {
     pub async fn insert_cursor(&mut self, collection: &str, mut cursor: Cursor, total: f64) -> BoxResult<()> {
         let coll = self.client.database(&self.db).collection(collection);
         self.counter.total(total);
-        log::info!("Inserting {} docs", total);
+        log::info!("Inserting {} docs to {}.{}", total, self.db, collection);
+
+        // Get timestamp
+        let start = Utc::now().timestamp();
+        
         while let Some(doc) = cursor.next().await {
             match doc {
                 Ok(doc) => {
                     match coll.insert_one(doc, None).await {
                         Ok(id) => {
-                            log::info!("Inserted id: {}", id.inserted_id.to_string());
+                            log::debug!("Inserted id: {}", id.inserted_id.to_string());
                         }
                         Err(e) => {
-                            log::error!("Got error: {}", e);
+                            log::debug!("Got error: {}", e);
+                        }
+                    }
+                    self.counter.incr(&self.db, collection, 1.0, start);
+                }
+                Err(e) => {
+                    log::error!("Caught error getting next doc, skipping: {}", e);
+                    continue;
+                }
+            };
+        }
+        log::info!("Completed {}.{}", self.db, collection);
+        Ok(())
+    }
+
+    pub async fn bulk_insert_cursor(&mut self, collection: &str, mut cursor: Cursor, total: f64, bulk_count: usize) -> BoxResult<()> {
+        let coll = self.client.database(&self.db).collection(collection);
+        self.counter.total(total);
+        log::info!("Bulk inserting {} docs to {}.{} in batches of {}", total, self.db, collection, bulk_count);
+
+        let insert_many_options = InsertManyOptions::builder()
+            .ordered(Some(false))
+            .build();
+
+        // Create vector of documents to bulk upload
+//        let mut bulk = Bulk::new(bulk_count);
+        let mut bulk: Vec<Document> = Vec::with_capacity(bulk_count);
+
+        // Get timestamp
+        let start = Utc::now().timestamp();
+        
+        while let Some(doc) = cursor.next().await {
+            match doc {
+                Ok(d) => {
+                    bulk.push(d);
+                    if bulk.len() >= bulk.capacity() {
+                        match coll.insert_many(*bulk, insert_many_options.clone()).await {
+                            Ok(_) => {
+                                log::debug!("Bulk inserted {} docs", bulk_count);
+                            }
+                            Err(e) => {
+                                log::debug!("Got error with insertMany: {}", e);
+                            }
                         }
+                        bulk.clear();
+                        self.counter.incr(&self.db, collection, bulk_count as f64, start);
+                    } else {
+                        continue
                     }
-                    self.counter.incr(&self.db, collection);
+//                    if bulk.push(d).await {
+//                        log::debug!("Bulk inserting {} docs", bulk_count);
+//                        let values = bulk.get().await;
+//                        match coll.insert_many(values, insert_many_options.clone()).await {
+//                            Ok(_) => {
+//                                log::debug!("Bulk inserted {} docs", bulk_count);
+//                            }
+//                            Err(e) => {
+//                                log::debug!("Got error with insertMany: {}", e);
+//                            }
+//                        }
+//                        bulk.clear().await;
+//                        self.counter.incr(&self.db, collection, bulk_count as f64, start);
+//                    } else {
+//                        log::debug!("inserted doc: {}/{}", bulk.len().await, bulk_count);
+//                        continue
+//                    }
                 }
                 Err(e) => {
                     log::error!("Caught error getting next doc, skipping: {}", e);
                     continue;
                 }
             };
         }
-        println!("Completed {}.{}", self.db, collection);
+        log::info!("Completed {}.{}", self.db, collection);
         Ok(())
     }
 
@@ -125,21 +207,67 @@ impl Counter {
         Counter {
             count: 0.0,
             marker: 0.0,
-            total: 0.0
+            total: 0.0,
         }
     }
 
     pub fn total(&mut self, total: f64) {
         self.total = total;
     }
 
-    pub fn incr(&mut self, db: &str, collection: &str) {
-        self.count += 1f64;
+    pub fn incr(&mut self, db: &str, collection: &str, count: f64, start: i64) {
+        self.count += count;
         let percent = self.count / self.total * 100.0;
 
+        // Get time elapsed
+        let now = Utc::now().timestamp();
+        let delta = now - start;
+
+        // Get insert rate
+        let rate = self.count / delta as f64;
+
         if percent - self.marker > 1.0 {
-            println!("Copying {}.{}: {:.2}%", db, collection, percent);
+            log::info!("{}.{}: {:.2}%, {:.2}/s, {}/{}", db, collection, percent, rate, self.count, self.total);
             self.marker += 1f64;
         };
     }
 }
+
+#[derive(Debug, Clone)]
+pub struct Bulk {
+    pub inner: Arc<Mutex<Vec<Document>>>
+}
+
+impl Bulk {
+    pub fn new(size: usize) -> Bulk {
+        Bulk { 
+            inner: Arc::new(Mutex::new(Vec::with_capacity(size))) 
+        }
+    }
+
+    pub async fn push(&mut self, doc: Document) -> bool {
+        let mut me = self.inner.lock().await;
+        me.push(doc);
+
+        if me.len() >= me.capacity() {
+            return true
+        } else {
+            return false
+        }
+    }
+
+    pub async fn len(&self) -> usize {
+        let me = self.inner.lock().await;
+        me.len()
+    }
+
+    pub async fn get(&self) -> Vec<Document> {
+        let me = self.inner.lock().await;
+        me.to_vec()
+    }
+
+    pub async fn clear(&self) {
+        let mut me = self.inner.lock().await;
+        me.clear()
+    }
+}
diff --git a/src/main.rs b/src/main.rs
@@ -55,6 +55,16 @@ async fn main() -> BoxResult<()> {
                 .help("MongoDB Collection")
                 .takes_value(true),
         )
+        .arg(
+            Arg::with_name("bulk")
+                .short("b")
+                .long("bulk")
+                .required(false)
+                .value_name("STREAM_BULK")
+                .env("STREAM_BULK")
+                .help("Bulk stream documents")
+                .takes_value(true)
+        )
         .get_matches();
 
     // Initialize log Builder
@@ -69,7 +79,7 @@ async fn main() -> BoxResult<()> {
             )
         })
         .target(Target::Stdout)
-        .filter_level(LevelFilter::Error)
+        .filter_level(LevelFilter::Info)
         .parse_default_env()
         .init();
 
@@ -78,6 +88,10 @@ async fn main() -> BoxResult<()> {
     let destination= &opts.value_of("destination_uri").unwrap();
     let collection = &opts.value_of("collection").unwrap();
     let db = &opts.value_of("db").unwrap();
+    let bulk = match opts.is_present("bulk") {
+        true => Some(opts.value_of("bulk").unwrap().parse::<u32>()?),
+        false => None
+    };
 
     println!(
         "Starting mongodb-stream-rs:{}", 
@@ -88,10 +102,22 @@ async fn main() -> BoxResult<()> {
     let mut source_db = DB::init(&source, &db).await?;
     let mut destination_db = DB::init(&destination, &db).await?;
 
-    // Acquire cursor from source
-    let (source_cursor,total) = source_db.find(collection, doc!{}).await?;
+    // If bulk flag is set, use insertMany
+    match bulk {
+        Some(bulk_size) => {
+            // Acquire cursor from source
+            let (source_cursor,total) = source_db.find(collection, doc!{}, Some(bulk_size as u64)).await?;
+
+            destination_db.bulk_insert_cursor(collection, source_cursor, total, bulk_size as usize).await?;
+        }
+        None => {
+            // Acquire cursor from source
+            let (source_cursor,total) = source_db.find(collection, doc!{}, None).await?;
+
+            destination_db.insert_cursor(collection, source_cursor, total).await?
+        }
 
-    destination_db.insert_cursor(collection, source_cursor, total).await?;
+    };
 
     Ok(())
 }
