flush the dump-writer only once everything has been inserted

dump/src/lib.rs

@@ -382,6 +382,7 @@ pub(crate) mod test {
         for document in &documents {
             index.push_document(document).unwrap();
         }
+        index.flush().unwrap();
         index.settings(&settings).unwrap();
 
         // ========== pushing the task queue
@@ -396,6 +397,7 @@ pub(crate) mod test {
                 }
             }
         }
+        task_queue.flush().unwrap();
 
         // ========== pushing the api keys
         let api_keys = create_test_api_keys();
@@ -404,6 +406,7 @@ pub(crate) mod test {
         for key in &api_keys {
             keys.push_key(key).unwrap();
         }
+        keys.flush().unwrap();
 
         // create the dump
         let mut file = tempfile::tempfile().unwrap();

dump/src/writer.rs

@@ -87,6 +87,11 @@ impl KeyWriter {
         self.keys.write_all(b"\n")?;
         Ok(())
     }
+
+    pub fn flush(mut self) -> Result<()> {
+        self.keys.flush()?;
+        Ok(())
+    }
 }
 
 pub struct TaskWriter {
@@ -113,12 +118,16 @@ impl TaskWriter {
     pub fn push_task(&mut self, task: &TaskDump) -> Result<UpdateFile> {
         self.queue.write_all(&serde_json::to_vec(task)?)?;
         self.queue.write_all(b"\n")?;
-        self.queue.flush()?;
 
         Ok(UpdateFile::new(
             self.update_files.join(format!("{}.jsonl", task.uid)),
         ))
     }
+
+    pub fn flush(mut self) -> Result<()> {
+        self.queue.flush()?;
+        Ok(())
+    }
 }
 
 pub struct UpdateFile {
@@ -135,7 +144,6 @@ impl UpdateFile {
         if let Some(writer) = self.writer.as_mut() {
             writer.write_all(&serde_json::to_vec(document)?)?;
             writer.write_all(b"\n")?;
-            writer.flush()?;
         } else {
             let file = File::create(&self.path).unwrap();
             self.writer = Some(BufWriter::new(file));
@@ -143,6 +151,13 @@ impl UpdateFile {
         }
         Ok(())
     }
+
+    pub fn flush(self) -> Result<()> {
+        if let Some(mut writer) = self.writer {
+            writer.flush()?;
+        }
+        Ok(())
+    }
 }
 
 pub struct IndexWriter {
@@ -167,8 +182,12 @@ impl IndexWriter {
     }
 
     pub fn push_document(&mut self, document: &Map<String, Value>) -> Result<()> {
-        self.documents.write_all(&serde_json::to_vec(document)?)?;
+        serde_json::to_writer(&mut self.documents, document)?;
         self.documents.write_all(b"\n")?;
+        Ok(())
+    }
+
+    pub fn flush(&mut self) -> Result<()> {
         self.documents.flush()?;
         Ok(())
     }