From addc2cbad0081809de8f9835fa94c63274c6d8ed Mon Sep 17 00:00:00 2001
From: "James R. Barlow" <james@purplerock.ca>
Date: Wed, 8 Jul 2020 00:43:13 -0700
Subject: [PATCH] Enable pikepdf mmap and set up signal handlers

---
 src/ocrmypdf/__main__.py    | 15 ++++++++++++++-
 src/ocrmypdf/_concurrent.py | 13 +++++++++++++
 src/ocrmypdf/_sync.py       |  7 +++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/ocrmypdf/__main__.py b/src/ocrmypdf/__main__.py
index 474c9678..9f15ba31 100755
--- a/src/ocrmypdf/__main__.py
+++ b/src/ocrmypdf/__main__.py
@@ -18,6 +18,7 @@
 
 import logging
 import os
+import signal
 import sys
 from multiprocessing import set_start_method
 
@@ -26,11 +27,20 @@ from ocrmypdf._plugin_manager import get_parser_options_plugins
 from ocrmypdf._sync import run_pipeline
 from ocrmypdf._validation import check_closed_streams, check_options
 from ocrmypdf.api import Verbosity, configure_logging
-from ocrmypdf.exceptions import BadArgsError, ExitCode, MissingDependencyError
+from ocrmypdf.exceptions import (
+    BadArgsError,
+    ExitCode,
+    InputFileError,
+    MissingDependencyError,
+)
 
 log = logging.getLogger('ocrmypdf')
 
 
+def sigbus(*args):
+    raise InputFileError("Lost access to the input file")
+
+
 def run(args=None):
     _parser, options, plugin_manager = get_parser_options_plugins(args=args)
 
@@ -62,6 +72,9 @@ def run(args=None):
         log.error(e)
         return ExitCode.missing_dependency
 
+    if hasattr(signal, 'SIGBUS'):
+        signal.signal(signal.SIGBUS, sigbus)
+
     result = run_pipeline(options=options, plugin_manager=plugin_manager)
     return result
 
diff --git a/src/ocrmypdf/_concurrent.py b/src/ocrmypdf/_concurrent.py
index f954d6d8..fdd9319e 100644
--- a/src/ocrmypdf/_concurrent.py
+++ b/src/ocrmypdf/_concurrent.py
@@ -28,6 +28,8 @@ from typing import Callable, Iterable, Optional
 
 from tqdm import tqdm
 
+from ocrmypdf.exceptions import InputFileError
+
 
 def log_listener(queue):
     """Listen to the worker processes and forward the messages to logging
@@ -53,12 +55,20 @@ def log_listener(queue):
             traceback.print_exc(file=sys.stderr)
 
 
+def process_sigbus(*args):
+    raise InputFileError("A worker process lost access to an input file")
+
+
 def process_init(queue, user_init):
     """Initialize a process pool worker"""
 
     # Ignore SIGINT (our parent process will kill us gracefully)
     signal.signal(signal.SIGINT, signal.SIG_IGN)
 
+    # Install SIGBUS handler (so our parent process can abort somewhat gracefully)
+    if hasattr(signal, 'SIGBUS'):
+        signal.signal(signal.SIGBUS, process_sigbus)
+
     # Reconfigure the root logger for this process to send all messages to a queue
     h = logging.handlers.QueueHandler(queue)
     root = logging.getLogger()
@@ -70,6 +80,9 @@ def process_init(queue, user_init):
 
 
 def thread_init(_queue, user_init):
+    # As a thread, block SIGBUS so the main thread deals with it...
+    if hasattr(signal, 'SIGBUS'):
+        signal.pthread_sigmask(signal.SIG_BLOCK, {signal.SIGBUS})
     if user_init:
         user_init()
 
diff --git a/src/ocrmypdf/_sync.py b/src/ocrmypdf/_sync.py
index c11b44fa..a3c64604 100644
--- a/src/ocrmypdf/_sync.py
+++ b/src/ocrmypdf/_sync.py
@@ -25,6 +25,7 @@ from pathlib import Path
 from tempfile import mkdtemp
 from typing import List, NamedTuple, Optional, Tuple
 
+import pikepdf
 import PIL
 
 from ocrmypdf._concurrent import exec_progress_pool
@@ -331,6 +332,12 @@ def run_pipeline(options, *, plugin_manager, api=False):
     ):
         debug_log_handler = configure_debug_logging(Path(work_folder) / "debug.log")
 
+    try:
+        if pikepdf._qpdf.set_access_default_mmap(True):
+            log.debug("pikepdf mmap enabled")
+    except AttributeError:
+        log.debug("pikepdf mmap not available")
+
     try:
         check_requested_output_file(options)
         start_input_file, original_filename = create_input_file(options, work_folder)