kreuzberg 4.0.0.pre.rc.6 → 4.0.0.rc1

data/vendor/kreuzberg/src/bin/profile_extract.rs

@@ -0,0 +1,455 @@
+//! Profile memory usage for extracting documents.
+//!
+//! This utility can run against a single file or a batch of files (via `--input-list`).
+//! For each input it prints or writes a JSON object containing duration, peak RSS,
+//! optional flamegraph path, and the top hot functions when CPU profiling is enabled.
+
+use std::env;
+use std::fs::{File, create_dir_all};
+use std::io::{BufRead, BufReader};
+use std::path::{Path, PathBuf};
+use std::time::Instant;
+
+use kreuzberg::core::config::ExtractionConfig;
+use kreuzberg::core::extractor::extract_file_sync;
+use serde::Serialize;
+
+#[cfg(feature = "profiling")]
+use pprof::{ProfilerGuardBuilder, Report};
+
+#[cfg(feature = "profiling")]
+use std::collections::HashMap;
+
+#[cfg(target_os = "macos")]
+fn normalize_rss(value: i64) -> i64 {
+    value / 1024
+}
+
+#[cfg(all(unix, not(target_os = "macos")))]
+fn normalize_rss(value: i64) -> i64 {
+    value
+}
+
+#[cfg(not(unix))]
+fn normalize_rss(value: i64) -> i64 {
+    value
+}
+
+#[cfg(unix)]
+fn max_rss_kb() -> Option<i64> {
+    use std::mem::MaybeUninit;
+
+    let mut usage = MaybeUninit::<libc::rusage>::uninit();
+    let rc = unsafe { libc::getrusage(libc::RUSAGE_SELF, usage.as_mut_ptr()) };
+    if rc != 0 {
+        return None;
+    }
+    let usage = unsafe { usage.assume_init() };
+    Some(normalize_rss(usage.ru_maxrss))
+}
+
+#[cfg(not(unix))]
+fn max_rss_kb() -> Option<i64> {
+    None
+}
+
+#[derive(Debug)]
+struct Options {
+    inputs: Vec<PathBuf>,
+    input_list: Option<PathBuf>,
+    flamegraph: Option<PathBuf>,
+    flamegraph_dir: Option<PathBuf>,
+    output_json: Option<PathBuf>,
+    output_dir: Option<PathBuf>,
+}
+
+#[derive(Serialize, Clone)]
+struct FunctionSample {
+    function: String,
+    samples: i64,
+    percentage: f64,
+}
+
+#[derive(Serialize, Clone)]
+struct ProfileOutput {
+    input: String,
+    duration_secs: f64,
+    peak_rss_kb: Option<i64>,
+    delta_rss_kb: Option<i64>,
+    flamegraph: Option<String>,
+    top_functions: Option<Vec<FunctionSample>>,
+}
+
+fn print_usage() {
+    eprintln!(
+        "Usage: profile_extract [options] <file ...>\n\nOptions:\n  --flamegraph <path>         Write flamegraph SVG (single input)\n  --flamegraph-dir <dir>      Write flamegraph SVGs for each input\n  --output-json <path>        Write JSON output (single input)\n  --output-dir <dir>          Write per-file JSON outputs to directory\n  --input-list <path>         File with newline-separated input paths\n  -h, --help                  Show this help message"
+    );
+}
+
+fn parse_options() -> Options {
+    let mut args = env::args().skip(1);
+    let mut inputs = Vec::new();
+    let mut input_list = None;
+    let mut flamegraph = None;
+    let mut flamegraph_dir = None;
+    let mut output_json = None;
+    let mut output_dir = None;
+
+    while let Some(arg) = args.next() {
+        match arg.as_str() {
+            "--flamegraph" => {
+                let path = args.next().unwrap_or_else(|| {
+                    print_usage();
+                    std::process::exit(64);
+                });
+                flamegraph = Some(PathBuf::from(path));
+            }
+            "--flamegraph-dir" => {
+                let path = args.next().unwrap_or_else(|| {
+                    print_usage();
+                    std::process::exit(64);
+                });
+                flamegraph_dir = Some(PathBuf::from(path));
+            }
+            "--output-json" => {
+                let path = args.next().unwrap_or_else(|| {
+                    print_usage();
+                    std::process::exit(64);
+                });
+                output_json = Some(PathBuf::from(path));
+            }
+            "--output-dir" => {
+                let path = args.next().unwrap_or_else(|| {
+                    print_usage();
+                    std::process::exit(64);
+                });
+                output_dir = Some(PathBuf::from(path));
+            }
+            "--input-list" => {
+                let path = args.next().unwrap_or_else(|| {
+                    print_usage();
+                    std::process::exit(64);
+                });
+                input_list = Some(PathBuf::from(path));
+            }
+            "-h" | "--help" => {
+                print_usage();
+                std::process::exit(0);
+            }
+            _ if arg.starts_with("--") => {
+                eprintln!("Unknown option: {arg}");
+                print_usage();
+                std::process::exit(64);
+            }
+            _ => inputs.push(PathBuf::from(arg)),
+        }
+    }
+
+    Options {
+        inputs,
+        input_list,
+        flamegraph,
+        flamegraph_dir,
+        output_json,
+        output_dir,
+    }
+}
+
+fn sanitize_file_name(path: &Path) -> String {
+    let name_owned;
+    let name = match path.file_name().and_then(|n| n.to_str()) {
+        Some(value) => value,
+        None => {
+            name_owned = path.to_string_lossy().into_owned();
+            &name_owned
+        }
+    };
+    let sanitized: String = name
+        .chars()
+        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
+        .collect();
+    if sanitized.is_empty() {
+        "output".to_string()
+    } else {
+        sanitized
+    }
+}
+
+fn read_inputs_from_file(list_path: &Path) -> Result<Vec<PathBuf>, String> {
+    let file = File::open(list_path).map_err(|e| format!("Failed to open input list {}: {e}", list_path.display()))?;
+    let reader = BufReader::new(file);
+    let mut inputs = Vec::new();
+    for line in reader.lines() {
+        let line = line.map_err(|e| format!("Failed to read line from {}: {e}", list_path.display()))?;
+        let trimmed = line.trim();
+        if trimmed.is_empty() || trimmed.starts_with('#') {
+            continue;
+        }
+        inputs.push(PathBuf::from(trimmed));
+    }
+    Ok(inputs)
+}
+
+fn main() {
+    let options = parse_options();
+
+    let mut targets = options.inputs.clone();
+    if let Some(list_path) = &options.input_list {
+        match read_inputs_from_file(list_path) {
+            Ok(mut list_inputs) => targets.append(&mut list_inputs),
+            Err(err) => {
+                eprintln!("{err}");
+                std::process::exit(66);
+            }
+        }
+    }
+
+    if targets.is_empty() {
+        eprintln!("No input files specified");
+        print_usage();
+        std::process::exit(64);
+    }
+
+    let multiple = targets.len() > 1;
+
+    if multiple && options.flamegraph.is_some() && options.flamegraph_dir.is_none() {
+        eprintln!("Note: --flamegraph applies to a single input. Use --flamegraph-dir for batch runs.");
+    }
+
+    if multiple && options.output_json.is_some() && options.output_dir.is_none() {
+        eprintln!("Note: --output-json applies to a single input. Use --output-dir for batch runs.");
+    }
+
+    if let Some(dir) = &options.output_dir
+        && let Err(err) = create_dir_all(dir)
+    {
+        eprintln!("Failed to create output directory {}: {err}", dir.display());
+    }
+
+    if let Some(dir) = &options.flamegraph_dir
+        && let Err(err) = create_dir_all(dir)
+    {
+        eprintln!("Failed to create flamegraph directory {}: {err}", dir.display());
+    }
+
+    let mut aggregated_results: Vec<ProfileOutput> = Vec::new();
+
+    for target in targets {
+        let flamegraph_path = if let Some(dir) = &options.flamegraph_dir {
+            Some(dir.join(format!("{}.svg", sanitize_file_name(&target))))
+        } else if !multiple {
+            options.flamegraph.clone()
+        } else {
+            None
+        };
+
+        let output_json_path = if let Some(dir) = &options.output_dir {
+            Some(dir.join(format!("{}.json", sanitize_file_name(&target))))
+        } else if !multiple {
+            options.output_json.clone()
+        } else {
+            None
+        };
+
+        match run_profile(&target, flamegraph_path.clone()) {
+            Ok(profile) => {
+                if let Some(json_path) = output_json_path {
+                    if let Some(parent) = json_path.parent()
+                        && let Err(err) = create_dir_all(parent)
+                    {
+                        eprintln!("Failed to create output directory {}: {err}", parent.display());
+                    }
+
+                    match File::create(&json_path) {
+                        Ok(file) => {
+                            if let Err(err) = serde_json::to_writer_pretty(file, &profile) {
+                                eprintln!("Failed to write JSON output {}: {err}", json_path.display());
+                            } else {
+                                eprintln!("Profile summary written to {}", json_path.display());
+                            }
+                        }
+                        Err(err) => eprintln!("Failed to create JSON output file {}: {err}", json_path.display()),
+                    }
+                } else {
+                    aggregated_results.push(profile);
+                }
+            }
+            Err(err) => {
+                eprintln!("{}: {err}", target.display());
+            }
+        }
+    }
+
+    if options.output_json.is_none() && options.output_dir.is_none() {
+        if aggregated_results.len() == 1 {
+            if let Ok(json) = serde_json::to_string_pretty(&aggregated_results[0]) {
+                println!("{json}");
+            }
+        } else if !aggregated_results.is_empty()
+            && let Ok(json) = serde_json::to_string_pretty(&aggregated_results)
+        {
+            println!("{json}");
+        }
+    }
+}
+
+fn run_profile(path: &Path, flamegraph_path: Option<PathBuf>) -> Result<ProfileOutput, String> {
+    if !path.exists() {
+        return Err("Input file does not exist".to_string());
+    }
+
+    #[cfg(feature = "profiling")]
+    let guard = if flamegraph_path.is_some() {
+        #[cfg_attr(not(target_os = "macos"), allow(unused_mut))]
+        let mut builder = ProfilerGuardBuilder::default().frequency(100);
+
+        #[cfg(target_os = "macos")]
+        {
+            builder = builder.blocklist(&[
+                "libsystem_kernel.dylib",
+                "libsystem_pthread.dylib",
+                "libsystem_platform.dylib",
+                "libdyld.dylib",
+            ]);
+        }
+
+        match builder.build() {
+            Ok(guard) => Some(guard),
+            Err(err) => {
+                eprintln!("Failed to start profiler: {err}");
+                None
+            }
+        }
+    } else {
+        None
+    };
+
+    #[cfg(not(feature = "profiling"))]
+    if flamegraph_path.is_some() {
+        eprintln!(
+            "--flamegraph requested but build missing 'profiling' feature; recompile with `--features profiling`."
+        );
+    }
+
+    let start_rss = max_rss_kb();
+    let start = Instant::now();
+
+    let config = ExtractionConfig::default();
+    let result = extract_file_sync(path, None, &config).map_err(|e| format!("Extraction failed: {e:?}"))?;
+    let _ = result;
+    let duration = start.elapsed();
+    let end_rss = max_rss_kb();
+
+    #[cfg(feature = "profiling")]
+    let (flamegraph_path_str, top_functions) = match (flamegraph_path.clone(), guard) {
+        (Some(path), Some(guard)) => match guard.report().build() {
+            Ok(report) => {
+                if let Some(parent) = path.parent()
+                    && let Err(err) = create_dir_all(parent)
+                {
+                    eprintln!("Failed to create flamegraph directory: {err}");
+                }
+
+                match File::create(&path) {
+                    Ok(mut file) => {
+                        if let Err(err) = report.flamegraph(&mut file) {
+                            eprintln!("Failed to write flamegraph: {err}");
+                        }
+                    }
+                    Err(err) => {
+                        eprintln!("Failed to create flamegraph file {}: {err}", path.display());
+                    }
+                }
+
+                let tops = summarize_top_functions(&report, 10);
+                let filtered = if tops.is_empty() { None } else { Some(tops) };
+                (Some(path.display().to_string()), filtered)
+            }
+            Err(err) => {
+                eprintln!("Failed to build profiling report: {err}");
+                (None, None)
+            }
+        },
+        (Some(path), None) => {
+            eprintln!("Profiler guard was not initialised; skipping flamegraph generation");
+            if let Some(parent) = path.parent() {
+                let _ = create_dir_all(parent);
+            }
+            (None, None)
+        }
+        _ => (None, None),
+    };
+
+    #[cfg(not(feature = "profiling"))]
+    let (flamegraph_path_str, top_functions): (Option<String>, Option<Vec<FunctionSample>>) = (None, None);
+
+    let peak_kb = end_rss.or(start_rss);
+    let delta_kb = match (start_rss, end_rss) {
+        (Some(before), Some(after)) => Some(after.saturating_sub(before)),
+        _ => None,
+    };
+
+    Ok(ProfileOutput {
+        input: path.display().to_string(),
+        duration_secs: duration.as_secs_f64(),
+        peak_rss_kb: peak_kb,
+        delta_rss_kb: delta_kb,
+        flamegraph: flamegraph_path_str,
+        top_functions,
+    })
+}
+
+#[cfg(feature = "profiling")]
+fn summarize_top_functions(report: &Report, limit: usize) -> Vec<FunctionSample> {
+    let mut totals: HashMap<String, i64> = HashMap::new();
+
+    for (frames, count) in &report.data {
+        let count = *count as i64;
+        if count <= 0 {
+            continue;
+        }
+
+        for frame_symbols in &frames.frames {
+            for symbol in frame_symbols {
+                let name = symbol.name();
+                *totals.entry(name).or_insert(0) += count;
+            }
+        }
+    }
+
+    let total_counts: i64 = totals.values().copied().sum();
+
+    let mut summary: Vec<FunctionSample> = totals
+        .into_iter()
+        .map(|(function, samples)| {
+            let percentage = if total_counts > 0 {
+                (samples as f64 / total_counts as f64) * 100.0
+            } else {
+                0.0
+            };
+            FunctionSample {
+                function,
+                samples,
+                percentage,
+            }
+        })
+        .collect();
+
+    summary.sort_by(|a, b| b.samples.cmp(&a.samples));
+
+    let filtered: Vec<FunctionSample> = summary
+        .iter()
+        .filter(|entry| {
+            let name = entry.function.as_str();
+            !name.starts_with("__") && !name.contains("libsystem") && !name.contains("dyld")
+        })
+        .take(limit)
+        .cloned()
+        .collect();
+
+    if filtered.is_empty() {
+        summary.into_iter().take(limit).collect()
+    } else {
+        filtered
+    }
+}

