ninjalog_uploader: Add is_cloudtop, gce_machine_type, is_cog to the Ninjalog metadata

Bug: 367856712
Change-Id: Ibaa96d779cc3fc896e8203525346e3ca27d9d8b0
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/5877849
Reviewed-by: Takuto Ikuta <tikuta@chromium.org>
Reviewed-by: Fumitoshi Ukai <ukai@google.com>
Commit-Queue: Junji Watanabe <jwata@google.com>

ninjalog_uploader.py

@@ -27,12 +27,14 @@ import multiprocessing
 import os
 import pathlib
 import platform
+import re
 import subprocess
 import sys
 import time
 import urllib.request
 
 import build_telemetry
+import gclient_utils
 
 # Configs that should not be uploaded as is.
 SENSITIVE_CONFIGS = (
@@ -160,11 +162,14 @@ def GetMetadata(cmdline, ninjalog, exit_code, build_duration, user):
         "build_duration_sec": build_duration,
         "platform": platform.system(),
         "cpu_core": multiprocessing.cpu_count(),
+        "is_cog": gclient_utils.IsEnvCog(),
         "build_configs": build_configs,
         "explicit_build_config_keys": explicit_keys,
         "targets": GetBuildTargetFromCommandLine(cmdline),
     }
 
+    metadata.update(GetGCEMetadata())
+
     invocation_id = os.environ.get("AUTONINJA_BUILD_ID")
     if invocation_id:
         metadata['invocation_id'] = invocation_id
@@ -175,6 +180,33 @@ def GetMetadata(cmdline, ninjalog, exit_code, build_duration, user):
     return metadata
 
 
+def GetGCEMetadata():
+    gce = _getGCEInfo()
+    if not gce:
+        return {}
+    md = {}
+    if "cloudtop" in gce.get("project", {}).get("projectId", ""):
+        md["is_cloudtop"] = True
+    match = re.search(r"machineTypes/([^/]+)",
+                      gce.get("instance", {}).get("machineType", ""))
+    if match:
+        md["gce_machine_type"] = match.group(1)
+    return md
+
+
+def _getGCEInfo():
+    url = "http://metadata.google.internal/computeMetadata/v1/?recursive=true"
+    request = urllib.request.Request(url, headers={"Metadata-Flavor": "Google"})
+    try:
+        response = urllib.request.urlopen(request)
+        meta = json.load(response)
+    except Exception as e:
+        # Only GCE machines can access to the metadata server.
+        logging.warning(e)
+        return
+    return meta
+
+
 def GetNinjalog(cmdline):
     """GetNinjalog returns the path to ninjalog from cmdline."""
     # ninjalog is in current working directory by default.

tests/ninjalog_uploader_test.py

@@ -3,11 +3,13 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import http
 import json
 import os
 import sys
 import unittest
 import unittest.mock
+import urllib.error
 
 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.insert(0, ROOT_DIR)
@@ -195,6 +197,43 @@ class NinjalogUploaderTest(unittest.TestCase):
 
         self.assertEqual(ninjalog_uploader.GetJflag(['ninja', '-j']), None)
 
+    def test_get_gce_metadata(self):
+        with unittest.mock.patch('urllib.request.urlopen') as urlopen_mock:
+            urlopen_mock.side_effect = urllib.error.HTTPError(
+                'http://test/not-found', http.HTTPStatus.NOT_FOUND, 'not found',
+                None, None)
+            self.assertEqual(ninjalog_uploader.GetGCEMetadata(), {})
+
+        with unittest.mock.patch(
+                'ninjalog_uploader._getGCEInfo') as getGCEInfo_mock:
+            getGCEInfo_mock.return_value = {
+                "instance": {
+                    "machineType":
+                    "projects/12345/machineTypes/n2d-standard-128",
+                },
+                "project": {
+                    "projectId": "cloudtop-test"
+                }
+            }
+            self.assertEqual(ninjalog_uploader.GetGCEMetadata(), {
+                'gce_machine_type': 'n2d-standard-128',
+                'is_cloudtop': True,
+            })
+
+        with unittest.mock.patch(
+                'ninjalog_uploader._getGCEInfo') as getGCEInfo_mock:
+            getGCEInfo_mock.return_value = {
+                "instance": {
+                    "machineType":
+                    "projects/12345/machineTypes/n2d-standard-128",
+                },
+                "project": {
+                    "projectId": "gce-project"
+                }
+            }
+            self.assertEqual(ninjalog_uploader.GetGCEMetadata(), {
+                'gce_machine_type': 'n2d-standard-128',
+            })
 
 if __name__ == '__main__':
     unittest.main()