diff --git a/cognite/extractorutils/unstable/core/checkin_worker.py b/cognite/extractorutils/unstable/core/checkin_worker.py index db4d28f1..c9c295d3 100644 --- a/cognite/extractorutils/unstable/core/checkin_worker.py +++ b/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.""" diff --git a/tests/test_unstable/test_checkin_worker.py b/tests/test_unstable/test_checkin_worker.py index 56bb59fa..63cc4110 100644 --- a/tests/test_unstable/test_checkin_worker.py +++ b/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,52 @@ 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) + if process.is_alive(): + # Only reached if the process failed to exit on its own (e.g. the assert above raised + # before the process could finish); avoids leaking a background process in that case. + process.terminate() + process.join(timeout=5) + + assert process.exitcode == 0