data/vendor/kreuzberg/src/cache/mod.rs

@@ -201,13 +201,6 @@ impl GenericCache {
         }
     }
 
-    #[cfg_attr(feature = "otel", tracing::instrument(
-        skip(self),
-        fields(
-            cache.hit = tracing::field::Empty,
-            cache.key = %cache_key,
-        )
-    ))]
     pub fn get(&self, cache_key: &str, source_file: Option<&str>) -> Result<Option<Vec<u8>>> {
         let cache_path = self.get_cache_path(cache_key);
 
@@ -217,24 +210,16 @@ impl GenericCache {
                 .lock()
                 .map_err(|e| KreuzbergError::LockPoisoned(format!("Deleting files mutex poisoned: {}", e)))?;
             if deleting.contains(&cache_path) {
-                #[cfg(feature = "otel")]
-                tracing::Span::current().record("cache.hit", false);
                 return Ok(None);
             }
         }
 
         if !self.is_valid(&cache_path, source_file) {
-            #[cfg(feature = "otel")]
-            tracing::Span::current().record("cache.hit", false);
             return Ok(None);
         }
 
         match fs::read(&cache_path) {
-            Ok(content) => {
-                #[cfg(feature = "otel")]
-                tracing::Span::current().record("cache.hit", true);
-                Ok(Some(content))
-            }
+            Ok(content) => Ok(Some(content)),
             Err(_) => {
                 // Best-effort cleanup of corrupted cache files ~keep
                 if let Err(e) = fs::remove_file(&cache_path) {
@@ -243,30 +228,21 @@ impl GenericCache {
                 if let Err(e) = fs::remove_file(self.get_metadata_path(cache_key)) {
                     tracing::debug!("Failed to remove corrupted metadata file: {}", e);
                 }
-                #[cfg(feature = "otel")]
-                tracing::Span::current().record("cache.hit", false);
                 Ok(None)
             }
         }
     }
 
-    #[cfg_attr(feature = "otel", tracing::instrument(
-        skip(self, data),
-        fields(
-            cache.key = %cache_key,
-            cache.size_bytes = data.len(),
-        )
-    ))]
     pub fn set(&self, cache_key: &str, data: Vec<u8>, source_file: Option<&str>) -> Result<()> {
         let cache_path = self.get_cache_path(cache_key);
 
-        fs::write(&cache_path, &data)
+        fs::write(&cache_path, data)
             .map_err(|e| KreuzbergError::cache(format!("Failed to write cache file: {}", e)))?;
 
         self.save_metadata(cache_key, source_file);
 
         let count = self.write_counter.fetch_add(1, Ordering::Relaxed);
-        if count.is_multiple_of(100)
+        if count % 100 == 0
             && let Some(cache_path_str) = self.cache_dir.to_str()
         {
             // Cache cleanup failure - safe to ignore, cache is optional fallback ~keep

data/vendor/kreuzberg/src/core/batch_mode.rs

@@ -33,63 +33,3 @@ where
 {
     BATCH_MODE.scope(Cell::new(true), future).await
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[tokio::test]
-    async fn test_batch_mode_not_set_by_default() {
-        let result = is_batch_mode();
-        assert!(!result, "batch mode should be false by default");
-    }
-
-    #[tokio::test]
-    async fn test_with_batch_mode_sets_flag() {
-        let result = with_batch_mode(async { is_batch_mode() }).await;
-
-        assert!(result, "batch mode should be true inside with_batch_mode");
-    }
-
-    #[tokio::test]
-    async fn test_batch_mode_scoped_to_future() {
-        assert!(!is_batch_mode(), "batch mode should be false before");
-
-        with_batch_mode(async {
-            assert!(is_batch_mode(), "batch mode should be true inside");
-        })
-        .await;
-
-        assert!(!is_batch_mode(), "batch mode should be false after future completes");
-    }
-
-    #[tokio::test]
-    async fn test_nested_batch_mode_calls() {
-        let result = with_batch_mode(async {
-            let outer = is_batch_mode();
-            let inner = with_batch_mode(async { is_batch_mode() }).await;
-            (outer, inner)
-        })
-        .await;
-
-        assert!(result.0, "outer batch mode should be true");
-        assert!(result.1, "inner batch mode should be true");
-    }
-
-    #[tokio::test]
-    async fn test_batch_mode_unaffected_after_with_batch_mode() {
-        with_batch_mode(async {
-            assert!(is_batch_mode(), "first call should set batch mode");
-        })
-        .await;
-
-        assert!(!is_batch_mode(), "batch mode should be false between calls");
-
-        with_batch_mode(async {
-            assert!(is_batch_mode(), "second call should set batch mode");
-        })
-        .await;
-
-        assert!(!is_batch_mode(), "batch mode should be false after all calls");
-    }
-}