Add a wait_for_samples method to the MovingWindow (frequenz-floss#1159)

Closes frequenz-floss#967

Author: shsms

RELEASE_NOTES.md

@@ -10,7 +10,7 @@
 
 ## New Features
 
-<!-- Here goes the main new features and examples or instructions on how to use them -->
+- The `MovingWindow` now has an async `wait_for_samples` method that waits for a given number of samples to become available in the moving window and then returns.
 
 ## Bug Fixes
 

src/frequenz/sdk/timeseries/_moving_window.py

@@ -190,6 +190,8 @@ def __init__(  # pylint: disable=too-many-arguments
             align_to=align_to,
         )
 
+        self._condition_new_sample = asyncio.Condition()
+
     def start(self) -> None:
         """Start the MovingWindow.
 
@@ -318,6 +320,44 @@ def window(
             start, end, force_copy=force_copy, fill_value=fill_value
         )
 
+    async def wait_for_samples(self, n: int) -> None:
+        """Wait until the next `n` samples are available in the MovingWindow.
+
+        This function returns after `n` new samples are available in the MovingWindow,
+        without considering whether the new samples are valid.  The validity of the
+        samples can be verified by calling the
+        [`count_valid`][frequenz.sdk.timeseries.MovingWindow.count_valid] method.
+
+        Args:
+            n: The number of samples to wait for.
+
+        Raises:
+            ValueError: If `n` is less than or equal to 0 or greater than the capacity
+                of the MovingWindow.
+        """
+        if n == 0:
+            return
+        if n < 0:
+            raise ValueError("The number of samples to wait for must be 0 or greater.")
+        if n > self.capacity:
+            raise ValueError(
+                "The number of samples to wait for must be less than or equal to the "
+                + f"capacity of the MovingWindow ({self.capacity})."
+            )
+        start_timestamp = (
+            # Start from the next expected timestamp.
+            self.newest_timestamp + self.sampling_period
+            if self.newest_timestamp is not None
+            else None
+        )
+        while True:
+            async with self._condition_new_sample:
+                # Every time a new sample is received, this condition gets notified and
+                # will wake up.
+                _ = await self._condition_new_sample.wait()
+            if self.count_covered(since=start_timestamp) >= n:
+                return
+
     async def _run_impl(self) -> None:
         """Awaits samples from the receiver and updates the underlying ring buffer.
 
@@ -331,6 +371,9 @@ async def _run_impl(self) -> None:
                     await self._resampler_sender.send(sample)
                 else:
                     self._buffer.update(sample)
+                    async with self._condition_new_sample:
+                        # Wake up all coroutines waiting for new samples.
+                        self._condition_new_sample.notify_all()
 
         except asyncio.CancelledError:
             _logger.info("MovingWindow task has been cancelled.")
@@ -343,8 +386,10 @@ def _configure_resampler(self) -> None:
         assert self._resampler is not None
 
         async def sink_buffer(sample: Sample[Quantity]) -> None:
-            if sample.value is not None:
-                self._buffer.update(sample)
+            self._buffer.update(sample)
+            async with self._condition_new_sample:
+                # Wake up all coroutines waiting for new samples.
+                self._condition_new_sample.notify_all()
 
         resampler_channel = Broadcast[Sample[Quantity]](name="average")
         self._resampler_sender = resampler_channel.new_sender()
@@ -355,23 +400,44 @@ async def sink_buffer(sample: Sample[Quantity]) -> None:
             asyncio.create_task(self._resampler.resample(), name="resample")
         )
 
-    def count_valid(self) -> int:
-        """
-        Count the number of valid samples in this `MovingWindow`.
+    def count_valid(
+        self, *, since: datetime | None = None, until: datetime | None = None
+    ) -> int:
+        """Count the number of valid samples in this `MovingWindow`.
+
+        If `since` and `until` are provided, the count is limited to the samples between
+        (and including) the given timestamps.
+
+        Args:
+            since: The timestamp from which to start counting.  If `None`, the oldest
+                timestamp of the buffer is used.
+            until: The timestamp until (and including) which to count.  If `None`, the
+                newest timestamp of the buffer is used.
 
         Returns:
             The number of valid samples in this `MovingWindow`.
         """
-        return self._buffer.count_valid()
+        return self._buffer.count_valid(since=since, until=until)
 
-    def count_covered(self) -> int:
+    def count_covered(
+        self, *, since: datetime | None = None, until: datetime | None = None
+    ) -> int:
         """Count the number of samples that are covered by the oldest and newest valid samples.
 
