bugs: emit reports in an order that bug_manager will respect

cl/508145633 has our bug_manager always iterate through new JSON reports
in sorted order (where "sorted" means "sorted by their file names").
This change has us emit reports from a process in a consistent order.

BUG=b:268368949
TEST=unittests

Fixed: b:268368949
Change-Id: I61d02079338e85b0b1e68749c51dafe1866d98e7
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/4234617
Commit-Queue: George Burgess <gbiv@chromium.org>
Reviewed-by: Jordan Abrahams-Whitehead <ajordanr@google.com>
Tested-by: George Burgess <gbiv@chromium.org>

2 files changed, 88 insertions, 11 deletions

diff --git a/cros_utils/bugs.py b/cros_utils/bugs.py
index 9a84ac40..25da130d 100755
--- a/cros_utils/bugs.py
+++ b/cros_utils/bugs.py
@@ -5,17 +5,16 @@
 
 """Utilities to file bugs."""
 
-import base64
 import datetime
 import enum
 import json
 import os
+import threading
 from typing import Any, Dict, List, Optional
 
 
 X20_PATH = "/google/data/rw/teams/c-compiler-chrome/prod_bugs"
 
-
 # These constants are sourced from
 # //google3/googleclient/chrome/chromeos_toolchain/bug_manager/bugs.go
 class WellKnownComponents(enum.IntEnum):
@@ -26,6 +25,47 @@ class WellKnownComponents(enum.IntEnum):
     AndroidRustToolchain = -3
 
 
+class _FileNameGenerator:
+    """Generates unique file names. This container is thread-safe.
+
+    The names generated have the following properties:
+        - successive, sequenced calls to `get_json_file_name()` will produce
+          names that sort later in lists over time (e.g.,
+          [generator.generate_json_file_name() for _ in range(10)] will be in
+          sorted order).
+        - file names cannot collide with file names generated on the same
+          machine (ignoring machines with unreasonable PID reuse).
+        - file names are incredibly unlikely to collide when generated on
+          multiple machines, as they have 8 bytes of entropy in them.
+    """
+
+    _RANDOM_BYTES = 8
+    _MAX_OS_ENTROPY_VALUE = 1 << _RANDOM_BYTES * 8
+    # The intent of this is "the maximum possible size of our entropy string,
+    # so we can zfill properly below." Double the value the OS hands us, since
+    # we add to it in `generate_json_file_name`.
+    _ENTROPY_STR_SIZE = len(str(2 * _MAX_OS_ENTROPY_VALUE))
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._entropy = int.from_bytes(
+            os.getrandom(self._RANDOM_BYTES), byteorder="little", signed=False
+        )
+
+    def generate_json_file_name(self, now: datetime.datetime):
+        with self._lock:
+            my_entropy = self._entropy
+            self._entropy += 1
+
+        now = now.isoformat("T", "seconds") + "Z"
+        entropy_str = str(my_entropy).zfill(self._ENTROPY_STR_SIZE)
+        pid = os.getpid()
+        return f"{now}_{entropy_str}_{pid}.json"
+
+
+_GLOBAL_NAME_GENERATOR = _FileNameGenerator()
+
+
 def _WriteBugJSONFile(object_type: str, json_object: Dict[str, Any]):
     """Writes a JSON file to X20_PATH with the given bug-ish object."""
     final_object = {
@@ -33,15 +73,10 @@ def _WriteBugJSONFile(object_type: str, json_object: Dict[str, Any]):
         "value": json_object,
     }
 
-    # The name of this has two parts:
-    # - An easily sortable time, to provide uniqueness and let our service send
-    #   things in the order they were put into the outbox.
-    # - 64 bits of entropy, so two racing bug writes don't clobber the same file.
-    now = datetime.datetime.utcnow().isoformat("T", "seconds") + "Z"
-    entropy = base64.urlsafe_b64encode(os.getrandom(8))
-    entropy_str = entropy.rstrip(b"=").decode("utf-8")
-    file_path = os.path.join(X20_PATH, f"{now}_{entropy_str}.json")
-
+    now = datetime.datetime.now(tz=datetime.timezone.utc)
+    file_path = os.path.join(
+        X20_PATH, _GLOBAL_NAME_GENERATOR.generate_json_file_name(now)
+    )
     temp_path = file_path + ".in_progress"
     try:
         with open(temp_path, "w") as f:
diff --git a/cros_utils/bugs_test.py b/cros_utils/bugs_test.py
index 5a07dbd8..53525b30 100755
--- a/cros_utils/bugs_test.py
+++ b/cros_utils/bugs_test.py
@@ -8,6 +8,7 @@
 
 """Tests bug filing bits."""
 
+import datetime
 import json
 import tempfile
 import unittest
@@ -16,6 +17,9 @@ from unittest.mock import patch
 import bugs
 
 
+_ARBITRARY_DATETIME = datetime.datetime(2020, 1, 1, 23, 0, 0, 0)
+
+
 class Tests(unittest.TestCase):
     """Tests for the bugs module."""
 
@@ -126,6 +130,44 @@ class Tests(unittest.TestCase):
             },
         )
 
+    def testFileNameGenerationProducesFileNamesInSortedOrder(self):
+        """Tests that _FileNameGenerator gives us sorted file names."""
+        gen = bugs._FileNameGenerator()
+        first = gen.generate_json_file_name(_ARBITRARY_DATETIME)
+        second = gen.generate_json_file_name(_ARBITRARY_DATETIME)
+        self.assertLess(first, second)
+
+    def testFileNameGenerationProtectsAgainstRipplingAdds(self):
+        """Tests that _FileNameGenerator gives us sorted file names."""
+        gen = bugs._FileNameGenerator()
+        gen._entropy = 9
+        first = gen.generate_json_file_name(_ARBITRARY_DATETIME)
+        second = gen.generate_json_file_name(_ARBITRARY_DATETIME)
+        self.assertLess(first, second)
+
+        gen = bugs._FileNameGenerator()
+        all_9s = "9" * (gen._ENTROPY_STR_SIZE - 1)
+        gen._entropy = int(all_9s)
+        third = gen.generate_json_file_name(_ARBITRARY_DATETIME)
+        self.assertLess(second, third)
+
+        fourth = gen.generate_json_file_name(_ARBITRARY_DATETIME)
+        self.assertLess(third, fourth)
+
+    @patch("os.getpid")
+    def testForkingProducesADifferentReport(self, mock_getpid):
+        """Tests that _FileNameGenerator gives us sorted file names."""
+        gen = bugs._FileNameGenerator()
+
+        mock_getpid.return_value = 1
+        gen._entropy = 0
+        parent_file = gen.generate_json_file_name(_ARBITRARY_DATETIME)
+
+        mock_getpid.return_value = 2
+        gen._entropy = 0
+        child_file = gen.generate_json_file_name(_ARBITRARY_DATETIME)
+        self.assertNotEqual(parent_file, child_file)
+
 
 if __name__ == "__main__":
     unittest.main()