diff --git a/runtime/src/file_data.rs b/runtime/src/file_data.rs index cdf0de63..0978bc65 100644 --- a/runtime/src/file_data.rs +++ b/runtime/src/file_data.rs @@ -3,7 +3,7 @@ use std::collections::VecDeque; use std::convert::Infallible; use crate::api_token::APIToken; use crate::pandoc::PandocProcessBuilder; -use crate::pdfium::PdfiumInit; +use crate::pdfium::{with_pdfium_access, PdfiumInit}; use crate::prompt_injection::{Finding as PromptInjectionFinding, Sanitizer}; use async_stream::stream; use axum::extract::Query; @@ -1139,7 +1139,7 @@ async fn stream_pdf(file_path: &str) -> Result { return; } }; - let doc = match pdfium.load_pdf_from_file(&path, None) { + let doc = match with_pdfium_access(|| pdfium.load_pdf_from_file(&path, None)) { Ok(document) => document, Err(e) => { let _ = tx.blocking_send(Err(classify_pdf_load_error(&e).into())); @@ -1152,11 +1152,27 @@ async fn stream_pdf(file_path: &str) -> Result { let mut number_of_failed_pages = 0; let mut receiver_gone = false; - for (num_page, page) in doc.pages().iter().enumerate() { - let page_number = num_page + 1; + // + // One page at a time, rather than the whole document: somebody else may be reading a PDF + // of their own, and holding PDFium for a thousand-page manual would make them wait for all + // of it. Between two pages, their pages get their turn. + // + let page_count = with_pdfium_access(|| doc.pages().len()); + + for page_index in 0..page_count { + let page_number = page_index as usize + 1; number_of_pages = page_number; - let content = match page.text().map(|t| t.all()) { + // + // The page and its text are opened and closed inside this call. Letting them outlive + // it would close them without PDFium to ourselves, which is a call like any other. + // + let extracted = with_pdfium_access(|| doc + .pages() + .get(page_index) + .and_then(|page| page.text().map(|text| text.all()))); + + let content = match extracted { Ok(text_content) => text_content, Err(e) => { // @@ -1193,23 +1209,29 @@ async fn stream_pdf(file_path: &str) -> Result { if receiver_gone { debug!("The consumer stopped reading the PDF stream of '{path}' after {number_of_pages} page(s)."); - return; + } else { + debug!("Extracted {number_of_characters} readable character(s) from {number_of_pages} page(s) of '{path}'; failed pages: {number_of_failed_pages}."); + + // + // Without this marker, a PDF without a text layer and a broken extraction both arrive + // as an empty document, and the AI would answer as if the file had no content at all. + // + if number_of_characters == 0 { + warn!("No text could be extracted from '{path}': {number_of_pages} page(s), {number_of_failed_pages} failed page(s). The PDF may consist of scanned images without a text layer."); + + let _ = tx.blocking_send(Ok(Chunk::from_error(&ExtractionError::new( + ExtractionErrorCode::NoTextExtracted, + format!("No text could be extracted from {number_of_pages} page(s). The PDF may consist of scanned images without a text layer."), + )))); + } } - debug!("Extracted {number_of_characters} readable character(s) from {number_of_pages} page(s) of '{path}'; failed pages: {number_of_failed_pages}."); - // - // Without this marker, a PDF without a text layer and a broken extraction both arrive as - // an empty document, and the AI would answer as if the file had no content at all. + // Closing the document calls PDFium as well, so it waits for its turn like everything else. + // This is why the code above says what it has to say instead of returning early: the + // document has to be closed on every way out of here. // - if number_of_characters == 0 { - warn!("No text could be extracted from '{path}': {number_of_pages} page(s), {number_of_failed_pages} failed page(s). The PDF may consist of scanned images without a text layer."); - - let _ = tx.blocking_send(Ok(Chunk::from_error(&ExtractionError::new( - ExtractionErrorCode::NoTextExtracted, - format!("No text could be extracted from {number_of_pages} page(s). The PDF may consist of scanned images without a text layer."), - )))); - } + with_pdfium_access(move || drop(doc)); }); Ok(Box::pin(ReceiverStream::new(rx))) diff --git a/runtime/src/pdfium.rs b/runtime/src/pdfium.rs index be4d9cf6..616d301c 100644 --- a/runtime/src/pdfium.rs +++ b/runtime/src/pdfium.rs @@ -7,6 +7,30 @@ use log::{error, info, warn}; pub static PDFIUM_LIB_PATH: Lazy>> = Lazy::new(|| Mutex::new(None)); static PDFIUM: OnceCell = OnceCell::new(); +/// Grants one caller at a time the right to talk to PDFium. +static PDFIUM_ACCESS: Mutex<()> = Mutex::new(()); + +/// Runs the given action with PDFium all to itself. +/// +/// PDFium is not thread-safe, and nothing else guarantees that for us: the `thread_safe` feature of +/// `pdfium-render` has only granted `Send` and `Sync` since its release 0.9.0 and no longer locks +/// anything, although its documentation still says so. Two documents read at the same time -- a +/// chat attachment while a data source is being indexed, say -- therefore corrupt PDFium's memory +/// and take the whole runtime down with a segmentation fault. +/// +/// Every call to PDFium belongs in here, and so does everything holding a page or a document open: +/// closing them calls PDFium as well. What does not belong in here is anything that waits, our own +/// work on the extracted text above all, because everybody else waits along with it. +pub fn with_pdfium_access(action: impl FnOnce() -> T) -> T { + // + // A panic while reading a document poisons this lock. Refusing every PDF from then on would + // turn one broken document into a broken feature, so we take the lock either way: what the + // panic left behind is inside PDFium, not inside the unit value we guard with. + // + let _access = PDFIUM_ACCESS.lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + action() +} + pub trait PdfiumInit { fn ai_studio_init() -> Result<&'static Pdfium, Box>; }