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/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'])