+        If `since` and `until` are provided, the count is limited to the samples between
+        (and including) the given timestamps.
+
+        Args:
+            since: The timestamp from which to start counting.  If `None`, the oldest
+                timestamp of the buffer is used.
+            until: The timestamp until (and including) which to count.  If `None`, the
+                newest timestamp of the buffer is used.
+
         Returns:
             The count of samples between the oldest and newest (inclusive) valid samples
                 or 0 if there are is no time range covered.
         """
-        return self._buffer.count_covered()
+        return self._buffer.count_covered(since=since, until=until)
 
     @overload
     def __getitem__(self, key: SupportsIndex) -> float:

src/frequenz/sdk/timeseries/_ringbuffer/buffer.py

@@ -651,9 +651,20 @@ def __getitem__(self, index_or_slice: SupportsIndex | slice) -> float | FloatArr
         """
         return self._buffer.__getitem__(index_or_slice)
 
-    def _covered_time_range(self) -> timedelta:
+    def _covered_time_range(
+        self, since: datetime | None = None, until: datetime | None = None
+    ) -> timedelta:
         """Return the time range that is covered by the oldest and newest valid samples.
 
+        If `since` and `until` are provided, the time range is limited to the items
+        between (and including) the given timestamps.
+
+        Args:
+            since: The timestamp from which to start counting.  If `None`, the oldest
+                timestamp in the buffer is used.
+            until: The timestamp until (and including) which to count.  If `None`, the
+                newest timestamp in the buffer is used.
+
         Returns:
             The time range between the oldest and newest valid samples or 0 if
                 there are is no time range covered.
@@ -664,45 +675,82 @@ def _covered_time_range(self) -> timedelta:
         assert (
             self.newest_timestamp is not None
         ), "Newest timestamp cannot be None here."
-        return self.newest_timestamp - self.oldest_timestamp + self._sampling_period
 
-    def count_covered(self) -> int:
+        if since is None or since < self.oldest_timestamp:
+            since = self.oldest_timestamp
+        if until is None or until > self.newest_timestamp:
+            until = self.newest_timestamp
+
+        if until < since:
+            return timedelta(0)
+
+        return until - since + self._sampling_period
+
+    def count_covered(
+        self, *, since: datetime | None = None, until: datetime | None = None
+    ) -> int:
         """Count the number of samples that are covered by the oldest and newest valid samples.
 
+        If `since` and `until` are provided, the count is limited to the items between
+        (and including) the given timestamps.
+
+        Args:
+            since: The timestamp from which to start counting.  If `None`, the oldest
+                timestamp in the buffer is used.
+            until: The timestamp until (and including) which to count.  If `None`, the
+                newest timestamp in the buffer is used.
+
         Returns:
             The count of samples between the oldest and newest (inclusive) valid samples
                 or 0 if there are is no time range covered.
         """
         return int(
-            self._covered_time_range().total_seconds()
+            self._covered_time_range(since, until).total_seconds()
             // self._sampling_period.total_seconds()
         )
 
-    def count_valid(self) -> int:
-        """Count the number of valid items that this buffer currently holds.
+    def count_valid(
+        self, *, since: datetime | None = None, until: datetime | None = None
+    ) -> int:
+        """Count the number of valid items in this buffer.
+
+        If `since` and `until` are provided, the count is limited to the items between
+        (and including) the given timestamps.
+
+        Args:
+            since: The timestamp from which to start counting.  If `None`, the oldest
+                timestamp in the buffer is used.
+            until: The timestamp until (and including) which to count.  If `None`, the
+                newest timestamp in the buffer is used.
 
         Returns:
             The number of valid items in this buffer.
         """
-        if self._timestamp_newest == self._TIMESTAMP_MIN:
+        if since is None or since < self._timestamp_oldest:
+            since = self._timestamp_oldest
+        if until is None or until > self._timestamp_newest:
+            until = self._timestamp_newest
+
+        if until == self._TIMESTAMP_MIN or until < since:
             return 0
 
         # Sum of all elements in the gap ranges
         sum_missing_entries = max(
             0,
             sum(
                 (
-                    gap.end
+                    min(gap.end, until + self._sampling_period)
                     # Don't look further back than oldest timestamp
-                    - max(gap.start, self._timestamp_oldest)
+                    - max(gap.start, since)
                 )
                 // self._sampling_period
                 for gap in self._gaps
+                if gap.start <= until and gap.end >= since
             ),
         )
 
-        start_pos = self.to_internal_index(self._timestamp_oldest)
-        end_pos = self.to_internal_index(self._timestamp_newest)
+        start_pos = self.to_internal_index(since)
+        end_pos = self.to_internal_index(until)
 
         if end_pos < start_pos:
             return len(self._buffer) - start_pos + end_pos + 1 - sum_missing_entries