Stats collector using context managers to hold the data

- Replaced list append with an explicit summation.
- Provided a variable as the divider for time calculation.

Author: kkalinowski-reef

CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Infrastructure
 * Fix nox's deprecated `session.install()` calls
 * Re-enable changelog validation in CI
+* StatsCollector is now a context manager
 
 ## [1.17.3] - 2022-07-15
 

b2sdk/transfer/inbound/downloader/parallel.py

@@ -8,22 +8,22 @@
 #
 ######################################################################
 
+import logging
+import queue
+import threading
 from concurrent import futures
 from io import IOBase
 from time import perf_counter_ns
 from typing import Optional
-import logging
-import queue
-import threading
 
 from requests.models import Response
 
-from .abstract import AbstractDownloader
-from .stats_collector import StatsCollector
 from b2sdk.encryption.setting import EncryptionSetting
 from b2sdk.file_version import DownloadVersion
 from b2sdk.session import B2Session
 from b2sdk.utils.range_ import Range
+from .abstract import AbstractDownloader
+from .stats_collector import StatsCollector
 
 logger = logging.getLogger(__name__)
 
@@ -219,27 +219,25 @@ def __init__(self, file, max_queue_depth):
     def run(self):
         file = self.file
         queue_get = self.queue.get
-        stats_collector_read_append = self.stats_collector.read.append
-        stats_collector_other_append = self.stats_collector.other.append
-        stats_collector_write_append = self.stats_collector.write.append
-        start = perf_counter_ns()
-        while 1:
+        stats_collector_read = self.stats_collector.read
+        stats_collector_other = self.stats_collector.other
+        stats_collector_write = self.stats_collector.write
 
-            before_read = perf_counter_ns()
-            shutdown, offset, data = queue_get()
-            stats_collector_read_append(perf_counter_ns() - before_read)
+        with self.stats_collector.total:
+            while 1:
+                with stats_collector_read:
+                    shutdown, offset, data = queue_get()
 
-            if shutdown:
-                break
-            before_seek = perf_counter_ns()
-            file.seek(offset)
-            after_seek = perf_counter_ns()
-            file.write(data)
-            after_write = perf_counter_ns()
-            stats_collector_other_append(after_seek - before_seek)
-            stats_collector_write_append(after_write - after_seek)
-            self.total += len(data)
-        self.stats_collector.total = perf_counter_ns() - start
+                if shutdown:
+                    break
+
+                with stats_collector_other:
+                    file.seek(offset)
+
+                with stats_collector_write:
+                    file.write(data)
+
+                self.total += len(data)
 
     def __enter__(self):
         self.start()
