fix: route insert-runner thread safety through a client capability

tests/test_rate_runner_thread_safety.py

@@ -0,0 +1,85 @@
+"""Unit tests for RatedMultiThreadingInsertRunner.send_insert_task routing.
+
+These exercise the thread-safety routing without a live database: a thread-safe
+client inserts through the shared object, while a non-thread-safe client is routed
+through copy_for_thread() + init() so each worker owns its own connection.
+"""
+
+from contextlib import contextmanager
+from copy import deepcopy
+
+from vectordb_bench.backend.clients.api import VectorDB
+from vectordb_bench.backend.runner.rate_runner import RatedMultiThreadingInsertRunner
+
+
+class FakeDB(VectorDB):
+    def __init__(self, thread_safe: bool = True, name: str = "Fake"):
+        self.thread_safe = thread_safe
+        self.name = name
+        self.init_calls = 0
+        self.inserted: list[int] = []
+
+    @contextmanager
+    def init(self):
+        self.init_calls += 1
+        yield
+
+    def insert_embeddings(self, embeddings: list, metadata: list, **kwargs):
+        self.inserted.append(len(embeddings))
+        return len(embeddings), None
+
+    def search_embedding(self, *args, **kwargs):
+        return []
+
+    def optimize(self, data_size: int | None = None):
+        return
+
+
+class CapturingNonThreadSafeDB(FakeDB):
+    """Records the per-thread copy so the test can assert against it."""
+
+    def __init__(self):
+        super().__init__(thread_safe=False, name="Capturing")
+        self.thread_copy: CapturingNonThreadSafeDB | None = None
+
+    def copy_for_thread(self) -> "VectorDB":
+        c = deepcopy(self)
+        self.thread_copy = c
+        return c
+
+
+def _runner(db: VectorDB):
+    return RatedMultiThreadingInsertRunner(rate=10, db=db, dataset_iter=None, batch_size=10)
+
+
+def test_thread_safe_client_uses_shared_object():
+    db = FakeDB(thread_safe=True)
+    _runner(db).send_insert_task(db, [[0.1], [0.2]], ["a", "b"])
+
+    # Shared client inserts directly, no per-thread copy/init from the runner.
+    assert db.inserted == [2]
+    assert db.init_calls == 0
+
+
+def test_non_thread_safe_client_routes_through_copy():
+    db = CapturingNonThreadSafeDB()
+    _runner(db).send_insert_task(db, [[0.1], [0.2]], ["a", "b"])
+
+    copy = db.thread_copy
+    assert copy is not None
+    assert copy is not db
+    # Insert + init happen on the thread-local copy, never the parent.
+    assert copy.init_calls == 1
+    assert copy.inserted == [2]
+    assert db.inserted == []
+    assert db.init_calls == 0
+
+
+def test_default_copy_for_thread_is_a_distinct_deep_copy():
+    db = FakeDB(thread_safe=False)
+    copy = db.copy_for_thread()
+
+    assert copy is not db
+    assert isinstance(copy, FakeDB)
+    copy.inserted.append(1)
+    assert db.inserted == []

vectordb_bench/backend/clients/adbpg/adbpg.py

@@ -3,6 +3,7 @@
 import logging
 from collections.abc import Generator, Sequence
 from contextlib import contextmanager
+from copy import copy
 from typing import Any
 
 import numpy as np
@@ -105,6 +106,14 @@ def _create_connection(**kwargs) -> tuple[Connection, Cursor]:
 
         return conn, cursor
 
