fix(odin): Fix checkin worker picklable

cognite/extractorutils/unstable/core/checkin_worker.py

@@ -101,6 +101,28 @@ def __init__(
         self._lock = RLock()
         self._flush_lock = RLock()
 
+    def __getstate__(self) -> dict[str, object]:
+        """
+        Return the worker's state for pickling, dropping the locks.
+
+        The worker is passed to ``multiprocessing.Process`` as a target argument, which pickles it
+        when the process is started with the ``spawn`` method (the default on macOS and Windows).
+        ``threading.RLock`` cannot be pickled, and a lock inherited from the parent process would be
+        meaningless in the child anyway, so the locks are recreated fresh in ``__setstate__`` instead.
+        """
+        state = self.__dict__.copy()
+        del state["_lock"]
+        del state["_flush_lock"]
+        return state
+
+    def __setstate__(self, state: dict[str, object]) -> None:
+        """
+        Restore the worker's state after unpickling, recreating the locks dropped by ``__getstate__``.
+        """
+        self.__dict__.update(state)
+        self._lock = RLock()
+        self._flush_lock = RLock()
+
     @property
     def active_revision(self) -> ConfigRevision:
         """Get the active configuration revision."""

tests/test_unstable/test_checkin_worker.py

@@ -1,13 +1,17 @@
 import logging
+import pickle
 from collections.abc import Callable
 from datetime import datetime, timezone
-from multiprocessing import Event, Queue
+from multiprocessing import Event, Queue, get_context
 from threading import Thread
 from time import sleep
 
 import faker
 import pytest
 import requests_mock
+from cognite.client import CogniteClient
+from cognite.client.config import ClientConfig
+from cognite.client.credentials import OAuthClientCredentials
 
 from cognite.extractorutils.threading import CancellationToken
 from cognite.extractorutils.unstable.configuration.models import ConnectionConfig
@@ -18,6 +22,39 @@
 from tests.test_unstable.conftest import TestConfig, TestExtractor
 
 
+def _make_local_cognite_client() -> CogniteClient:
+    """
+    Build a CogniteClient that never talks to the network, for tests that only exercise pickling.
+
+    Avoids the live-integration `connection_config` fixture, which provisions a real integration
+    in a CDF project and therefore requires dev credentials to be set in the environment.
+    """
+    return CogniteClient(
+        ClientConfig(
+            client_name="test-checkin-worker-pickling",
+            project="test-project",
+            base_url="https://api.cognitedata.com",
+            credentials=OAuthClientCredentials(
+                token_url="https://example.com/token",
+                client_id="client-id",
+                client_secret="client-secret",
+                scopes=["scope"],
+            ),
+        )
+    )
+
+
+def _put_active_revision_in_queue(worker: CheckinWorker, result_queue: "Queue[int]") -> None:
+    """
+    Module-level target for the spawn-based multiprocessing test below.
+
+    A spawned child process re-imports this function by qualified name, so it must be a
+    top-level, module-scoped callable rather than a closure or a lambda.
+    """
+    with worker._lock, worker._flush_lock:
+        result_queue.put(worker.active_revision)
+
+
 def test_report_startup_request(
     connection_config: ConnectionConfig,
     application_config: TestConfig,
@@ -536,3 +573,47 @@ def test_run_report_periodic_checkin_requeue(
     # initial 2 requests for auth and startup, then 1 for expected number of check-ins
     assert requests_mock.call_count == 2 + 1
     assert len(worker._errors) == 2
+
+
+def test_checkin_worker_pickle_round_trip() -> None:
+    """
+    __getstate__/__setstate__ must drop the un-picklable RLocks and recreate fresh, usable ones,
+    while preserving the rest of the worker's state.
+    """
+    worker = CheckinWorker(_make_local_cognite_client(), "test-integration", logging.getLogger(__name__))
+    worker.active_revision = 5
+    worker.report_task_start("some-task")
+
+    restored: CheckinWorker = pickle.loads(pickle.dumps(worker))  # noqa: S301 (data is produced in-process, not untrusted)
+
+    assert restored.active_revision == 5
+    assert len(restored._task_updates) == 1
+    assert restored._lock is not worker._lock
+    assert restored._flush_lock is not worker._flush_lock
+
+    # The recreated locks must actually be usable, not left in some broken or shared state.
+    with restored._lock, restored._flush_lock:
+        pass
+
+
+def test_checkin_worker_is_picklable_with_spawn_start_method() -> None:
+    """
+    Runtime._spawn_extractor passes a live CheckinWorker to multiprocessing.Process as a target argument.
+
+    On macOS and Windows, the default start method is "spawn", which pickles every argument passed to
+    Process(...). This reproduces that path directly: threading.RLock instances used to make the whole
+    worker unpicklable, crashing the extractor with a TypeError before any provider/destination code ran.
+    """
+    worker = CheckinWorker(_make_local_cognite_client(), "test-integration", logging.getLogger(__name__))
+    worker.active_revision = 3
+
+    ctx = get_context("spawn")
+    result_queue: Queue[int] = ctx.Queue()
+    process = ctx.Process(target=_put_active_revision_in_queue, args=(worker, result_queue))
+    process.start()
+    try:
+        assert result_queue.get(timeout=10) == 3
+    finally:
+        process.join(timeout=10)
+
+    assert process.exitcode == 0