Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ Changes are grouped as follows
## Unreleased

### Fixed
* In the `unstable` package: Action status updates (`running`, final statuses) are now reported to Integrations API shortly after they are queued (about one second, debounced) instead of waiting for the next periodic check-in.
* In the `unstable` package: Fixed the check-in thread dying silently when the startup request could not be built or reported (for example when more than 100 actions are available). The failure is now logged at critical level and treated as fatal, and the extractor process exits with a non-zero exit code when it stops because of a fatal check-in error.
* In the `unstable` package: Fixed transient network errors during check-ins (timeouts, connection errors, including raw `requests`/`httpx` errors from the OAuth token refresh) being treated as fatal and shutting the extractor down. They are now logged as warnings and retried at the next check-in. Only 401/403 errors and bad credentials remain fatal.

Expand Down
28 changes: 25 additions & 3 deletions cognite/extractorutils/unstable/core/checkin_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@
DEFAULT_SLEEP_INTERVAL = STARTUP_BACKOFF_SECONDS = 30.0
MAX_ERRORS_PER_CHECKIN = MAX_TASK_UPDATES_PER_CHECKIN = 1000
MAX_ACTION_UPDATES_PER_CHECKIN = 100
# After an action update is queued, wait this long before the early check-in so a burst of updates is sent together.
ACTION_UPDATE_DEBOUNCE_SECONDS = 1.0
rng = SystemRandom()

