fix(datapoints): restore backpressure in insert_multiple to fix memory regression vs v7 (#2594)

Author: haakonvt

cognite/client/_api/datapoints.py

@@ -2419,17 +2419,18 @@ def _create_payload_tasks(
             yield payload
 
     async def _insert_datapoints(self, payload: list[dict[str, Any]]) -> None:
-        # Convert to memory intensive format as late as possible (and clean up after insert)
-        for dct in payload:
-            dct["datapoints"] = [dp.dump() for dp in dct["datapoints"]]
-        headers: dict[str, str] | None = None
-
-        await self.dps_client._post(
-            url_path=self.dps_client._RESOURCE_PATH,
-            json={"items": payload},
-            headers=headers,
-            semaphore=self.dps_client._get_semaphore("write"),
-        )
+        # Acquire the semaphore before converting to memory-intensive format:
+        async with self.dps_client._get_semaphore("write"):
+            for dct in payload:
+                dct["datapoints"] = [dp.dump() for dp in dct["datapoints"]]
+            await self.dps_client._post(
+                url_path=self.dps_client._RESOURCE_PATH,
+                json={"items": payload},
+                headers=None,
+                semaphore=None,  # Already holding the semaphore above
+            )
+        # ...and clean up after insert
+        # (needed because a ref preventing gc is held to these until the whole job is done):
         for dct in payload:
             dct["datapoints"].clear()
 

cognite/client/_sync_api/datapoints.py

@@ -1,6 +1,6 @@
 """
 ===============================================================================
-dfa197641edb7840e903b3b65c021f58
+830266fb5fd8d205e7c1163175ba904e
 This file is auto-generated from the Async API modules, - do not edit manually!
 ===============================================================================
 """

tests/tests_unit/test_api/test_datapoints.py

@@ -353,6 +353,46 @@ def test_insert_multiple_ts_single_call__above_dps_limit_below_ts_limit(
         cognite_client.time_series.data.insert_multiple(dps_objects)
         assert 2 == len(mock_post_datapoints.get_requests())
 
+    def test_insert_multiple_dump_only_called_within_semaphore(
+        self,
+        cognite_client: CogniteClient,
+        mock_post_datapoints: HTTPXMock,
+        monkeypatch: MonkeyPatch,
+        async_client: AsyncCogniteClient,
+    ) -> None:
+        # Bug in 8.0.0 to 8.1.0: all asyncio tasks were created at once and dp.dump() (synchronous,
+        # no await) ran in every task before any of them reached the semaphore, causing all payloads
+        # to be materialised in memory simultaneously (memory regression vs v7).
+        sem_held = False
+        dump_sem_held: list[bool] = []
+
+        class _MockSemaphore:
+            async def __aenter__(self) -> _MockSemaphore:
+                nonlocal sem_held
+                sem_held = True
+                return self
+
+            async def __aexit__(self, *_: Any) -> None:
+                nonlocal sem_held
+                sem_held = False
+
+        monkeypatch.setattr(async_client.time_series.data, "_get_semaphore", lambda _op: _MockSemaphore())
+
+        original_dump = _InsertDatapoint.dump
+
+        def tracking_dump(self: _InsertDatapoint) -> dict:
+            dump_sem_held.append(sem_held)
+            return original_dump(self)
+
+        monkeypatch.setattr(_InsertDatapoint, "dump", tracking_dump)
+
+        monkeypatch.setattr(async_client.time_series.data, "_DPS_INSERT_LIMIT", 5)
+        dps = [(i * 1e11, i) for i in range(50)]
+        cognite_client.time_series.data.insert(dps, id=1)
+
+        assert len(dump_sem_held) == 50
+        assert all(dump_sem_held), "dp.dump() was called outside the semaphore context"
+
 
 @pytest.fixture
 def mock_delete_datapoints(