+    def copy_for_thread(self) -> "VectorDB":
+        # psycopg holds an open socket that can't be deep-copied; shallow-copy
+        # and drop the connection so init() reconnects inside the worker thread.
+        db_copy = copy(self)
+        db_copy.conn = None
+        db_copy.cursor = None
+        return db_copy
+
     def _generate_search_query(self) -> sql.Composed:
         search_param = self.case_config.search_param()
         distance_operator = {

vectordb_bench/backend/clients/api.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
+from copy import deepcopy
 from enum import StrEnum
 from typing import ClassVar
 
@@ -199,6 +200,17 @@ def prepare_filter(self, filters: Filter):
         (All search tests in a case use consistent filtering conditions.)"""
         return
 
+    def copy_for_thread(self) -> "VectorDB":
+        """Return a per-thread copy of this client for non-thread-safe backends.
+
+        Runners call this (instead of branching on db.name) when thread_safe is
+        False, then init() the copy inside the worker so each thread owns its own
+        connection. Defaults to a deep copy; clients whose live connection can't be
+        deep-copied (e.g. an open DB-API socket) override this to shallow-copy and
+        drop their connection handles so init() re-establishes them per thread.
+        """
+        return deepcopy(self)
+
     @abstractmethod
     def __init__(
         self,

vectordb_bench/backend/clients/doris/doris.py

@@ -1,6 +1,7 @@
 import logging
 import os
 from contextlib import contextmanager
+from copy import deepcopy
 from typing import Any
 
 import pandas as pd
@@ -113,6 +114,14 @@ def _ensure_client_initialized(self):
                 # Table might not exist yet; leave it to ready_to_load
                 self.table = None
 
+    def copy_for_thread(self) -> "VectorDB":
+        # DorisVectorClient isn't thread-safe; hand each worker its own copy and force
+        # a fresh client/table so init() rebuilds them instead of sharing the parent's.
+        db_copy = deepcopy(self)
+        db_copy.client = None
+        db_copy.table = None
+        return db_copy
+
     @contextmanager
     def init(self):
         try:

vectordb_bench/backend/clients/oceanbase/oceanbase.py

@@ -3,6 +3,7 @@
 import time
 from collections.abc import Generator
 from contextlib import contextmanager
+from copy import copy
 from typing import Any
 
 import mysql.connector as mysql
@@ -82,6 +83,14 @@ def _disconnect(self):
             self._conn.close()
             self._conn = None
 
+    def copy_for_thread(self) -> "VectorDB":
+        # mysql.connector holds an open socket that can't be deep-copied; shallow-copy
+        # and drop the connection so init() reconnects inside the worker thread.
+        db_copy = copy(self)
+        db_copy._conn = None
+        db_copy._cursor = None
+        return db_copy
+
     @contextmanager
     def init(self) -> Generator[None, None, None]:
         try:

vectordb_bench/backend/clients/pgvector/pgvector.py

@@ -3,6 +3,7 @@
 import logging
 from collections.abc import Generator, Sequence
 from contextlib import contextmanager
+from copy import copy
 from typing import Any
 
 import numpy as np
@@ -99,6 +100,14 @@ def _create_connection(**kwargs) -> tuple[Connection, Cursor]:
 
         return conn, cursor
 
+    def copy_for_thread(self) -> "VectorDB":
+        # psycopg holds an open socket that can't be deep-copied; shallow-copy
+        # and drop the connection so init() reconnects inside the worker thread.
+        db_copy = copy(self)
+        db_copy.conn = None
+        db_copy.cursor = None
+        return db_copy
+
     def _generate_search_query(self) -> sql.Composed:
         index_param = self.case_config.index_param()
         reranking = self.case_config.search_param()["reranking"]

vectordb_bench/backend/clients/seekdb/seekdb.py

@@ -3,6 +3,7 @@
 import struct
 from collections.abc import Generator
 from contextlib import contextmanager
+from copy import copy
 from typing import Any
 
 import mysql.connector as mysql
@@ -125,6 +126,14 @@ def _init_session_settings(self):
             # SeekDB uses OceanBase-style session vars (not plain hnsw_ef_search).
             self._cursor.execute(f"SET ob_hnsw_ef_search={ef_search}")
 
+    def copy_for_thread(self) -> "VectorDB":
+        # mysql.connector holds an open socket that can't be deep-copied; shallow-copy
+        # and drop the connection so init() reconnects inside the worker thread.
+        db_copy = copy(self)
+        db_copy._conn = None
+        db_copy._cursor = None
+        return db_copy
+
     @contextmanager
     def init(self) -> Generator[None, None, None]:
         try:

vectordb_bench/backend/clients/vectorchord/vectorchord.py

@@ -3,6 +3,7 @@
 import logging
 from collections.abc import Generator
 from contextlib import contextmanager
+from copy import copy
 from typing import Any
 
 import numpy as np
@@ -99,6 +100,14 @@ def _create_connection(**kwargs) -> tuple[Connection, Cursor]:
 
         return conn, cursor
 
+    def copy_for_thread(self) -> "VectorDB":
+        # psycopg holds an open socket that can't be deep-copied; shallow-copy
+        # and drop the connection so init() reconnects inside the worker thread.
+        db_copy = copy(self)
+        db_copy.conn = None
+        db_copy.cursor = None
+        return db_copy
+
     @contextmanager
     def init(self) -> Generator[None, None, None]:
         self.conn, self.cursor = self._create_connection(**self.db_config)

vectordb_bench/backend/clients/volc_mysql/volc_mysql.py

@@ -5,6 +5,7 @@
 import struct
 import tempfile
 from contextlib import contextmanager
+from copy import copy
 from pathlib import Path
 
 import mysql.connector as mysql
@@ -206,6 +207,15 @@ def _probe_binary_support(self) -> bool:
                 log.debug("Failed to drop binary-probe temp table", exc_info=True)
             cur.close()
 
+    def copy_for_thread(self) -> "VectorDB":
+        # mysql.connector holds an open socket that can't be deep-copied; shallow-copy
+        # and drop the connection so init() reconnects inside the worker thread.
+        db_copy = copy(self)
+        db_copy.conn = None
+        db_copy.cursor = None
+        db_copy.admin_cursor = None
+        return db_copy
+
     @contextmanager
     def init(self):
         """create and destory connections to database.

vectordb_bench/backend/runner/rate_runner.py

@@ -3,7 +3,6 @@
 import multiprocessing as mp
 import time
 from concurrent.futures import ThreadPoolExecutor
-from copy import copy, deepcopy
 
 from vectordb_bench import config
 from vectordb_bench.backend.clients import api
@@ -60,49 +59,15 @@ def _insert_embeddings(db: api.VectorDB, emb: list[list[float]], metadata: list[
                     msg = f"Insert failed and retried more than {config.MAX_INSERT_RETRY} times"
                     raise RuntimeError(msg) from None
 
-        if db.name == "PgVector":
-            # pgvector is not thread-safe for concurrent insert,
-            #   so we need to copy the db object, make sure each thread has its own connection
-            db_copy = deepcopy(db)
-            with db_copy.init():
-                _insert_embeddings(db_copy, emb, metadata, retry_idx=0)
-        elif db.name == "Doris":
-            # DorisVectorClient is not thread-safe. Similar to pgvector, create a per-thread client
-            # by deep-copying the wrapper and forcing lazy re-init inside the thread.
-            db_copy = deepcopy(db)
-            # Ensure a fresh client/table will be created in this thread
-            try:
-                db_copy.client = None
-                db_copy.table = None
-            except Exception:
-                log.debug("Failed to reset Doris client or table on thread-local copy", exc_info=True)
-            with db_copy.init():
-                _insert_embeddings(db_copy, emb, metadata, retry_idx=0)
-        elif db.name == "SeekDB":
-            # mysql.connector is not thread-safe; do not share one connection across workers.
-            # deepcopy() fails on an open _conn (socket is not picklable / not copy-safe in spawn workers).
-            db_copy = copy(db)
-            try:
-                db_copy._conn = None
-                db_copy._cursor = None
-            except Exception:
-                log.debug("Failed to reset SeekDB connection on thread-local copy", exc_info=True)
-            with db_copy.init():
-                _insert_embeddings(db_copy, emb, metadata, retry_idx=0)
-        elif db.name == "VolcMySQL":
-            # mysql.connector is not thread-safe; do not share one connection across workers.
-            # deepcopy() fails on an open conn (socket is not picklable / not copy-safe in spawn workers).
-            db_copy = copy(db)
-            try:
-                db_copy.conn = None
-                db_copy.cursor = None
-                db_copy.admin_cursor = None
-            except Exception:
-                log.debug("Failed to reset VolcMySQL connection on thread-local copy", exc_info=True)
+        if db.thread_safe:
+            _insert_embeddings(db, emb, metadata, retry_idx=0)
+        else:
+            # Non-thread-safe clients can't share one connection across insert workers.
+            # Each client owns how to make a thread-local copy (see VectorDB.copy_for_thread);
+            # init() then re-establishes its connection inside this thread.
+            db_copy = db.copy_for_thread()
             with db_copy.init():
                 _insert_embeddings(db_copy, emb, metadata, retry_idx=0)
-        else:
-            _insert_embeddings(db, emb, metadata, retry_idx=0)
 
     @time_it
     def run_with_rate(self, q: mp.Queue):