# Errors that mean "could not reach CDF right now" rather than "the request or credentials are wrong". The Cognite SDK
Expand Down Expand Up @@ -115,6 +117,8 @@ def __init__(
self._action_dispatcher: Callable[[list[Action]], None] | None = None
self._lock = RLock()
self._flush_lock = RLock()
# Child of the check-in loop's token; cancelling it wakes the loop early. Only set while the loop runs.
self._wake_token: CancellationToken | None = None

def __getstate__(self) -> dict[str, object]:
"""
Expand All @@ -128,6 +132,7 @@ def __getstate__(self) -> dict[str, object]:
state = self.__dict__.copy()
del state["_lock"]
del state["_flush_lock"]
del state["_wake_token"]
return state

def __setstate__(self, state: dict[str, object]) -> None:
Expand All @@ -137,6 +142,7 @@ def __setstate__(self, state: dict[str, object]) -> None:
self.__dict__.update(state)
self._lock = RLock()
self._flush_lock = RLock()
self._wake_token = None

@property
def active_revision(self) -> ConfigRevision:
Expand Down Expand Up @@ -167,15 +173,20 @@ def set_action_dispatcher(self, dispatcher: Callable[[list[Action]], None]) -> N

def queue_action_update(self, update: ActionUpdate) -> None:
"""
Queue an action update to be included in the next checkin request.
Queue an action update and wake the periodic check-in loop so it is sent promptly.

This method is thread-safe and may be called from action handler threads.
The early check-in happens after ``ACTION_UPDATE_DEBOUNCE_SECONDS``, so a burst of updates results in a
single extra check-in. This method is thread-safe and may be called from action handler threads.

Arguments:
update (ActionUpdate): The action update to queue.
"""
with self._lock:
self._action_updates.append(update)
wake_token = self._wake_token
# Cancel outside of self._lock: it only takes the token's own condition, never the worker locks.
if wake_token is not None:
wake_token.cancel()

def set_on_revision_change_handler(self, on_revision_change: Callable[[int], None]) -> None:
"""
Expand Down Expand Up @@ -244,10 +255,21 @@ def run_periodic_checkin(
report_interval = interval or DEFAULT_SLEEP_INTERVAL

while not cancellation_token.is_cancelled:
# Created before the flush so updates queued while flushing wake the next wait. Updates queued during the
# debounce below hit the previous (already cancelled) token and are picked up by this flush.
wake_token = cancellation_token.create_child_token()
with self._lock:
self._wake_token = wake_token
self._logger.debug("Running periodic check-in with interval %.2f seconds", report_interval)
self.flush(cancellation_token)
self._logger.debug(f"Check-in worker finished check-in, sleeping for {report_interval:.2f} seconds")
cancellation_token.wait(report_interval)
wake_token.wait(report_interval)
if wake_token.is_cancelled and not cancellation_token.is_cancelled:
self._logger.debug("Action update queued, checking in early")
cancellation_token.wait(ACTION_UPDATE_DEBOUNCE_SECONDS)

with self._lock:
self._wake_token = None

def _run_startup_report(
self, cancellation_token: CancellationToken, startup_request: StartupRequest, interval: float | None = None
Expand Down
104 changes: 104 additions & 0 deletions tests/test_unstable/test_early_checkin.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
"""Tests that queuing an action update wakes the periodic check-in loop early (debounced)."""

import logging
from collections.abc import Callable
from threading import Event, Thread
from time import monotonic, sleep
from unittest.mock import MagicMock

import pytest

from cognite.extractorutils.threading import CancellationToken
from cognite.extractorutils.unstable.core import checkin_worker as cw
from cognite.extractorutils.unstable.core._dto import ActionStatus, ActionUpdate, StartupRequest
from cognite.extractorutils.unstable.core.checkin_worker import CheckinWorker

DEBOUNCE = 0.2
LONG_INTERVAL = 60.0


class _Harness:
def __init__(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(cw, "ACTION_UPDATE_DEBOUNCE_SECONDS", DEBOUNCE)
self.worker = CheckinWorker(MagicMock(), "integration", logging.getLogger("test.early"))
self.flush_times: list[float] = []
self.sent: list[ActionUpdate] = []
self.token = CancellationToken()
self.first_flush = Event()

def fake_flush(token: CancellationToken) -> None:
self.flush_times.append(monotonic())
with self.worker._lock:
self.sent.extend(self.worker._action_updates)
self.worker._action_updates = []
self.first_flush.set()

monkeypatch.setattr(self.worker, "flush", fake_flush)
monkeypatch.setattr(self.worker, "_run_startup_report", lambda *a, **k: None)
self.thread = Thread(
target=self.worker.run_periodic_checkin,
args=(self.token, MagicMock(spec=StartupRequest), LONG_INTERVAL),
daemon=True,
)

def __enter__(self) -> "_Harness":
self.thread.start()
assert self.first_flush.wait(5)
return self

def __exit__(self, *_: object) -> None:
self.token.cancel()
self.thread.join(5)
assert not self.thread.is_alive(), "loop must honor cancellation"


def _update(i: int = 0) -> ActionUpdate:
return ActionUpdate(external_id=f"a{i}", status=ActionStatus.running)


def _wait_for(predicate: Callable[[], bool], timeout: float = 5.0) -> bool:
end = monotonic() + timeout
while monotonic() < end:
if predicate():
return True
sleep(0.01)
return False


def test_update_triggers_early_checkin_within_debounce_window(monkeypatch: pytest.MonkeyPatch) -> None:
with _Harness(monkeypatch) as h:
queued_at = monotonic()
h.worker.queue_action_update(_update())
assert _wait_for(lambda: len(h.flush_times) == 2), "update waited for the periodic interval"
assert h.flush_times[1] - queued_at < DEBOUNCE + 1.0
assert h.flush_times[1] - queued_at >= DEBOUNCE * 0.9
assert len(h.sent) == 1


def test_burst_of_updates_produces_one_extra_checkin(monkeypatch: pytest.MonkeyPatch) -> None:
with _Harness(monkeypatch) as h:
for i in range(50):
h.worker.queue_action_update(_update(i))
assert _wait_for(lambda: len(h.flush_times) == 2)
sleep(DEBOUNCE * 3)
assert len(h.flush_times) == 2
assert len(h.sent) == 50


def test_no_update_keeps_periodic_cadence(monkeypatch: pytest.MonkeyPatch) -> None:
with _Harness(monkeypatch) as h:
sleep(DEBOUNCE * 3)
assert len(h.flush_times) == 1


def test_cancellation_during_wait_stops_loop_without_extra_flush(monkeypatch: pytest.MonkeyPatch) -> None:
h = _Harness(monkeypatch)
with h:
pass
assert len(h.flush_times) == 1


def test_queue_update_before_loop_starts_does_not_fail() -> None:
worker = CheckinWorker(MagicMock(), "integration", logging.getLogger("test.early"))
worker.queue_action_update(_update())
assert len(worker._action_updates) == 1
Loading