fix: memory leak on apple native OCR #278

Cargo.toml

@@ -5,12 +5,9 @@ members = [
     "screenpipe-audio",
     "screenpipe-server", 
     "screenpipe-integrations",
-
 ]
 exclude = [
     "examples/apps/screenpipe-app-tauri/src-tauri",
-
-
 ]
 resolver = "2"
 
@@ -30,7 +27,7 @@ candle-nn = { package = "candle-nn", version = "0.6.0" }
 candle-transformers = { package = "candle-transformers", version = "0.6.0" }
 tokenizers = "0.19.1"
 tracing = "0.1.40"
-tracing-subscriber = "0.3.18"
+tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 tokio = { version = "1.15", features = ["full", "tracing"] }
 hf-hub = "0.3.0"
 crossbeam = "0.8.4"

examples/apps/screenpipe-app-tauri/src-tauri/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "screenpipe-app"
-version = "0.1.86"
+version = "0.1.87"
 description = ""
 authors = ["you"]
 license = ""

screenpipe-integrations/src/unstructured_ocr.rs

@@ -5,15 +5,14 @@ use std::collections::HashMap;
 use std::env;
 use std::io::Cursor;
 use std::io::Write;
-use std::sync::Arc;
 use tokio::time::{timeout, Duration};
 use reqwest::blocking::Client;
 use serde_json::Value;
 use tempfile::NamedTempFile;
 use anyhow::{Result, anyhow};
 use log::error;
 
-pub async fn perform_ocr_cloud(image: &Arc<DynamicImage>) -> Result<(String, String, Option<f64>)> {
+pub async fn perform_ocr_cloud(image: &DynamicImage) -> Result<(String, String, Option<f64>)> {
     let api_key = match env::var("UNSTRUCTURED_API_KEY") {
         Ok(key) => key,
         Err(_) => {

screenpipe-server/src/video.rs

@@ -46,7 +46,7 @@ impl VideoCapture {
                 result_sender,
                 Duration::from_secs_f64(1.0 / fps),
                 save_text_files,
-                ocr_engine,
+                *ocr_engine,
                 monitor_id,
             )
             .await;

screenpipe-vision/Cargo.toml

@@ -35,13 +35,18 @@ clap = { version = "4.0", features = ["derive"] }
 # Integrations
 screenpipe-integrations = { path = "../screenpipe-integrations" }
 
+tracing-subscriber = { workspace = true }
+tracing = { workspace = true }
+
 [dev-dependencies]
 tempfile = "3.3.0"
 criterion = { workspace = true }
 assert_cmd = "2.0.14"
 predicates = "3.1.0"
 assert_fs = "1.1.1"
 strsim = "0.10.0"
+memory-stats = "1.2.0"
+
 
 [build-dependencies]
 cc = "1.0"
@@ -61,6 +66,10 @@ harness = false
 name = "ocr_benchmark"
 harness = false
 
+[[bench]]
+name = "apple_leak_bench"
+harness = false
+
 
 [target.'cfg(target_os = "windows")'.dependencies]
 windows = { version = "0.58", features = ["Graphics_Imaging", "Media_Ocr", "Storage", "Storage_Streams"] }

screenpipe-vision/benches/apple_leak_bench.rs

@@ -0,0 +1,96 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use image::GenericImageView;
+use memory_stats::memory_stats;
+use screenpipe_vision::perform_ocr_apple;
+use std::path::PathBuf;
+
+fn bytes_to_mb(bytes: usize) -> f64 {
+    bytes as f64 / (1024.0 * 1024.0)
+}
+
+fn bytes_to_gb(bytes: usize) -> f64 {
+    bytes as f64 / (1024.0 * 1024.0 * 1024.0)
+}
+
+fn apple_ocr_benchmark(c: &mut Criterion) {
+    let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+    path.push("tests");
+    path.push("testing_OCR.png");
+
+    let image = image::open(&path).expect("Failed to open image");
+    println!("Image dimensions: {:?}", image.dimensions());
+
+    let mut group = c.benchmark_group("apple_ocr");
+    group.sample_size(100); // Increased sample size
+    group.measurement_time(std::time::Duration::from_secs(60)); // Run for at least 60 seconds
+    group.bench_function("perform_ocr_apple", |b| {
+        b.iter_custom(|iters| {
+            let start = std::time::Instant::now();
+            let mut initial_memory = 0;
+            let mut final_memory = 0;
+            let mut max_memory = 0;
+
+            for i in 0..iters {
+                if i == 0 {
+                    if let Some(usage) = memory_stats() {
+                        initial_memory = usage.physical_mem;
+                        max_memory = initial_memory;
+                    }
+                }
+
+                let result = perform_ocr_apple(&image);
+                assert!(
+                    result.contains("receiver_count"),
+                    "OCR failed: {:?}",
+                    result
+                );
+
+                if let Some(usage) = memory_stats() {
+                    final_memory = usage.physical_mem;
+                    max_memory = max_memory.max(final_memory);
+                }
+
+                if i % 10 == 0 {
+                    println!(
+                        "Iteration {}: Current memory usage: {:.2} MB ({:.3} GB)",
+                        i,
+                        bytes_to_mb(final_memory),
+                        bytes_to_gb(final_memory)
+                    );
+                }
+            }
+
+            println!(
+                "Initial memory usage: {:.2} MB ({:.3} GB)",
+                bytes_to_mb(initial_memory),
+                bytes_to_gb(initial_memory)
+            );
+            println!(
+                "Final memory usage: {:.2} MB ({:.3} GB)",
+                bytes_to_mb(final_memory),
+                bytes_to_gb(final_memory)
+            );
+            println!(
+                "Max memory usage: {:.2} MB ({:.3} GB)",
+                bytes_to_mb(max_memory),
+                bytes_to_gb(max_memory)
+            );
+            println!(
+                "Total memory difference: {:.2} MB ({:.3} GB)",
+                bytes_to_mb(final_memory - initial_memory),
+                bytes_to_gb(final_memory - initial_memory)
+            );
+            println!(
+                "Max memory difference: {:.2} MB ({:.3} GB)",
+                bytes_to_mb(max_memory - initial_memory),
+                bytes_to_gb(max_memory - initial_memory)
+            );
+
+            start.elapsed()
+        });
+    });
+    group.finish();
+}
+
+criterion_group!(benches, apple_ocr_benchmark);
+criterion_main!(benches);

screenpipe-vision/src/apple.rs

@@ -3,6 +3,20 @@ use log::error;
 use std::ffi::CStr;
 use std::os::raw::{c_char, c_uchar};
 
+use std::ops::Drop;
+
+struct OcrResultGuard(*mut c_char);
+
+impl Drop for OcrResultGuard {
+    fn drop(&mut self) {
+        unsafe {
+            if !self.0.is_null() {
+                free_string(self.0);
+            }
+        }
+    }
+}
+
 #[cfg(target_os = "macos")]
 #[link(name = "screenpipe")]
 extern "C" {
@@ -12,6 +26,7 @@ extern "C" {
         width: i32,
         height: i32,
     ) -> *mut c_char;
+    fn free_string(ptr: *mut c_char);
 }
 #[cfg(target_os = "macos")]
 pub fn perform_ocr_apple(image: &DynamicImage) -> String {
@@ -26,8 +41,8 @@ pub fn perform_ocr_apple(image: &DynamicImage) -> String {
             width as i32,
             height as i32,
         );
+        let _guard = OcrResultGuard(result_ptr);
         let result = CStr::from_ptr(result_ptr).to_string_lossy().into_owned();
-        libc::free(result_ptr as *mut libc::c_void);
         result
     }
 }
@@ -51,8 +66,7 @@ pub fn parse_apple_ocr_result(json_result: &str) -> (String, String, Option<f64>
         .as_array()
         .unwrap_or(&vec![])
         .clone();
-    let overall_confidence = parsed_result["overallConfidence"]
-        .as_f64();
+    let overall_confidence = parsed_result["overallConfidence"].as_f64();
 
     let json_output: Vec<serde_json::Value> = text_elements
         .iter()
@@ -80,4 +94,4 @@ pub fn parse_apple_ocr_result(json_result: &str) -> (String, String, Option<f64>
     });
 
     (text, json_output_string, overall_confidence)
-}
+}

screenpipe-vision/src/bin/screenpipe-vision.rs

@@ -1,7 +1,8 @@
 use clap::Parser;
 use screenpipe_vision::{continuous_capture, monitor::get_default_monitor, OcrEngine};
-use std::{sync::Arc, time::Duration};
+use std::time::Duration;
 use tokio::sync::mpsc::channel;
+use tracing_subscriber::{fmt::format::FmtSpan, EnvFilter};
 
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
@@ -17,6 +18,14 @@ struct Cli {
 
 #[tokio::main]
 async fn main() {
+    tracing_subscriber::fmt()
+        .with_env_filter(
+            EnvFilter::from_default_env()
+                .add_directive(tracing::Level::DEBUG.into())
+                .add_directive("tokenizers=error".parse().unwrap()),
+        )
+        .with_span_events(FmtSpan::CLOSE)
+        .init();
     let cli = Cli::parse();
 
     let (result_tx, mut result_rx) = channel(512);
@@ -31,11 +40,7 @@ async fn main() {
             result_tx,
             Duration::from_secs_f32(1.0 / cli.fps),
             save_text_files,
-            Arc::new(if cfg!(target_os = "macos") {
-                OcrEngine::AppleNative
-            } else {
-                OcrEngine::Tesseract
-            }),
+            OcrEngine::AppleNative,
             id,
         )
         .await