diff --git a/Cargo.lock b/Cargo.lock index 4ac80bd..41ad4c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -301,6 +301,7 @@ dependencies = [ "indicatif", "itertools", "memdb", + "mimalloc", "native_db", "native_model", "serde", diff --git a/crates/admin/Cargo.toml b/crates/admin/Cargo.toml index eddbfb2..0d8ba94 100644 --- a/crates/admin/Cargo.toml +++ b/crates/admin/Cargo.toml @@ -17,4 +17,5 @@ serde_json.workspace = true itertools.workspace = true snafu.workspace = true -memdb.workspace = true \ No newline at end of file +memdb.workspace = true +mimalloc = "0.1.50" \ No newline at end of file diff --git a/crates/admin/src/main.rs b/crates/admin/src/main.rs index 00364e2..75344fd 100644 --- a/crates/admin/src/main.rs +++ b/crates/admin/src/main.rs @@ -18,6 +18,7 @@ use console::style; use dialoguer::{theme::ColorfulTheme, Select}; +use mimalloc::MiMalloc; use crate::{migrate::handle_migration, reset::handle_reset_password}; @@ -25,8 +26,15 @@ pub mod meta; pub mod migrate; pub mod reset; +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + run_interactive(); +} + #[tokio::main] -async fn main() { +async fn run_interactive() { let theme = ColorfulTheme::default(); println!( "\n{}\n", diff --git a/crates/admin/src/migrate.rs b/crates/admin/src/migrate.rs index 4d870e3..6ecdb9a 100644 --- a/crates/admin/src/migrate.rs +++ b/crates/admin/src/migrate.rs @@ -1,7 +1,7 @@ use std::path::{Path, PathBuf}; use bichon_core::migrate::{ - do_migrate, is_tantivy_index_dir, + count_eml_segments, do_migrate_segment, is_tantivy_index_dir, store::{LegacyDirs, NewDirs}, }; use console::style; @@ -234,8 +234,7 @@ pub fn handle_migration(theme: &ColorfulTheme) { ); eprintln!( "{}", - style("Aborting migration. No changes have been made to Tantivy data.") - .yellow() + style("Aborting migration. No changes have been made to Tantivy data.").yellow() ); return; } @@ -246,48 +245,179 @@ pub fn handle_migration(theme: &ColorfulTheme) { style("⌛").yellow(), style("Step 2: Migrating email index and blob data...").cyan() ); - let pb = ProgressBar::new(0); - pb.set_style(ProgressStyle::default_bar() - .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta}) {msg}") - .unwrap() - .progress_chars("#>-")); - let legacy = LegacyDirs::new(index_path, data_path); - let new_dirs = NewDirs::new(new_index_path, new_data_path); - if let Err(e) = do_migrate(legacy, new_dirs, |msg| { - if let Some(data) = msg.strip_prefix("PROGRESS:") { - let parts: Vec<&str> = data.split(':').collect(); - if parts.len() == 2 { - let migrated = parts[0].parse::().unwrap_or(0); - let skipped = parts[1].parse::().unwrap_or(0); - pb.set_position(migrated + skipped); - pb.set_message(format!( - "Migrated: {}, {} {}", - style(migrated).green(), - style(skipped).red(), - style("skipped").dim() - )); - } - } else if let Some(total) = msg.strip_prefix("TOTAL:") { - pb.set_length(total.parse().unwrap_or(0)); - } else if msg.starts_with("WARN:") { - pb.println(format!("{} {}", style("⚠").yellow(), &msg[5..])); - } else if let Some(done_data) = msg.strip_prefix("DONE:") { - let parts: Vec<&str> = done_data.split(':').collect(); - pb.finish_with_message(format!( - "Migration finished. Total: {}, Skipped: {}", - parts.get(0).unwrap_or(&"0"), - parts.get(1).unwrap_or(&"0") - )); + println!( + "\n{} {}", + style("ℹ").blue(), + style("Batch size controls memory usage during migration:").dim() + ); + println!( + " {} 1000 — ~500MB RAM (slower, low memory)", + style("•").dim() + ); + println!(" {} 3000 — ~1GB RAM (recommended)", style("•").dim()); + println!( + " {} 5000 — ~2GB RAM (faster, high memory)", + style("•").dim() + ); + println!( + " {} Note: actual memory usage depends on your average email size.", + style("•").yellow() + ); + println!( + " {} If your mailbox contains many large attachments, use a smaller batch size.\n", + style(" ").dim() + ); + + let batch_size: u32 = { + let input: String = Input::with_theme(&ColorfulTheme::default()) + .with_prompt("Enter batch size (affects memory usage, see notes above)") + .default("3000".to_string()) + .validate_with(|s: &String| match s.trim().parse::() { + Ok(n) if n > 0 => Ok(()), + _ => Err("Please enter a valid positive number"), + }) + .interact_text() + .unwrap_or("3000".to_string()); + input.trim().parse::().unwrap_or(3000) + }; + + println!( + "{} Using batch size: {}\n", + style("✓").green(), + style(batch_size).cyan().bold() + ); + + println!( + "{} Using batch size: {}\n", + style("✓").green(), + style(batch_size).cyan().bold() + ); + + let legacy = LegacyDirs::new(index_path.clone(), data_path.clone()); + let total_segments = match count_eml_segments(&legacy) { + Ok(n) => n, + Err(e) => { + eprintln!( + "\n{} Failed to count EML segments:\n{:?}", + style("✘").red().bold(), + e + ); + return; } - }) { - eprintln!( - "\n{} Migration failed:\n{:?}", - style("✘").red().bold(), - style(e).red() + }; + + if total_segments == 0 { + println!( + "{} {}", + style("✔").green(), + style("No EML segments found. Nothing to migrate.").bold() ); return; } + + println!( + "{} EML segments to migrate: {}", + style("⌛").yellow(), + style(total_segments).cyan() + ); + + let pb = ProgressBar::new(total_segments as u64); + pb.set_style( + ProgressStyle::default_bar() + .template( + "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta}) {msg}", + ) + .unwrap() + .progress_chars("#>-"), + ); + + let mut grand_total_migrated: usize = 0; + let mut grand_total_skipped: usize = 0; + + for seg_idx in 0..total_segments { + let seg_total: std::cell::Cell = std::cell::Cell::new(0); + + pb.set_message(format!("Segment {}/{}", seg_idx + 1, total_segments)); + let legacy = LegacyDirs::new(index_path.clone(), data_path.clone()); + match do_migrate_segment( + batch_size, + legacy, + NewDirs::new(new_index_path.clone(), new_data_path.clone()), + seg_idx, + |msg| { + if let Some(data) = msg.strip_prefix("TOTAL:") { + seg_total.set(data.parse().unwrap_or(0)); + } else if let Some(data) = msg.strip_prefix("PHASE1:") { + let parts: Vec<&str> = data.split('/').collect(); + let scanned: usize = parts.get(0).and_then(|s| s.parse().ok()).unwrap_or(0); + let total: usize = parts + .get(1) + .and_then(|s| s.split_once(" skipped:").map(|(n, _)| n)) + .and_then(|s| s.parse().ok()) + .unwrap_or(0); + let skipped: usize = data + .split_once("skipped:") + .and_then(|(_, s)| s.parse().ok()) + .unwrap_or(0); + let pct = if total > 0 { + (scanned * 100) / total + } else { + 0 + }; + pb.set_message(format!( + "Segment {}/{} [scanning {}/{} skipped:{} {}%]", + seg_idx + 1, + total_segments, + scanned, + total, + skipped, + pct, + )); + } else if let Some(data) = msg.strip_prefix("PROGRESS:") { + let parts: Vec<&str> = data.split(':').collect(); + let migrated: usize = parts.get(0).and_then(|s| s.parse().ok()).unwrap_or(0); + let total = seg_total.get(); + let pct = if total > 0 { + (migrated * 100) / total + } else { + 0 + }; + pb.set_message(format!( + "Segment {}/{} [migrating {}/{} {}%]", + seg_idx + 1, + total_segments, + migrated, + total, + pct, + )); + } else if let Some(warn) = msg.strip_prefix("WARN:") { + pb.println(format!("{} {}", style("⚠").yellow(), warn)); + } else if let Some(done_data) = msg.strip_prefix("DONE:") { + let parts: Vec<&str> = done_data.split(':').collect(); + let migrated: usize = parts.get(0).and_then(|s| s.parse().ok()).unwrap_or(0); + let skipped: usize = parts.get(1).and_then(|s| s.parse().ok()).unwrap_or(0); + grand_total_migrated += migrated; + grand_total_skipped += skipped; + } + }, + ) { + Ok(()) => {} + Err(e) => { + pb.finish_with_message(format!("{}", style("Migration failed.").red())); + eprintln!("\n{} {:?}", style("✘").red().bold(), e); + return; + } + } + + pb.set_position((seg_idx + 1) as u64); + } + + pb.finish_with_message(format!( + "Migration finished. Total: {}, Skipped: {}", + grand_total_migrated, grand_total_skipped + )); + println!( "{} {}", style("✔").green(), diff --git a/crates/core/src/migrate/mod.rs b/crates/core/src/migrate/mod.rs index 444cb3b..02f15aa 100644 --- a/crates/core/src/migrate/mod.rs +++ b/crates/core/src/migrate/mod.rs @@ -1,4 +1,4 @@ -use std::path::PathBuf; +use std::{collections::HashMap, path::PathBuf}; use crate::{ error::{code::ErrorCode, BichonResult}, @@ -10,7 +10,11 @@ use crate::{ settings::cli::SETTINGS, }; use tantivy::{ - collector::TopDocs, query::AllQuery, schema::Value, DocAddress, Index, TantivyDocument, + collector::TopDocs, + columnar::Column, + query::TermQuery, + schema::{IndexRecordOption, Value}, + DocAddress, Index, TantivyDocument, Term, }; pub mod legacy; @@ -43,6 +47,18 @@ pub fn is_tantivy_index_dir(dir: &PathBuf) -> std::io::Result { Ok(has_meta_json && match_count >= 3) } +/// Return the number of segments in the legacy EML Tantivy index. +/// Each segment can be passed to `do_migrate_segment` for bounded-memory batch migration. +pub fn count_eml_segments(legacy: &LegacyDirs) -> BichonResult { + let eml_index = Index::open_in_dir(&legacy.eml_dir) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + let reader = eml_index + .reader() + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + let searcher = reader.searcher(); + Ok(searcher.segment_readers().len()) +} + pub fn check_data_status() -> std::io::Result { let root_dir = PathBuf::from(&SETTINGS.bichon_root_dir); @@ -97,12 +113,22 @@ fn is_dir_not_empty(path: &PathBuf) -> std::io::Result { Ok(entries.next().is_some()) } -const PAGE_SIZE: usize = 100; - -pub fn do_migrate(legacy: LegacyDirs, new_dirs: NewDirs, mut on_progress: F) -> BichonResult<()> +/// Migrate all documents from a single EML segment to the new storage layout. +/// +/// This is the core of the batch migration strategy: each Process B invocation +/// handles exactly one EML segment, so peak memory is bounded by that segment's +/// size regardless of the total archive size. +pub fn do_migrate_segment( + batch_size: u32, + legacy: LegacyDirs, + new_dirs: NewDirs, + segment_index: usize, + mut on_progress: F, +) -> BichonResult<()> where F: FnMut(&str), { + // ── open legacy indices ──────────────────────────────────────────── let envelope_index = Index::open_in_dir(&legacy.envelope_dir) .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; let eml_index = Index::open_in_dir(&legacy.eml_dir) @@ -118,114 +144,172 @@ where let envelope_searcher = envelope_reader.searcher(); let eml_searcher = eml_reader.searcher(); - let total_count = envelope_searcher.num_docs(); - on_progress(&format!("TOTAL:{}", total_count)); - let ef = SchemaTools::envelope_fields(); let mf = SchemaTools::eml_fields(); - let mut writer = NewIndexWriter::open(new_dirs)?; + let eml_segments = eml_searcher.segment_readers(); + let eml_segment = eml_segments.get(segment_index).ok_or_else(|| { + raise_error!( + format!( + "segment index {} out of range ({} segments)", + segment_index, + eml_segments.len() + ), + ErrorCode::InternalError + ) + })?; - let mut offset = 0usize; - let mut total_migrated = 0usize; - let mut total_skipped = 0usize; + let num_docs = eml_segment.num_docs(); + if num_docs == 0 { + on_progress("TOTAL:0"); + on_progress("DONE:0:0"); + return Ok(()); + } - loop { - let page: Vec<(_, DocAddress)> = envelope_searcher - .search( - &AllQuery, - &TopDocs::with_limit(PAGE_SIZE) - .and_offset(offset) - .order_by_score(), - ) + on_progress(&format!("TOTAL:{}", num_docs)); + + let max_doc = eml_segment.max_doc(); + let ff = eml_segment.fast_fields(); + let f_id_col: Column = ff.u64("id").map_err(|e| { + raise_error!( + format!("failed to open f_id fast field: {e:#?}"), + ErrorCode::InternalError + ) + })?; + + // ── Phase 1: build eid → (uid, internal_date) from envelope, then drop it ── + let mut envelope_map: HashMap = HashMap::with_capacity(num_docs as usize); + + let mut env_scanned = 0u32; + let mut env_skipped = 0u32; + for doc_id in 0..max_doc { + if eml_segment.is_deleted(doc_id) { + continue; + } + let eid = f_id_col.values.get_val(doc_id); + + let term = Term::from_field_u64(ef.f_id, eid); + let query = TermQuery::new(term, IndexRecordOption::Basic); + let hits: Vec<(_, DocAddress)> = envelope_searcher + .search(&query, &TopDocs::with_limit(1).order_by_score()) .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; - if page.is_empty() { - break; - } - let fetched = page.len(); - - for (_, doc_address) in page { - let doc: TantivyDocument = envelope_searcher - .doc(doc_address) + if let Some((_, addr)) = hits.first() { + let env_doc: TantivyDocument = envelope_searcher + .doc(*addr) .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; - - let eid = match doc.get_first(ef.f_id).and_then(|v| v.as_u64()) { - Some(v) => v, - None => { - total_skipped += 1; - continue; - } - }; - let account_id = match doc.get_first(ef.f_account_id).and_then(|v| v.as_u64()) { - Some(v) => v, - None => { - total_skipped += 1; - continue; - } - }; - let mailbox_id = doc - .get_first(ef.f_mailbox_id) - .and_then(|v| v.as_u64()) - .unwrap_or(0); - let uid = doc + let uid = env_doc .get_first(ef.f_uid) .and_then(|v| v.as_u64()) .unwrap_or(0) as u32; - let internal_date = doc + let internal_date = env_doc .get_first(ef.f_internal_date) .and_then(|v| v.as_i64()) .unwrap_or(0); + envelope_map.insert(eid, (uid, internal_date)); + env_scanned += 1; + } else { + env_skipped += 1; + } - let eml_term = tantivy::Term::from_field_u64(mf.f_id, eid); - let eml_query = - tantivy::query::TermQuery::new(eml_term, tantivy::schema::IndexRecordOption::Basic); - let eml_hits: Vec<(_, DocAddress)> = eml_searcher - .search(&eml_query, &TopDocs::with_limit(1).order_by_score()) - .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + if env_scanned % 10 == 0 { + on_progress(&format!( + "PHASE1:{}/{} skipped:{}", + env_scanned, max_doc, env_skipped + )); + } + } - let eml_bytes = match eml_hits.first() { - Some((_, addr)) => { - let eml_doc: TantivyDocument = eml_searcher - .doc(*addr) - .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; - match eml_doc.get_first(mf.f_eml).and_then(|v| v.as_bytes()) { - Some(b) => b.to_vec(), - None => { - on_progress(&format!("WARN: Account {} ID {} eml field missing", account_id, eid)); - total_skipped += 1; - continue; - } - } - } + // Free the envelope index before the heavy EML processing. + drop(envelope_searcher); + drop(envelope_reader); + drop(envelope_index); + + // ── Phase 2: process EML docs, streaming one at a time ───────────── + let mut writer = NewIndexWriter::open(new_dirs)?; + + let mut total_migrated = 0usize; + let mut total_skipped = 0usize; + + // Recreate the StoreReader periodically to bound any internal caches. + //const CHUNK_SIZE: u32 = 3000; + let mut chunk_start = 0u32; + + while chunk_start < max_doc { + let chunk_end = (chunk_start + batch_size).min(max_doc); + let store_reader = eml_segment + .get_store_reader(2) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + + for doc_id in chunk_start..chunk_end { + if eml_segment.is_deleted(doc_id) { + continue; + } + + let eid = f_id_col.values.get_val(doc_id); + + let (uid, internal_date) = match envelope_map.get(&eid) { + Some(v) => *v, None => { - on_progress(&format!("WARN:Account {} ID {} eml not found", account_id, eid)); + on_progress(&format!("WARN: eid {} envelope not found", eid)); total_skipped += 1; continue; } }; - if let Err(e) = writer.ingest(&eml_bytes, account_id, mailbox_id, uid, internal_date) { - on_progress(&format!("ERROR:Account {} ID {} ingest failed: {}", account_id, eid, e)); + let eml_doc: TantivyDocument = store_reader + .get(doc_id) + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + + let account_id = match eml_doc.get_first(mf.f_account_id).and_then(|v| v.as_u64()) { + Some(v) => v, + None => { + on_progress(&format!("WARN: eid {} account_id missing", eid)); + total_skipped += 1; + continue; + } + }; + let mailbox_id = eml_doc + .get_first(mf.f_mailbox_id) + .and_then(|v| v.as_u64()) + .unwrap_or(0); + + // Borrow directly from eml_doc — no .to_vec() clone. + let eml_bytes = match eml_doc.get_first(mf.f_eml).and_then(|v| v.as_bytes()) { + Some(b) => b, + None => { + on_progress(&format!("WARN: eid {} eml bytes missing", eid)); + total_skipped += 1; + continue; + } + }; + + if let Err(e) = writer.ingest(eml_bytes, account_id, mailbox_id, uid, internal_date) { + on_progress(&format!( + "ERROR: Account {} eid {} ingest failed: {}", + account_id, eid, e + )); total_skipped += 1; continue; } total_migrated += 1; - if total_migrated % 100 == 0 || total_migrated == total_count as usize { - on_progress(&format!("PROGRESS:{}:{}", total_migrated, total_skipped)); + if total_migrated % 10 == 0 || total_migrated as u32 == num_docs { + on_progress(&format!("PROGRESS:{}:{}", total_migrated, num_docs)); } } - offset += fetched; - if fetched < PAGE_SIZE { - break; - } + drop(store_reader); + + // Flush Fjall buffers via ingestion API — bypasses memtable/WAL. + writer.flush_fjall_buffers()?; + + chunk_start = chunk_end; } - writer.commit()?; - + writer.finish_writers()?; + on_progress(&format!("DONE:{}:{}", total_migrated, total_skipped)); Ok(()) } diff --git a/crates/core/src/migrate/store.rs b/crates/core/src/migrate/store.rs index 691b204..a343e16 100644 --- a/crates/core/src/migrate/store.rs +++ b/crates/core/src/migrate/store.rs @@ -1,4 +1,4 @@ -use std::path::PathBuf; +use std::{path::PathBuf, time::Instant}; use bytes::Bytes; use mail_parser::MimeHeaders; @@ -13,7 +13,10 @@ use fjall::{ CompressionType, Database, Keyspace, KeyspaceCreateOptions, KvSeparationOptions, }; use mail_parser::MessageParser; -use tantivy::{Index, IndexWriter, TantivyDocument}; +use tantivy::{ + indexer::{LogMergePolicy, NoMergePolicy}, + Index, IndexWriter, TantivyDocument, +}; use uuid::Uuid; use crate::{ @@ -121,14 +124,16 @@ pub fn detach_attachments_standalone( } pub struct NewIndexWriter { - pub envelope_writer: IndexWriter, - pub attachment_writer: IndexWriter, + pub envelope_writer: Option, + pub attachment_writer: Option, pub email_ks: Keyspace, pub attachment_ks: Keyspace, pending: usize, + email_buf: Vec<(String, Vec)>, + attachment_buf: Vec<(String, Vec)>, } -const COMMIT_THRESHOLD: usize = 500; +//const COMMIT_THRESHOLD: usize = 500; impl NewIndexWriter { pub fn open(dirs: NewDirs) -> BichonResult { @@ -153,9 +158,15 @@ impl NewIndexWriter { .register("euro", EuroTokenizer::new()); let envelope_writer = envelope_index - .writer_with_num_threads(2, 128 * 1024 * 1024) + .writer_with_num_threads(3, 256 * 1024 * 1024) .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + // let mut merge_policy = LogMergePolicy::default(); + // merge_policy.set_min_num_segments(25); + // merge_policy.set_min_layer_size(10_000); + // merge_policy.set_max_docs_before_merge(100_000); + + envelope_writer.set_merge_policy(Box::new(NoMergePolicy)); // ── attachment index ───────────────────────────────────────────── std::fs::create_dir_all(&dirs.attachment_dir) .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; @@ -176,21 +187,30 @@ impl NewIndexWriter { .tokenizers() .register("euro", EuroTokenizer::new()); let attachment_writer = attachment_index - .writer_with_num_threads(2, 64 * 1024 * 1024) + .writer_with_num_threads(3, 256 * 1024 * 1024) .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + // let mut merge_policy = LogMergePolicy::default(); + // merge_policy.set_min_num_segments(25); + // merge_policy.set_min_layer_size(10_000); + // merge_policy.set_max_docs_before_merge(100_000); + + attachment_writer.set_merge_policy(Box::new(NoMergePolicy)); + // ── blob store ─────────────────────────────────────────────────── std::fs::create_dir_all(&dirs.storage_dir) .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; let db = Database::builder(&dirs.storage_dir) - .cache_size(64 * 1024 * 1024) + .cache_size(8 * 1024 * 1024) + .journal_compression(CompressionType::None) + .max_journaling_size(64 * 1024 * 1024) .open() .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; let email_ks = db .keyspace("email", || { KeyspaceCreateOptions::default() - .max_memtable_size(16 * 1024 * 1024) + .max_memtable_size(4 * 1024 * 1024) .data_block_size_policy(BlockSizePolicy::all(4 * 1024)) .data_block_compression_policy(CompressionPolicy::all(CompressionType::Lz4)) .with_kv_separation(Some( @@ -205,7 +225,7 @@ impl NewIndexWriter { let attachment_ks = db .keyspace("attachments", || { KeyspaceCreateOptions::default() - .max_memtable_size(16 * 1024 * 1024) + .max_memtable_size(4 * 1024 * 1024) .data_block_size_policy(BlockSizePolicy::all(4 * 1024)) .data_block_compression_policy(CompressionPolicy::all(CompressionType::Lz4)) .with_kv_separation(Some( @@ -218,11 +238,13 @@ impl NewIndexWriter { .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; Ok(Self { - envelope_writer, - attachment_writer, + envelope_writer: Some(envelope_writer), + attachment_writer: Some(attachment_writer), email_ks, attachment_ks, pending: 0, + email_buf: Vec::new(), + attachment_buf: Vec::new(), }) } @@ -299,23 +321,11 @@ impl NewIndexWriter { // ── detach attachments → blob ────────────────────────────────────── let (stripped_eml, attachment_output) = detach_attachments_standalone(eml_bytes, &message); - if !self - .email_ks - .contains_key(&email_content_hash) - .unwrap_or(false) - { - self.email_ks - .insert(&email_content_hash, stripped_eml.as_slice()) - .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; - } - - // write attachment blobs + // Buffer for bulk ingestion — sorted + flushed later. + self.email_buf + .push((email_content_hash.clone(), stripped_eml)); for (hash, data) in &attachment_output.blobs { - if !self.attachment_ks.contains_key(hash).unwrap_or(false) { - self.attachment_ks - .insert(hash, data.as_ref()) - .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; - } + self.attachment_buf.push((hash.clone(), data.to_vec())); } // ── build envelope doc ──────────────────────────────────────────── @@ -390,35 +400,151 @@ impl NewIndexWriter { let envelope_doc = ea.to_document(&text, 0)?; self.envelope_writer + .as_mut() + .unwrap() .add_document(envelope_doc) .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; for doc in attachment_docs { self.attachment_writer + .as_mut() + .unwrap() .add_document(doc) .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; } self.pending += 1; - if self.pending >= COMMIT_THRESHOLD { - self.commit()?; - } + // if self.pending >= COMMIT_THRESHOLD { + // self.commit()?; + // } Ok(()) } - pub fn commit(&mut self) -> BichonResult<()> { + /// Commit pending Tantivy documents (mid-stream) — frees the in-memory + /// term dictionary / postings that accumulate in the IndexWriter. + fn commit_tantivy(&mut self) -> BichonResult<()> { if self.pending == 0 { return Ok(()); } - self.envelope_writer - .commit() - .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; - self.attachment_writer - .commit() - .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; - tracing::info!(count = self.pending, "committed batch"); + println!("Tantivy committing... this may take 2-3 minutes, please wait."); + let start = Instant::now(); + if let Some(writer) = self.envelope_writer.as_mut() { + writer + .commit() + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + } + if let Some(writer) = self.attachment_writer.as_mut() { + writer + .commit() + .map_err(|e| raise_error!(format!("{e:#?}"), ErrorCode::InternalError))?; + } + println!("tantivy commit elasped: {:#?}", start.elapsed()); + tracing::info!(count = self.pending, "committed tantivy batch"); self.pending = 0; Ok(()) } + + /// Final commit + segment merge for Tantivy writers (called once at end). + pub fn finish_writers(&mut self) -> BichonResult<()> { + self.commit_tantivy()?; + + for (name, writer_opt) in [ + ("envelope", &mut self.envelope_writer), + ("attachment", &mut self.attachment_writer), + ] { + if let Some(writer) = writer_opt.as_mut() { + let reader = writer + .index() + .reader() + .map_err(|e| raise_error!(format!("{:#?}", e), ErrorCode::InternalError))?; + let seg_ids: Vec<_> = reader + .searcher() + .segment_readers() + .iter() + .map(|r| r.segment_id()) + .collect(); + println!("merging {} {} segments...", seg_ids.len(), name); + if seg_ids.len() > 1 { + let _ = writer.merge(&seg_ids); + } + } + + if let Some(writer) = writer_opt.take() { + println!("waiting for {} merge to finish...", name); + let start = std::time::Instant::now(); + let _ = writer.wait_merging_threads(); + println!("{} merge done: {:#?}", name, start.elapsed()); + } + } + + Ok(()) + } + + /// Sort buffered (hash, data) pairs, dedup, and write via Fjall's + /// ingestion API — writes SSTables directly, bypassing memtable and WAL. + /// Also commits the Tantivy writers to bound their in-memory state. + pub fn flush_fjall_buffers(&mut self) -> BichonResult<()> { + self.commit_tantivy()?; + + if !self.email_buf.is_empty() { + self.email_buf.sort_by(|a, b| a.0.cmp(&b.0)); + self.email_buf.dedup_by(|a, b| a.0 == b.0); + + let mut ingestion = self.email_ks.start_ingestion().map_err(|e| { + raise_error!( + format!("email ingestion start: {e:#?}"), + ErrorCode::InternalError + ) + })?; + for (hash, data) in &self.email_buf { + ingestion + .write(hash.as_bytes(), data.as_slice()) + .map_err(|e| { + raise_error!( + format!("email ingestion write: {e:#?}"), + ErrorCode::InternalError + ) + })?; + } + ingestion.finish().map_err(|e| { + raise_error!( + format!("email ingestion finish: {e:#?}"), + ErrorCode::InternalError + ) + })?; + self.email_buf.clear(); + } + + if !self.attachment_buf.is_empty() { + self.attachment_buf.sort_by(|a, b| a.0.cmp(&b.0)); + self.attachment_buf.dedup_by(|a, b| a.0 == b.0); + + let mut ingestion = self.attachment_ks.start_ingestion().map_err(|e| { + raise_error!( + format!("attachment ingestion start: {e:#?}"), + ErrorCode::InternalError + ) + })?; + for (hash, data) in &self.attachment_buf { + ingestion + .write(hash.as_bytes(), data.as_slice()) + .map_err(|e| { + raise_error!( + format!("attachment ingestion write: {e:#?}"), + ErrorCode::InternalError + ) + })?; + } + ingestion.finish().map_err(|e| { + raise_error!( + format!("attachment ingestion finish: {e:#?}"), + ErrorCode::InternalError + ) + })?; + self.attachment_buf.clear(); + } + + Ok(()) + } }