perf: Adding thread local s3 cache db connections to fix performance bottleneck during bundle submission (#896)

perf: Enabling WAL mode by default for improved concurrency

perf: Adding a get_connection_entry method which accepts a db connection
parameter to allow for greater read concurrency.  get_connection_entry
doesn't take the main db_lock which protects the primary db_connection.

perf: Adding a get_local_connection entry which creates (if necessary) a
thread local db conneciton suitable for reading entries in
get_connection_entry.

Signed-off-by: Brian Axelson <86568017+baxeaz@users.noreply.github.com>

Author: baxeaz

src/deadline/job_attachments/caches/cache_db.py

@@ -6,6 +6,7 @@
 
 import logging
 import os
+import threading as _threading
 from abc import ABC
 from threading import Lock
 from typing import Optional
@@ -34,6 +35,8 @@ def __init__(
         self.cache_name: str = cache_name
         self.table_name: str = table_name
         self.create_query: str = create_query
+        self._local = _threading.local()
+        self._local_connections: set = set()
 
         try:
             # SQLite is included in Python installers, but might not exist if building python from source.
@@ -64,6 +67,7 @@ def __enter__(self):
                 self.db_connection: sqlite3.Connection = sqlite3.connect(
                     self.cache_dir, check_same_thread=False
                 )
+                self.db_connection.execute("PRAGMA journal_mode=WAL")
             except sqlite3.OperationalError as oe:
                 raise JobAttachmentsError(
                     f"Could not access cache file in {self.cache_dir}"
@@ -81,8 +85,35 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
         """Called when exiting the context manager."""
+
         if self.enabled:
+            import sqlite3
+
             self.db_connection.close()
+            for conn in self._local_connections:
+                try:
+                    conn.close()
+                except sqlite3.Error as e:
+                    logger.warning(f"SQLite connection failed to close with error {e}")
+
+            self._local_connections.clear()
+
+    def get_local_connection(self):
+        """Create and/or returns a thread local connection to the SQLite database."""
+        if not self.enabled:
+            return None
+        import sqlite3
+
+        if not hasattr(self._local, "connection"):
+            try:
+                self._local.connection = sqlite3.connect(self.cache_dir, check_same_thread=False)
+                self._local_connections.add(self._local.connection)
+            except sqlite3.OperationalError as oe:
+                raise JobAttachmentsError(
+                    f"Could not create connection to cache in {self.cache_dir}"
+                ) from oe
+
+        return self._local.connection
 
     @classmethod
     def get_default_cache_db_file_dir(cls) -> Optional[str]:
@@ -99,12 +130,23 @@ def remove_cache(self) -> None:
         """
         Removes the underlying cache contents from the file system.
         """
+
         if self.enabled:
+            import sqlite3
+
             self.db_connection.close()
+            conn_list = list(self._local_connections)
+            for conn in conn_list:
+                try:
+                    conn.close()
+                    self._local_connections.remove(conn)
+                except sqlite3.Error as e:
+                    logger.warning(f"SQLite connection failed to close with error {e}")
 
         logger.debug(f"The cache {self.cache_dir} will be removed")
         try:
             os.remove(self.cache_dir)
         except Exception as e:
             logger.error(f"Error occurred while removing the cache file {self.cache_dir}: {e}")
+
             raise e

src/deadline/job_attachments/caches/hash_cache.py

@@ -59,6 +59,35 @@ def __init__(self, cache_dir: Optional[str] = None) -> None:
             cache_dir=cache_dir,
         )
 
+    def get_connection_entry(
+        self, file_path_key: str, hash_algorithm: HashAlgorithm, connection
+    ) -> Optional[HashCacheEntry]:
+        """
+        Returns an entry from the hash cache, if it exists.  This is the "lockless" (Doesn't take
+        the main db_lock protecting db_connection) version of get_entry which expects a connection
+        parameter for the connection which will be used to read from the DB - this can generally
+        be the thread local connection returned by get_local_connection()
+        """
+        if not self.enabled:
+            return None
+
+        entry_vals = connection.execute(
+            f"SELECT * FROM {self.table_name} WHERE file_path=? AND hash_algorithm=?",
+            [
+                file_path_key.encode(encoding="utf-8", errors="surrogatepass"),
+                hash_algorithm.value,
+            ],
+        ).fetchone()
+        if entry_vals:
+            return HashCacheEntry(
+                file_path=str(entry_vals[0], encoding="utf-8", errors="surrogatepass"),
+                hash_algorithm=HashAlgorithm(entry_vals[1]),
+                file_hash=entry_vals[2],
+                last_modified_time=str(entry_vals[3]),
+            )
+        else:
+            return None
+
     def get_entry(
         self, file_path_key: str, hash_algorithm: HashAlgorithm
     ) -> Optional[HashCacheEntry]:
@@ -69,22 +98,7 @@ def get_entry(
             return None
 
         with self.db_lock, self.db_connection:
-            entry_vals = self.db_connection.execute(
-                f"SELECT * FROM {self.table_name} WHERE file_path=? AND hash_algorithm=?",
-                [
-                    file_path_key.encode(encoding="utf-8", errors="surrogatepass"),
-                    hash_algorithm.value,
-                ],
-            ).fetchone()
-            if entry_vals:
-                return HashCacheEntry(
-                    file_path=str(entry_vals[0], encoding="utf-8", errors="surrogatepass"),
-                    hash_algorithm=HashAlgorithm(entry_vals[1]),
-                    file_hash=entry_vals[2],
-                    last_modified_time=str(entry_vals[3]),
-                )
-            else:
-                return None
+            return self.get_connection_entry(file_path_key, hash_algorithm, self.db_connection)
 
     def put_entry(self, entry: HashCacheEntry) -> None:
         """Inserts or replaces an entry into the hash cache database after acquiring the lock."""

src/deadline/job_attachments/caches/s3_check_cache.py

@@ -56,6 +56,32 @@ def __init__(self, cache_dir: Optional[str] = None) -> None:
             cache_dir=cache_dir,
         )
 
+    def get_connection_entry(self, s3_key: str, connection) -> Optional[S3CheckCacheEntry]:
+        """
+        Returns an entry from the hash cache, if it exists.  This is the "lockless" (Doesn't take
+        the main db_lock protecting db_connection) version of get_entry which expects a connection
+        parameter for the connection which will be used to read from the DB - this can generally
+        be the thread local connection returned by get_local_connection()
+        """
+
+        entry_vals = connection.execute(
+            f"SELECT * FROM {self.table_name} WHERE s3_key=?",
+            [s3_key],
+        ).fetchone()
+        if entry_vals:
+            entry = S3CheckCacheEntry(
+                s3_key=entry_vals[0],
+                last_seen_time=str(entry_vals[1]),
+            )
+            try:
+                last_seen = datetime.fromtimestamp(float(entry.last_seen_time))
+                if (datetime.now() - last_seen).days < self.ENTRY_EXPIRY_DAYS:
+                    return entry
+            except ValueError:
+                logger.warning(f"Timestamp for S3 key {s3_key} is not valid. Ignoring.")
+
+        return None
+
     def get_entry(self, s3_key: str) -> Optional[S3CheckCacheEntry]:
         """
         Checks if an entry exists in the cache, and returns it if it hasn't expired.
@@ -64,23 +90,7 @@ def get_entry(self, s3_key: str) -> Optional[S3CheckCacheEntry]:
             return None
 
         with self.db_lock, self.db_connection:
-            entry_vals = self.db_connection.execute(
-                f"SELECT * FROM {self.table_name} WHERE s3_key=?",
-                [s3_key],
-            ).fetchone()
-            if entry_vals:
-                entry = S3CheckCacheEntry(
-                    s3_key=entry_vals[0],
-                    last_seen_time=str(entry_vals[1]),
-                )
-                try:
-                    last_seen = datetime.fromtimestamp(float(entry.last_seen_time))
-                    if (datetime.now() - last_seen).days < self.ENTRY_EXPIRY_DAYS:
-                        return entry
-                except ValueError:
-                    logger.warning(f"Timestamp for S3 key {s3_key} is not valid. Ignoring.")
-
-            return None
+            return self.get_connection_entry(s3_key, self.db_connection)
 
     def put_entry(self, entry: S3CheckCacheEntry) -> None:
         """Inserts or replaces an entry into the cache database."""

src/deadline/job_attachments/upload.py

@@ -593,8 +593,11 @@ def verify_hash_cache_integrity(
         random.shuffle(s3_upload_keys)
         sampled_cache_entries: List[S3CheckCacheEntry] = []
         with S3CheckCache(s3_check_cache_dir) as s3_cache:
+            local_connection = s3_cache.get_local_connection()
             for upload_key in s3_upload_keys:
-                this_entry = s3_cache.get_entry(s3_key=f"{s3_bucket}/{upload_key}")
+                this_entry = s3_cache.get_connection_entry(
+                    s3_key=f"{s3_bucket}/{upload_key}", connection=local_connection
+                )
                 if this_entry is not None:
                     sampled_cache_entries.append(this_entry)
                     if len(sampled_cache_entries) >= 30:
@@ -651,7 +654,9 @@ def upload_object_to_cas(
         s3_upload_key = self._generate_s3_upload_key(file, hash_algorithm, s3_cas_prefix)
         is_uploaded = False
 
-        if s3_check_cache.get_entry(s3_key=f"{s3_bucket}/{s3_upload_key}"):
+        if s3_check_cache.get_connection_entry(
+            s3_key=f"{s3_bucket}/{s3_upload_key}", connection=s3_check_cache.get_local_connection()
+        ):
             logger.debug(
                 f"skipping {local_path} because {s3_bucket}/{s3_upload_key} exists in the cache"
             )
@@ -1068,7 +1073,9 @@ def _process_input_path(
         file_status: FileStatus = FileStatus.UNCHANGED
         actual_modified_time = str(datetime.fromtimestamp(path.stat().st_mtime))
 
-        entry: Optional[HashCacheEntry] = hash_cache.get_entry(full_path, hash_alg)
+        entry: Optional[HashCacheEntry] = hash_cache.get_connection_entry(
+            full_path, hash_alg, connection=hash_cache.get_local_connection()
+        )
         if entry is not None:
             # If the file was modified, we need to rehash it
             if actual_modified_time != entry.last_modified_time: