Reduz o pico de memória na conversão e exportação de métricas diárias

VERSION

@@ -1 +1 @@
-2.3.3
+2.3.4

config/settings/base.py

@@ -433,6 +433,14 @@
     "OPENSEARCH_VERIFY_CERTS",
     default=False,
 )
+OPENSEARCH_HTTP_COMPRESS = env.bool(
+    "OPENSEARCH_HTTP_COMPRESS",
+    default=True,
+)
+OPENSEARCH_BULK_CHUNK_SIZE = env.int(
+    "OPENSEARCH_BULK_CHUNK_SIZE",
+    default=500,
+)
 
 # Resources
 # ------------------------------------------------------------------------------

metrics/counter/access/daily_accumulator.py

@@ -167,6 +167,18 @@ def iter_materialized_values(self, consume=False):
         finally:
             self.clear()
 
+    def iter_materialized_record_items(self):
+        for record_key, record in self._records.items():
+            yield record_key, record.as_dict(self)
+
+    def iter_materialized_record_keys(self, record_keys, consume=False):
+        for record_key in record_keys:
+            if consume:
+                record = self._records.pop(record_key)
+            else:
+                record = self._records[record_key]
+            yield record.as_dict(self)
+
     def clear(self):
         self._records.clear()
         self._documents.clear()

metrics/counter/indexing/converter.py

@@ -1,3 +1,5 @@
+import hashlib
+
 from metrics.counter.indexing.engines.article import ArticlePipeline
 from metrics.counter.indexing.engines.base import DocumentPipeline
 from metrics.counter.indexing.engines.book import BookPipeline
@@ -12,22 +14,92 @@
     "chapter": BookPipeline(),
 }
 _DEFAULT = DocumentPipeline()
+_DEFAULT_PARTITION_COUNT = 64
+
+
+def iter_partitioned_documents(
+    accumulator,
+    granularity,
+    partition_count=_DEFAULT_PARTITION_COUNT,
+):
+    if partition_count <= 0:
+        raise ValueError("Partition count must be greater than zero.")
+
+    consume = granularity == "year"
+    partitions = [[] for _ in range(partition_count)]
+    for record_key, value in accumulator.iter_materialized_record_items():
+        partition = _partition_for_value(value, granularity, partition_count)
+        partitions[partition].append(record_key)
+
+    try:
+        for record_keys in partitions:
+            values = accumulator.iter_materialized_record_keys(
+                record_keys,
+                consume=consume,
+            )
+            yield from _convert_partition(values, granularity)
+            record_keys.clear()
+    finally:
+        for record_keys in partitions:
+            record_keys.clear()
+        partitions.clear()
+        if consume:
+            accumulator.clear()
+
+
+def iter_partitioned_values(
+    values,
+    granularity,
+    partition_count=_DEFAULT_PARTITION_COUNT,
+):
+    if partition_count <= 0:
+        raise ValueError("Partition count must be greater than zero.")
 
+    partitions = [[] for _ in range(partition_count)]
+    for value in values:
+        partition = _partition_for_value(value, granularity, partition_count)
+        partitions[partition].append(value)
+
+    try:
+        for partition_values in partitions:
+            yield from _convert_partition(partition_values, granularity)
+            partition_values.clear()
+    finally:
+        for partition_values in partitions:
+            partition_values.clear()
+        partitions.clear()
+
+
+def _partition_for_value(value, granularity, partition_count):
+    pipeline = _get_pipeline(value)
+    partition_key = pipeline.partition_key(value, granularity)
+    digest = hashlib.blake2b(
+        partition_key.encode("utf-8"),
+        digest_size=8,
+    ).digest()
+    return int.from_bytes(digest, "big") % partition_count
 
-def convert_granularity(values, granularity):
+
+def _convert_partition(values, granularity):
     converted_data = {}
     unique_state = _initialize_unique_state()
 
-    for value in values:
-        pipeline = _get_pipeline(value)
-        pipeline.accumulate(
-            data=converted_data,
-            unique_state=unique_state,
-            value=value,
-            granularity=granularity,
-        )
-
-    return converted_data
+    try:
+        for value in values:
+            pipeline = _get_pipeline(value)
+            pipeline.accumulate(
+                data=converted_data,
+                unique_state=unique_state,
+                value=value,
+                granularity=granularity,
+            )
+
+        for document_id in sorted(converted_data):
+            yield document_id, converted_data[document_id]
+    finally:
+        converted_data.clear()
+        for bucket in unique_state.values():
+            bucket.clear()
 
 
 def _get_pipeline(value):

metrics/counter/indexing/engines/base.py

@@ -29,6 +29,9 @@ def accumulate(self, data, unique_state, value, granularity):
             is_request_event=is_request(value.get("content_type")),
         )
 
+    def partition_key(self, value, granularity):
+        return self._generate_document_id(value, granularity)
+
     def _generate_document_id(
         self, value, granularity, metric_scope=None, pid_generic=None
     ):

metrics/counter/indexing/engines/book.py

@@ -8,6 +8,17 @@
 
 
 class BookPipeline(DocumentPipeline):
+    def partition_key(self, value, granularity):
+        title_pid_generic = _extract_title_pid_generic(value)
+        if title_pid_generic:
+            return self._generate_document_id(
+                value,
+                granularity,
+                metric_scope="title",
+                pid_generic=title_pid_generic,
+            )
+        return self._generate_document_id(value, granularity)
+
     def accumulate(self, data, unique_state, value, granularity):
         if not isinstance(value, dict):
             return

metrics/opensearch/client.py

@@ -10,13 +10,25 @@
     merge_metric_document,
 )
 
-_BULK_CHUNK_SIZE = 500
-
 
 class OpenSearchUsageClient:
     def __init__(self, url=None, basic_auth=None, api_key=None, verify_certs=None):
+        self.bulk_chunk_size = getattr(
+            settings,
+            "OPENSEARCH_BULK_CHUNK_SIZE",
+            500,
+        )
+        if self.bulk_chunk_size <= 0:
+            raise ValueError("OpenSearch bulk chunk size must be greater than zero.")
+
         self.client = self.get_opensearch_client(url, basic_auth, api_key, verify_certs)
-        logging.info("OpenSearch HTTP request compression is enabled.")
+        logging.info(
+            "OpenSearch HTTP request compression is %s; bulk chunk size is %s.",
+            "enabled"
+            if getattr(settings, "OPENSEARCH_HTTP_COMPRESS", True)
+            else "disabled",
+            self.bulk_chunk_size,
+        )
 
     def get_opensearch_client(
         self,
@@ -30,25 +42,26 @@ def get_opensearch_client(
         api_key = api_key or getattr(settings, "OPENSEARCH_API_KEY", None)
         if verify_certs is None:
             verify_certs = getattr(settings, "OPENSEARCH_VERIFY_CERTS", False)
+        http_compress = getattr(settings, "OPENSEARCH_HTTP_COMPRESS", True)
 
         if basic_auth:
             return OpenSearch(
                 url,
                 http_auth=tuple(basic_auth),
                 verify_certs=verify_certs,
-                http_compress=True,
+                http_compress=http_compress,
             )
         if api_key:
             return OpenSearch(
                 url,
                 api_key=api_key,
                 verify_certs=verify_certs,
-                http_compress=True,
+                http_compress=http_compress,
             )
         return OpenSearch(
             url,
             verify_certs=verify_certs,
-            http_compress=True,
+            http_compress=http_compress,
         )
 
     def ping(self):
@@ -138,7 +151,7 @@ def increment_document_items_for_daily_job(
                 )
                 for doc_id, document in document_items
             ),
-            chunk_size=_BULK_CHUNK_SIZE,
+            chunk_size=self.bulk_chunk_size,
         )
         return succeeded
 

metrics/services/daily_payloads.py

@@ -57,18 +57,28 @@ def __enter__(self):
         self._write_text(',"documents":{"month":')
         return self
 
-    def write_documents(self, granularity, documents):
+    def write_document_items(self, granularity, document_items):
         if granularity != self.next_granularity:
             raise RuntimeError(
                 f"Expected {self.next_granularity} documents, got {granularity}."
             )
 
-        self._write_json(documents)
+        document_count = 0
+        self._write_text("{")
+        for document_id, document in document_items:
+            if document_count:
+                self._write_text(",")
+            self._write_json(document_id)
+            self._write_text(":")
+            self._write_json(document)
+            document_count += 1
+        self._write_text("}")
         if granularity == "month":
             self._write_text(',"year":')
             self.next_granularity = "year"
         else:
             self.next_granularity = None
+        return document_count
 
     def finalize(self, input_log_hashes, summary):
         if self.next_granularity is not None:

metrics/services/export.py

@@ -1,13 +1,12 @@
 import logging
-import resource
 from itertools import chain
 from time import monotonic
 
 from django.conf import settings
 
 from metrics.opensearch.mappings import get_index_mappings
 from metrics.opensearch.names import generate_month_index_name, generate_year_index_name
-from metrics.services import daily_payloads
+from metrics.services import daily_payloads, memory
 
 
 def daily_metric_payload_exists(job):
@@ -40,12 +39,12 @@ def export_daily_metric_payload(search_client, job):
         )
         logging.info(
             "Daily metric job %s %s OpenSearch export completed in %.3f "
-            "seconds; %s documents; peak RSS %.1f MiB.",
+            "seconds; %s documents; %s.",
             job.pk,
             granularity,
             monotonic() - started,
             exported,
-            _peak_rss_mib(),
+            memory.format_snapshot(),
         )
 
 
@@ -86,7 +85,3 @@ def _sync_documents_group(
         document_items=chain((first_item,), document_items),
         job_id=job_id,
     )
-
-
-def _peak_rss_mib():
-    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024

metrics/services/memory.py

@@ -0,0 +1,55 @@
+import resource
+from pathlib import Path
+
+_CGROUP_CURRENT_PATHS = (
+    Path("/sys/fs/cgroup/memory.current"),
+    Path("/sys/fs/cgroup/memory/memory.usage_in_bytes"),
+)
+_CGROUP_PEAK_PATHS = (
+    Path("/sys/fs/cgroup/memory.peak"),
+    Path("/sys/fs/cgroup/memory/memory.max_usage_in_bytes"),
+)
+_MIB = 1024 * 1024
+
+
+def snapshot():
+    return {
+        "rss_mib": _current_rss_mib(),
+        "peak_rss_mib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,
+        "cgroup_current_mib": _read_cgroup_mib(_CGROUP_CURRENT_PATHS),
+        "cgroup_peak_mib": _read_cgroup_mib(_CGROUP_PEAK_PATHS),
+    }
+
+
+def format_snapshot(values=None):
+    values = values or snapshot()
+    parts = [
+        f"RSS {values['rss_mib']:.1f} MiB",
+        f"peak RSS {values['peak_rss_mib']:.1f} MiB",
+    ]
+    if values["cgroup_current_mib"] is not None:
+        parts.append(f"cgroup current {values['cgroup_current_mib']:.1f} MiB")
+    if values["cgroup_peak_mib"] is not None:
+        parts.append(f"cgroup peak {values['cgroup_peak_mib']:.1f} MiB")
+    return "; ".join(parts)
+
+
+def _current_rss_mib():
+    try:
+        for line in Path("/proc/self/status").read_text().splitlines():
+            if line.startswith("VmRSS:"):
+                return int(line.split()[1]) / 1024
+    except (OSError, ValueError, IndexError):
+        pass
+    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
+
+
+def _read_cgroup_mib(paths):
+    for path in paths:
+        try:
+            value = path.read_text().strip()
+            if value != "max":
+                return int(value) / _MIB
+        except (OSError, ValueError):
+            continue
+    return None