roll-prebuilt: update also heap_profile and record_android_trace

Updates the //tools/roll-prebuilt to deal also with
injection into more complex scripts like heap_profile and
record_android_trace.
The script now works as follows:
- In all the files that need prebuilts we inject two
  markers:
  # BEGIN_SECTION_GENERATED_BY(roll-prebuilts)
  # END_SECTION_GENERATED_BY(roll-prebuilts)
- tools/roll-prebuilt injects the get_perfetto_prebuilt()
  function and the manifest within those markers, leaving
  the rest of the file intact.

This CL does NOT yet run the script and make the matching
changes in the other script files, aosp/1780040 does that.

Bug: 177349647

Change-Id: I78a625c9027eca5973a871e42b7eabad1f4dc8fa
diff --git a/tools/get_perfetto_prebuilt.py b/tools/get_perfetto_prebuilt.py
new file mode 100644
index 0000000..28e9806
--- /dev/null
+++ b/tools/get_perfetto_prebuilt.py
@@ -0,0 +1,122 @@
+# Copyright (C) 2021 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This source defines a self-contained function to fetch a perfetto prebuilt.
+
+This function is copy/pasted by //tools/roll-prebuilts in different places:
+- Into the //tools/{trace_processor, traceconv} scripts, which are just plain
+  wrappers around executables.
+- Into the //tools/{heap_profiler, record_android_trace} scripts, which contain
+  some other hand-written python code.
+In both cases toll-prebuilts copies this source (together with a manifest) into
+a section annotated with "BEGIN_SECTION_GENERATED_BY(roll-prebuilts)" / END... .
+The automated-copy-paste is to keep those script hermetic, so people can just
+download and run them without checking out the repo.
+
+The manifest argument looks as follows in the generated files:
+PERFETTO_PREBUILT_MANIFEST = [{
+    'tool': 'trace_to_text',
+    'arch': 'mac-amd64',
+    'file_name': 'trace_to_text',
+    'file_size': 7087080,
+    'url': https://commondatastorage.googleapis.com/.../trace_to_text',
+    'sha256': 7d957c005b0dc130f5bd855d6cec27e060d38841b320d04840afc569f9087490',
+    'platform': 'darwin',
+    'machine': 'x86_64'
+  },
+  ...
+]
+
+The intended usage is:
+
+bin_path = get_perfetto_prebuilt('trace_processor_shell')
+subprocess.call(bin_path, ...)
+"""
+
+from logging import exception
+
+PERFETTO_PREBUILT_MANIFEST = []
+
+# COPIED_SECTION_START_MARKER
+
+
+# DO NOT EDIT. If you wish to make edits to this code, you need to change only
+# //tools/get_perfetto_prebuilt.py and run /tools/roll-prebuilts to regenerate
+# all the others scripts this is embedded into.
+def get_perfetto_prebuilt(tool_name, soft_fail=False, arch=None):
+  """ Downloads the prebuilt, if necessary, and returns its path on disk. """
+
+  # The first time this is invoked, it downloads the |url| and caches it into
+  # ~/.perfetto/prebuilts/$tool_name. On subsequent invocations it just runs the
+  # cached version.
+  def download_or_get_cached(file_name, url, sha256):
+    import os, hashlib, subprocess
+    dir = os.path.join(
+        os.path.expanduser('~'), '.local', 'share', 'perfetto', 'prebuilts')
+    os.makedirs(dir, exist_ok=True)
+    bin_path = os.path.join(dir, file_name)
+    sha256_path = os.path.join(dir, file_name + '.sha256')
+    needs_download = True
+
+    # Avoid recomputing the SHA-256 on each invocation. The SHA-256 of the last
+    # download is cached into file_name.sha256, just check if that matches.
+    if os.path.exists(bin_path) and os.path.exists(sha256_path):
+      with open(sha256_path, 'rb') as f:
+        digest = f.read().decode()
+        if digest == sha256:
+          needs_download = False
+
+    if needs_download:
+      # Either the filed doesn't exist or the SHA256 doesn't match.
+      tmp_path = bin_path + '.tmp'
+      print('Downloading ' + url)
+      subprocess.check_call(['curl', '-f', '-L', '-#', '-o', tmp_path, url])
+      with open(tmp_path, 'rb') as fd:
+        actual_sha256 = hashlib.sha256(fd.read()).hexdigest()
+      if actual_sha256 != sha256:
+        raise Exception('Checksum mismatch for %s (actual: %s, expected: %s)' %
+                        (url, actual_sha256, sha256))
+      os.chmod(tmp_path, 0o755)
+      os.rename(tmp_path, bin_path)
+      with open(sha256_path, 'w') as f:
+        f.write(sha256)
+    return bin_path
+    # --- end of download_or_get_cached() ---
+
+  # --- get_perfetto_prebuilt() function starts here. ---
+  import os, platform, sys
+  plat = sys.platform.lower()
+  machine = platform.machine().lower()
+  manifest_entry = None
+  for entry in PERFETTO_PREBUILT_MANIFEST:
+    # If the caller overrides the arch, just match that (for Android prebuilts).
+    if arch and entry.get('arch') == arch:
+      manifest_entry = entry
+      break
+    # Otherwise guess the local machine arch.
+    if entry.get('tool') == tool_name and entry.get(
+        'platform') == plat and entry.get('machine') == machine:
+      manifest_entry = entry
+      break
+  if manifest_entry is None:
+    if soft_fail:
+      return None
+    raise Exception(
+        ('No prebuilts available for %s-%s\n' % (plat, machine)) +
+        'See https://perfetto.dev/docs/contributing/build-instructions')
+
+  return download_or_get_cached(
+      file_name=manifest_entry['file_name'],
+      url=manifest_entry['url'],
+      sha256=manifest_entry['sha256'])
diff --git a/tools/prebuilt_template.py b/tools/prebuilt_template.py
deleted file mode 100755
index aefb83a..0000000
--- a/tools/prebuilt_template.py
+++ /dev/null
@@ -1,104 +0,0 @@
-#!/usr/bin/env python3
-# Copyright (C) 2021 The Android Open Source Project
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# This is the template used by tools/roll-prebuilts to generate the script
-# tools/{trace_processor, traceconv, tracebox} which in turn are served by
-# https://get.perfetto.dev/{trace_processor, traceconv, tracebox}
-#
-# This file should do the same thing when being invoked in any of these ways:
-# ./tool_name
-# python tool_name
-# bash tool_name
-# cat ./tool_name | bash
-# cat ./tool_name | python -
-
-BASH_FALLBACK = """ "
-exec python3 - "$@" <<'#'EOF
-#"""
-
-import hashlib
-import os
-import platform
-import subprocess
-import sys
-
-# The placeholder below will be replaced with something like:
-# TOOL_NAME = 'trace_processor_shell'
-# MANIFEST = [{'tool': 'trace_processor_shell', ...]
-TOOL_NAME = ''
-MANIFEST = []
-
-# REPLACEMENT_PLACEHOLDER
-
-
-# The first time this is invoked, it downloads the |url| and caches it into
-# ~/.perfetto/prebuilts/$tool_name. On subsequent invocations it just runs the
-# cached version.
-def download_or_get_cached(file_name, url, sha256):
-  dir = os.path.join(os.path.expanduser('~'), '.local', 'share', 'perfetto', 'prebuilts')
-  os.makedirs(dir, exist_ok=True)
-  bin_path = os.path.join(dir, file_name)
-  sha256_path = os.path.join(dir, file_name + '.sha256')
-  needs_download = True
-
-  # Avoid recomputing the SHA-256 on each invocation. The SHA-256 of the last
-  # download is cached into file_name.sha256, just check if that matches.
-  if os.path.exists(bin_path) and os.path.exists(sha256_path):
-    with open(sha256_path, 'rb') as f:
-      digest = f.read().decode()
-      if digest == sha256:
-        needs_download = False
-
-  if needs_download:
-    # Either the filed doesn't exist or the SHA256 doesn't match.
-    tmp_path = bin_path + '.tmp'
-    print('Downloading ' + url)
-    subprocess.check_call(['curl', '-f', '-L', '-#', '-o', tmp_path, url])
-    with open(tmp_path, 'rb') as fd:
-      actual_sha256 = hashlib.sha256(fd.read()).hexdigest()
-    if actual_sha256 != sha256:
-      raise 'Checksum mismatch for %s (actual: %s, expected: %s)' % (
-          url, actual_sha256, sha256)
-    os.chmod(tmp_path, 0o755)
-    os.rename(tmp_path, bin_path)
-    with open(sha256_path, 'w') as f:
-      f.write(sha256)
-  return bin_path
-
-
-def main(argv):
-  plat = sys.platform.lower()
-  machine = platform.machine().lower()
-  manifest = None
-  for entry in MANIFEST:
-    if entry.get('tool') == TOOL_NAME and entry.get(
-        'platform') == plat and entry.get('machine') == machine:
-      manifest = entry
-      break
-  if manifest is None:
-    print('No prebuilts available for %s/%s' % (plat, machine))
-    print('See https://perfetto.dev/docs/contributing/build-instructions')
-    return 1
-  bin_path = download_or_get_cached(
-      file_name=manifest['file_name'],
-      url=manifest['url'],
-      sha256=manifest['sha256'])
-  os.execv(bin_path, [bin_path] + argv[1:])
-
-
-if __name__ == '__main__':
-  sys.exit(main(sys.argv))
-
-#EOF
diff --git a/tools/roll-prebuilts b/tools/roll-prebuilts
index fb38ddf..5cd3aee 100755
--- a/tools/roll-prebuilts
+++ b/tools/roll-prebuilts
@@ -14,11 +14,11 @@
 # limitations under the License.
 """Updates the python scripts in tools/{trace_processor, traceconv, tracebox}
 
-This script does the following, for each prebuilt in PREBUILTS:
+This script does the following, for each entry in SCRIPTS_TO_UPDATE:
   - Downloads the artifact by the LUCI infrastructure, one for each arch.
   - Computes the SHA-256 of each artifact.
   - Generates a manifest with URL, SHA-256 and other details.
-  - Merges tools/prebuilt_template.py with the manifest and writes tools/xxx.
+  - Merges get_perfetto_prebuilt.py with the manifest and writes tools/xxx.
 
 This script is supposed to be run by Perfetto OWNERS after every monthly release
 after the LUCI jobs have completed.
@@ -39,26 +39,36 @@
 TOOLS_DIR = os.path.join(ROOT_DIR, 'tools')
 GIT_REV = subprocess.check_output(['git', 'rev-parse', 'head']).decode().strip()
 
-PREBUILTS = [
+SCRIPTS_TO_UPDATE = [
     {
-        'tool': 'trace_processor_shell',
         'script': 'trace_processor',
+        'tool': 'trace_processor_shell',
         'archs': ['mac-amd64', 'linux-amd64', 'windows-amd64']
     },
     {
-        'tool': 'trace_to_text',
         'script': 'traceconv',
+        'tool': 'trace_to_text',
         'archs': ['mac-amd64', 'linux-amd64', 'windows-amd64']
     },
     {
-        'tool': 'tracebox',
         'script': 'tracebox',
+        'tool': 'tracebox',
         'archs': ['mac-amd64', 'linux-amd64']
     },
+    {
+        'script': 'heap_profile',
+        'tool': 'trace_to_text',
+        'archs': ['mac-amd64', 'linux-amd64', 'windows-amd64']
+    },
+    {
+        'script': 'record_android_trace',
+        'tool': 'tracebox',
+        'archs': ['android-arm', 'android-arm64', 'android-x86', 'android-x64']
+    },
 ]
 
-# Maps a CIPD 'os-arch' string into corresponding tuples that match against
-# python's platform / machine API (see prebuilt_template.py).
+# Maps a 'os-arch' string into corresponding tuples that match against
+# python's platform / machine API (see get_perfetto_prebuilt.py).
 ARCH_TO_PYTHON = {
     'mac-amd64': {
         'platform': 'darwin',
@@ -93,25 +103,44 @@
   return manifest
 
 
-# Takes tool/prebuilt_template.py, replaces the manifest in it and writes the
-# result into tools/$script_name.
+# Returns the section of get_perfetto_prebuilt.py which should be copy/pasted
+# in the various scripts.
+def read_get_perfetto_prebuilt_script():
+  in_file = os.path.join(TOOLS_DIR, 'get_perfetto_prebuilt.py')
+  with open(in_file, 'r') as f:
+    contents = f.read()
+  return contents.split('COPIED_SECTION_START_MARKER')[1]
+
+
 def update_script(git_revision, tool_name, script_name, archs):
-  with open(os.path.join(TOOLS_DIR, 'prebuilt_template.py'), 'r') as f:
-    template = f.read()
   with ThreadPoolExecutor(max_workers=8) as executor:
     manifests = list(
         executor.map(lambda arch: make_manifest(git_revision, tool_name, arch),
                      archs))
+  out_file = os.path.join(TOOLS_DIR, script_name)
+  with open(out_file) as f:
+    script = f.read()
 
-    repl = '# Generated by %s @ %s\n' % (__file__, GIT_REV)
-    repl += 'TOOL_NAME = \'%s\'\n' % tool_name
-    repl += 'MANIFEST = %s\n' % str(manifests)
-    script = template.replace('# REPLACEMENT_PLACEHOLDER', repl)
-    out_file = os.path.join(TOOLS_DIR, script_name)
-    with open(out_file + '.tmp', 'w') as f:
-      f.write(script)
-    os.rename(out_file + '.tmp', out_file)
-    os.chmod(out_file, 0o755)
+  begin_marker = '\n# BEGIN_SECTION_GENERATED_BY(roll-prebuilts)\n'
+  end_marker = '\n# END_SECTION_GENERATED_BY(roll-prebuilts)\n'
+  before = script.partition(begin_marker)[0]
+  after = script.partition(end_marker)[2]
+
+  content = '# Revision: {git_revision}\n'
+  content += 'PERFETTO_PREBUILT_MANIFEST = {manifests}\n'
+  content += '{fn_body}\n'
+  content = content.format(
+      git_revision=git_revision,
+      manifests=str(manifests),
+      fn_body=read_get_perfetto_prebuilt_script())
+
+  script = before + begin_marker + content + end_marker + after
+
+  with open(out_file + '.tmp', 'w') as f:
+    f.write(script)
+  subprocess.check_call(['yapf', '-i', out_file + '.tmp'])
+  os.rename(out_file + '.tmp', out_file)
+  os.chmod(out_file, 0o755)
 
 
 def main():
@@ -127,7 +156,7 @@
     return 1
 
   git_revision = args.revision
-  for spec in PREBUILTS:
+  for spec in SCRIPTS_TO_UPDATE:
     logging.info('Updating %s', spec['script'])
     update_script(git_revision, spec['tool'], spec['script'], spec['archs'])