@@ -294,67 +292,73 @@ def download_first_part(
     stop = False
 
     stats_collector = StatsCollector(response.url, f'{first_offset}:{last_offset}', 'hash')
-    stats_collector_read_append = stats_collector.read.append
-    stats_collector_other_append = stats_collector.other.append
-    stats_collector_write_append = stats_collector.write.append
-    start = before_read = perf_counter_ns()
-    for data in response.iter_content(chunk_size=chunk_size):
-        stats_collector_read_append(perf_counter_ns() - before_read)
-        if first_offset + bytes_read + len(data) >= last_offset:
-            to_write = data[:last_offset - bytes_read]
-            stop = True
-        else:
-            to_write = data
-        before_put = perf_counter_ns()
-        writer_queue_put((False, first_offset + bytes_read, to_write))
-
-        before_hash = perf_counter_ns()
-        hasher_update(to_write)
-        after_hash = perf_counter_ns()
-
-        stats_collector_write_append(before_hash - before_put)
-        stats_collector_other_append(after_hash - before_hash)
-
-        bytes_read += len(to_write)
-        if stop:
-            break
-
-    # since we got everything we need from original response, close the socket and free the buffer
-    # to avoid a timeout exception during hashing and other trouble
-    response.close()
-
-    url = response.request.url
-    tries_left = 5 - 1  # this is hardcoded because we are going to replace the entire retry interface soon, so we'll avoid deprecation here and keep it private
-    while tries_left and bytes_read < actual_part_size:
-        cloud_range = starting_cloud_range.subrange(
-            bytes_read, actual_part_size - 1
-        )  # first attempt was for the whole file, but retries are bound correctly
-        logger.debug(
-            'download attempts remaining: %i, bytes read already: %i. Getting range %s now.',
-            tries_left, bytes_read, cloud_range
-        )
-        with session.download_file_from_url(
-            url,
-            cloud_range.as_tuple(),
-            encryption=encryption,
-        ) as response:
-            before_read = perf_counter_ns()
-            for to_write in response.iter_content(chunk_size=chunk_size):
-                stats_collector_read_append(perf_counter_ns() - before_read)
+    stats_collector_read = stats_collector.read
+    stats_collector_other = stats_collector.other
+    stats_collector_write = stats_collector.write
+
+    with stats_collector.total:
+        response_iterator = response.iter_content(chunk_size=chunk_size)
+
+        while True:
+            with stats_collector_read:
+                try:
+                    data = next(response_iterator)
+                except StopIteration:
+                    break
+
+            if first_offset + bytes_read + len(data) >= last_offset:
+                to_write = data[:last_offset - bytes_read]
+                stop = True
+            else:
+                to_write = data
 
-                before_put = perf_counter_ns()
+            with stats_collector_write:
                 writer_queue_put((False, first_offset + bytes_read, to_write))
-                before_hash = perf_counter_ns()
+
+            with stats_collector_other:
                 hasher_update(to_write)
-                after_hash = perf_counter_ns()
 
-                stats_collector_write_append(before_hash - before_put)
-                stats_collector_other_append(after_hash - before_hash)
+            bytes_read += len(to_write)
+            if stop:
+                break
+
+        # since we got everything we need from original response, close the socket and free the buffer
+        # to avoid a timeout exception during hashing and other trouble
+        response.close()
+
+        url = response.request.url
+        tries_left = 5 - 1  # this is hardcoded because we are going to replace the entire retry interface soon, so we'll avoid deprecation here and keep it private
+        while tries_left and bytes_read < actual_part_size:
+            cloud_range = starting_cloud_range.subrange(
+                bytes_read, actual_part_size - 1
+            )  # first attempt was for the whole file, but retries are bound correctly
+            logger.debug(
+                'download attempts remaining: %i, bytes read already: %i. Getting range %s now.',
+                tries_left, bytes_read, cloud_range
+            )
+            with session.download_file_from_url(
+                url,
+                cloud_range.as_tuple(),
+                encryption=encryption,
+            ) as response:
+                response_iterator = response.iter_content(chunk_size=chunk_size)
+
+                while True:
+                    with stats_collector_read:
+                        try:
+                            to_write = next(response_iterator)
+                        except StopIteration:
+                            break
+
+                    with stats_collector_write:
+                        writer_queue_put((False, first_offset + bytes_read, to_write))
+
+                    with stats_collector_other:
+                        hasher_update(to_write)
+
+                    bytes_read += len(to_write)
+            tries_left -= 1
 
-                bytes_read += len(to_write)
-                before_read = perf_counter_ns()
-        tries_left -= 1
-    stats_collector.total = perf_counter_ns() - start
     stats_collector.report()
 
 
@@ -389,25 +393,30 @@ def download_non_first_part(
             retries_left, bytes_read, cloud_range
         )
         stats_collector = StatsCollector(url, f'{cloud_range.start}:{cloud_range.end}', 'none')
-        stats_collector_read_append = stats_collector.read.append
-        stats_collector_write_append = stats_collector.write.append
-        start = before_read = perf_counter_ns()
-        with session.download_file_from_url(
-            url,
-            cloud_range.as_tuple(),
-            encryption=encryption,
-        ) as response:
-            before_read = perf_counter_ns()
-            for to_write in response.iter_content(chunk_size=chunk_size):
-                after_read = perf_counter_ns()
-                writer_queue_put((False, start_range + bytes_read, to_write))
-                after_write = perf_counter_ns()
-                stats_collector_read_append(after_read - before_read)
-                stats_collector_write_append(after_write - after_read)
-                bytes_read += len(to_write)
-                before_read = perf_counter_ns()
-        retries_left -= 1
-        stats_collector.total = perf_counter_ns() - start
+        stats_collector_read = stats_collector.read
+        stats_collector_write = stats_collector.write
+
+        with stats_collector.total:
+            with session.download_file_from_url(
+                url,
+                cloud_range.as_tuple(),
+                encryption=encryption,
+            ) as response:
+                response_iterator = response.iter_content(chunk_size=chunk_size)
+
+                while True:
+                    with stats_collector_read:
+                        try:
+                            to_write = next(response_iterator)
+                        except StopIteration:
+                            break
+
+                    with stats_collector_write:
+                        writer_queue_put((False, start_range + bytes_read, to_write))
+
+                    bytes_read += len(to_write)
+            retries_left -= 1
+
         stats_collector.report()
 
 

b2sdk/transfer/inbound/downloader/stats_collector.py

@@ -9,43 +9,81 @@
 ######################################################################
 
 import logging
-from dataclasses import dataclass, field
-from typing import List  # 3.7 doesn't understand `list` vs `List`
-from typing import Optional
+from dataclasses import (
+    dataclass,
+    field,
+)
+from time import perf_counter_ns
+from typing import (
+    Any,
+    Optional,
+    Type,
+)
 
 logger = logging.getLogger(__name__)
 
 
+class SingleStatsCollector:
+    TO_MS = 1_000_000
+
+    def __init__(self):
+        self.latest_entry: Optional[int] = None
+        self.sum_of_all_entries: int = 0
+        self.started_perf_timer: Optional[int] = None
+
+    def __enter__(self) -> None:
+        self.started_perf_timer = perf_counter_ns()
+
+    def __exit__(self, exc_type: Type, exc_val: Exception, exc_tb: Any) -> None:
+        time_diff = perf_counter_ns() - self.started_perf_timer
+        self.latest_entry = time_diff
+        self.sum_of_all_entries += time_diff
+        self.started_perf_timer = None
+
+    @property
+    def sum_ms(self) -> float:
+        return self.sum_of_all_entries / self.TO_MS
+
+    @property
+    def latest_ms(self) -> float:
+        return self.latest_entry / self.TO_MS
+
+    @property
+    def has_any_entry(self) -> bool:
+        return self.latest_entry is not None
+
+
 @dataclass
 class StatsCollector:
     name: str  #: file name or object url
     detail: str  #: description of the thread, ex. "10000000:20000000" or "writer"
     other_name: str  #: other statistic, typically "seek" or "hash"
-    total: Optional[int] = None
-    other: List[int] = field(default_factory=list)
-    write: List[int] = field(default_factory=list)
-    read: List[int] = field(default_factory=list)
+    total: SingleStatsCollector = field(default_factory=SingleStatsCollector)
+    other: SingleStatsCollector = field(default_factory=SingleStatsCollector)
+    write: SingleStatsCollector = field(default_factory=SingleStatsCollector)
+    read: SingleStatsCollector = field(default_factory=SingleStatsCollector)
 
     def report(self):
-        if self.read:
-            logger.info('download stats | %s | TTFB: %.3f ms', self, self.read[0] / 1000000)
+        if self.read.has_any_entry:
+            logger.info('download stats | %s | TTFB: %.3f ms', self, self.read.latest_ms)
             logger.info(
                 'download stats | %s | read() without TTFB: %.3f ms', self,
-                sum(self.read[1:]) / 1000000
+                (self.read.sum_of_all_entries - self.read.latest_entry) / self.read.TO_MS
             )
-        if self.other:
+        if self.other.has_any_entry:
             logger.info(
-                'download stats | %s | %s total: %.3f ms', self, self.other_name,
-                sum(self.other) / 1000000
+                'download stats | %s | %s total: %.3f ms', self, self.other_name, self.other.sum_ms
             )
-        if self.write:
+        if self.write.has_any_entry:
+            logger.info('download stats | %s | write() total: %.3f ms', self, self.write.sum_ms)
+        if self.total.has_any_entry:
+            basic_operation_time = self.write.sum_of_all_entries \
+                                   + self.other.sum_of_all_entries \
+                                   + self.read.sum_of_all_entries
+            overhead = self.total.sum_of_all_entries - basic_operation_time
             logger.info(
-                'download stats | %s | write() total: %.3f ms', self,
-                sum(self.write) / 1000000
+                'download stats | %s | overhead: %.3f ms', self, overhead / self.total.TO_MS
             )
-        if self.total is not None:
-            overhead = self.total - sum(self.write) - sum(self.other) - sum(self.read)
-            logger.info('download stats | %s | overhead: %.3f ms', self, overhead / 1000000)
 
     def __str__(self):
         return f'{self.name}[{self.detail}]'