From f9a079d6f6c8391307d1226c5e92b4209f45bcc6 Mon Sep 17 00:00:00 2001
From: Einar <einar.omang@cognite.com>
Date: Tue, 20 Feb 2024 11:08:26 +0100
Subject: [PATCH] Make `CancellationToken.wait` only wait for a limited time
 (#300)

`Event` didn't have this problem, since it assumed any notify implied
that the event had been set (frankly a slightly insane assumption, but
whatever).
---
 CHANGELOG.md                        |  6 ++++++
 cognite/extractorutils/threading.py | 17 ++++++++++++++---
 pyproject.toml                      |  2 +-
 3 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9c77eea4..66a96ee7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,12 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## 7.0.1
+
+### Fixed
+
+ * Ensure that `CancellationToken.wait(timeout)` only waits for at most `timeout`, even if it is notified in that time.
+
 ## 7.0.0
 
 ### Changed
diff --git a/cognite/extractorutils/threading.py b/cognite/extractorutils/threading.py
index 9ac354f7..30c92fbf 100644
--- a/cognite/extractorutils/threading.py
+++ b/cognite/extractorutils/threading.py
@@ -1,6 +1,7 @@
 import logging
 import signal
 from threading import Condition
+from time import time
 from typing import Any, Optional
 
 
@@ -59,11 +60,21 @@ def set(self) -> None:
         self.cancel()
 
     def wait(self, timeout: Optional[float] = None) -> bool:
+        endtime = None
+        if timeout is not None:
+            endtime = time() + timeout
+
         while not self.is_cancelled:
             with self._cv:
-                did_not_time_out = self._cv.wait(timeout)
-                if not did_not_time_out:
-                    return False
+                if endtime is not None:
+                    remaining_time = endtime - time()
+                    if remaining_time <= 0.0:
+                        return True
+                    timed_out = not self._cv.wait(remaining_time)
+                    if timed_out:
+                        return False
+                else:
+                    self._cv.wait()
         return True
 
     def create_child_token(self) -> "CancellationToken":
diff --git a/pyproject.toml b/pyproject.toml
index ccb1dde2..4fc35c6a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "7.0.0"
+version = "7.0.1"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"