Snap for 11819063 from 53dd8e309607e80ba27c58bc91106b1c6fb1a30d to sdk-releasesdk-release

Change-Id: Ifaed9e18e83c145cf5e132e252366d4b2799cf86

95 files changed, 2132 insertions, 1029 deletions

diff --git a/Android.bp b/Android.bp
index ae8225dae..c43694efc 100644
--- a/Android.bp
+++ b/Android.bp
@@ -13477,6 +13477,7 @@ filegroup {
         "src/trace_redaction/redact_process_free.cc",
         "src/trace_redaction/redact_sched_switch.cc",
         "src/trace_redaction/redact_task_newtask.cc",
+        "src/trace_redaction/remap_scheduling_events.cc",
         "src/trace_redaction/scrub_ftrace_events.cc",
         "src/trace_redaction/scrub_process_stats.cc",
         "src/trace_redaction/scrub_process_trees.cc",
@@ -13505,6 +13506,7 @@ filegroup {
         "src/trace_redaction/redact_process_free_unittest.cc",
         "src/trace_redaction/redact_sched_switch_unittest.cc",
         "src/trace_redaction/redact_task_newtask_unittest.cc",
+        "src/trace_redaction/remap_scheduling_events_unittest.cc",
         "src/trace_redaction/suspend_resume_unittest.cc",
     ],
 }
diff --git a/CHANGELOG b/CHANGELOG
index eb1dab57a..407a08dad 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -2,7 +2,8 @@ Unreleased:
   Tracing service and probes:
     *
   Trace Processor:
-    *
+    * Optimised single column `DISTINCT` queries.
+    * Optimised `LIMIT` and `OFFSET` queries.
   SQL Standard library:
     * Improved support for querying startups on Android 9 (API level 28) and
       below. Available in `android.startup.startups` module.
diff --git a/python/generators/diff_tests/runner.py b/python/generators/diff_tests/runner.py
index 7c77d544b..0fe0982ff 100644
--- a/python/generators/diff_tests/runner.py
+++ b/python/generators/diff_tests/runner.py
@@ -450,7 +450,7 @@ class DiffTestsRunner:
                          override_sql_module_paths))
 
   def run_all_tests(self, metrics_descriptor_paths: List[str],
-                    chrome_extensions: str, test_extensions: str,
+                    chrome_extensions: str, test_extensions: str, winscope_extensions: str,
                     keep_input: bool, rebase: bool) -> TestResults:
     perf_results = []
     failures = []
@@ -459,7 +459,7 @@ class DiffTestsRunner:
 
     with concurrent.futures.ProcessPoolExecutor() as e:
       fut = [
-          e.submit(test.execute, [chrome_extensions, test_extensions],
+          e.submit(test.execute, [chrome_extensions, test_extensions, winscope_extensions],
                    metrics_descriptor_paths, keep_input, rebase)
           for test in self.test_runners
       ]
diff --git a/python/tools/install_test_reporter_app.py b/python/tools/install_test_reporter_app.py
new file mode 100755
index 000000000..8673bbe19
--- /dev/null
+++ b/python/tools/install_test_reporter_app.py
@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+# Copyright (C) 2024 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.
+
+import argparse
+import os
+import re
+import sys
+import tempfile
+import time
+
+from perfetto.prebuilts.perfetto_prebuilts import *
+
+PERMSISION_REGEX = re.compile(r'''uses-permission: name='(.*)'.*''')
+NAME_REGEX = re.compile(r'''package: name='(.*?)' .*''')
+
+
+def cmd(args: list[str]):
+  print('Running command ' + ' '.join(args))
+  return subprocess.check_output(args)
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--apk', help='Local APK to use instead of builtin')
+
+  args = parser.parse_args()
+
+  if args.apk:
+    apk = args.apk
+  else:
+    apk = download_or_get_cached(
+        'CtsPerfettoReporterApp.apk',
+        'https://storage.googleapis.com/perfetto/CtsPerfettoReporterApp.apk',
+        'f21dda36668c368793500b13724ab2a6231d12ded05746f7cfaaba4adedd7d46')
+
+  # Figure out the package name and the permissions we need
+  aapt = subprocess.check_output(['aapt', 'dump', 'badging', apk]).decode()
+  permission_names = []
+  name = ''
+  for l in aapt.splitlines():
+    name_match = NAME_REGEX.match(l)
+    if name_match:
+      name = name_match[1]
+      continue
+
+    permission_match = PERMSISION_REGEX.match(l)
+    if permission_match:
+      permission_names.append(permission_match[1])
+      continue
+
+  # Root and remount the device.
+  cmd(['adb', 'root'])
+  cmd(['adb', 'wait-for-device'])
+  cmd(['adb', 'remount', '-R'])
+  input('The device might now reboot. If so, please unlock the device then '
+        'press enter to continue')
+  cmd(['adb', 'wait-for-device'])
+  cmd(['adb', 'root'])
+  cmd(['adb', 'wait-for-device'])
+  cmd(['adb', 'remount', '-R'])
+
+  # Write out the permission file on device.
+  permissions = '\n'.join(
+      f'''<permission name='{p}' />''' for p in permission_names)
+  permission_file_contents = f'''
+    <permissions>
+      <privapp-permissions package="{name}">
+        {permissions}
+      </privapp-permissions>
+    </permissions>
+  '''
+  with tempfile.NamedTemporaryFile() as f:
+    f.write(permission_file_contents.encode())
+    f.flush()
+
+    cmd([
+        'adb', 'push', f.name,
+        f'/system/etc/permissions/privapp-permissions-{name}.xml'
+    ])
+
+  # Stop system_server, push the apk on device and restart system_server
+  priv_app_path = f'/system/priv-app/{name}/{name}.apk'
+  cmd(['adb', 'shell', 'stop'])
+  cmd(['adb', 'push', apk, priv_app_path])
+  cmd(['adb', 'shell', 'start'])
+  cmd(['adb', 'wait-for-device'])
+  time.sleep(10)
+
+  # Wait for system_server and package manager to come up.
+  while 'system_server' not in cmd(['adb', 'shell', 'ps']).decode():
+    time.sleep(1)
+  while True:
+    ps = set([
+        l.strip()
+        for l in cmd(['adb', 'shell', 'dumpsys', '-l']).decode().splitlines()
+    ])
+    if 'storaged' in ps and 'settings' in ps and 'package' in ps:
+      break
+    time.sleep(1)
+
+  # Install the actual APK.
+  cmd(['adb', 'shell', 'pm', 'install', '-r', '-d', '-g', '-t', priv_app_path])
+
+  return 0
+
+
+sys.exit(main())
diff --git a/src/trace_processor/importers/common/process_tracker.cc b/src/trace_processor/importers/common/process_tracker.cc
index f03486772..3fc95fcbd 100644
--- a/src/trace_processor/importers/common/process_tracker.cc
+++ b/src/trace_processor/importers/common/process_tracker.cc
@@ -16,12 +16,21 @@
 
 #include "src/trace_processor/importers/common/process_tracker.h"
 
+#include <algorithm>
+#include <cstdint>
+#include <optional>
 #include <utility>
+#include <vector>
 
+#include "perfetto/base/logging.h"
+#include "perfetto/ext/base/string_view.h"
+#include "perfetto/public/compiler.h"
 #include "src/trace_processor/storage/stats.h"
+#include "src/trace_processor/storage/trace_storage.h"
+#include "src/trace_processor/tables/metadata_tables_py.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
-namespace perfetto {
-namespace trace_processor {
+namespace perfetto::trace_processor {
 
 ProcessTracker::ProcessTracker(TraceProcessorContext* context)
     : context_(context), args_tracker_(context) {
@@ -187,11 +196,8 @@ bool ProcessTracker::IsThreadAlive(UniqueTid utid) {
 
   // If the process has been replaced in |pids_|, this thread is dead.
   uint32_t current_pid = processes->pid()[current_upid];
-  auto pid_it = pids_.Find(current_pid);
-  if (pid_it && *pid_it != current_upid)
-    return false;
-
-  return true;
+  auto* pid_it = pids_.Find(current_pid);
+  return !pid_it || *pid_it == current_upid;
 }
 
 std::optional<UniqueTid> ProcessTracker::GetThreadOrNull(
@@ -200,7 +206,7 @@ std::optional<UniqueTid> ProcessTracker::GetThreadOrNull(
   auto* threads = context_->storage->mutable_thread_table();
   auto* processes = context_->storage->mutable_process_table();
 
-  auto vector_it = tids_.Find(tid);
+  auto* vector_it = tids_.Find(tid);
   if (!vector_it)
     return std::nullopt;
 
@@ -576,5 +582,4 @@ void ProcessTracker::UpdateNamespacedThread(uint32_t pid,
   namespaced_threads_[tid] = {pid, tid, std::move(nstid)};
 }
 
-}  // namespace trace_processor
-}  // namespace perfetto
+}  // namespace perfetto::trace_processor
diff --git a/src/trace_processor/importers/common/process_tracker.h b/src/trace_processor/importers/common/process_tracker.h
index e5ccd3999..fa2ffa47b 100644
--- a/src/trace_processor/importers/common/process_tracker.h
+++ b/src/trace_processor/importers/common/process_tracker.h
@@ -17,10 +17,11 @@
 #ifndef SRC_TRACE_PROCESSOR_IMPORTERS_COMMON_PROCESS_TRACKER_H_
 #define SRC_TRACE_PROCESSOR_IMPORTERS_COMMON_PROCESS_TRACKER_H_
 
-#include <stdint.h>
-
+#include <cstdint>
 #include <optional>
+#include <unordered_map>
 #include <unordered_set>
+#include <utility>
 #include <vector>
 
 #include "perfetto/ext/base/flat_hash_map.h"
@@ -96,7 +97,7 @@ class ProcessTracker {
   // Called when a thread is seen the process tree. Retrieves the matching utid
   // for the tid and the matching upid for the tgid and stores both.
   // Virtual for testing.
-  virtual UniqueTid UpdateThread(uint32_t tid, uint32_t tgid);
+  virtual UniqueTid UpdateThread(uint32_t tid, uint32_t pid);
 
   // Associates trusted_pid with track UUID.
   void UpdateTrustedPid(uint32_t trusted_pid, uint64_t uuid);
diff --git a/src/trace_processor/importers/etw/BUILD.gn b/src/trace_processor/importers/etw/BUILD.gn
index 23d8965fd..5e14c2fe6 100644
--- a/src/trace_processor/importers/etw/BUILD.gn
+++ b/src/trace_processor/importers/etw/BUILD.gn
@@ -51,6 +51,7 @@ source_set("full") {
     "../common:parser_types",
     "../i2c:full",
     "../proto:minimal",
+    "../proto:packet_sequence_state_generation_hdr",
     "../syscalls:full",
   ]
 }
diff --git a/src/trace_processor/importers/etw/etw_module_impl.cc b/src/trace_processor/importers/etw/etw_module_impl.cc
index 7c44d049e..6da884fbb 100644
--- a/src/trace_processor/importers/etw/etw_module_impl.cc
+++ b/src/trace_processor/importers/etw/etw_module_impl.cc
@@ -20,6 +20,7 @@
 #include "src/trace_processor/importers/etw/etw_tokenizer.h"
 
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 
 namespace perfetto {
 namespace trace_processor {
@@ -35,13 +36,13 @@ ModuleResult EtwModuleImpl::TokenizePacket(
     const protos::pbzero::TracePacket::Decoder& decoder,
     TraceBlobView* packet,
     int64_t /*packet_timestamp*/,
-    PacketSequenceState* seq_state,
+    RefPtr<PacketSequenceStateGeneration> seq_state,
     uint32_t field_id) {
   switch (field_id) {
     case TracePacket::kEtwEventsFieldNumber: {
       auto etw_field = decoder.etw_events();
       tokenizer_.TokenizeEtwBundle(
-          packet->slice(etw_field.data, etw_field.size), seq_state);
+          packet->slice(etw_field.data, etw_field.size), std::move(seq_state));
       return ModuleResult::Handled();
     }
   }
diff --git a/src/trace_processor/importers/etw/etw_module_impl.h b/src/trace_processor/importers/etw/etw_module_impl.h
index af05fc8aa..4a27846f4 100644
--- a/src/trace_processor/importers/etw/etw_module_impl.h
+++ b/src/trace_processor/importers/etw/etw_module_impl.h
@@ -21,6 +21,7 @@
 #include "src/trace_processor/importers/etw/etw_module.h"
 #include "src/trace_processor/importers/etw/etw_parser.h"
 #include "src/trace_processor/importers/etw/etw_tokenizer.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
@@ -38,7 +39,7 @@ class EtwModuleImpl : public EtwModule {
       const protos::pbzero::TracePacket::Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
   void ParseEtwEventData(uint32_t cpu,
diff --git a/src/trace_processor/importers/etw/etw_tokenizer.cc b/src/trace_processor/importers/etw/etw_tokenizer.cc
index 2474454b3..e3b4a747b 100644
--- a/src/trace_processor/importers/etw/etw_tokenizer.cc
+++ b/src/trace_processor/importers/etw/etw_tokenizer.cc
@@ -22,7 +22,7 @@
 #include "perfetto/ext/base/status_or.h"
 #include "perfetto/protozero/proto_decoder.h"
 #include "perfetto/protozero/proto_utils.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/sorter/trace_sorter.h"
 #include "src/trace_processor/storage/trace_storage.h"
 
@@ -41,8 +41,9 @@ using protos::pbzero::BuiltinClock;
 using protos::pbzero::EtwTraceEventBundle;
 
 PERFETTO_ALWAYS_INLINE
-base::Status EtwTokenizer::TokenizeEtwBundle(TraceBlobView bundle,
-                                             PacketSequenceState* state) {
+base::Status EtwTokenizer::TokenizeEtwBundle(
+    TraceBlobView bundle,
+    RefPtr<PacketSequenceStateGeneration> state) {
   protos::pbzero::EtwTraceEventBundle::Decoder decoder(bundle.data(),
                                                        bundle.length());
   // Cpu id can either be in the etw bundle or inside the individual
@@ -61,7 +62,7 @@ PERFETTO_ALWAYS_INLINE
 base::Status EtwTokenizer::TokenizeEtwEvent(
     std::optional<uint32_t> fallback_cpu,
     TraceBlobView event,
-    PacketSequenceState* state) {
+    RefPtr<PacketSequenceStateGeneration> state) {
   const uint8_t* data = event.data();
   const size_t length = event.length();
   ProtoDecoder decoder(data, length);
@@ -104,7 +105,7 @@ base::Status EtwTokenizer::TokenizeEtwEvent(
   }
 
   context_->sorter->PushEtwEvent(cpu, *timestamp, std::move(event),
-                                 state->current_generation());
+                                 std::move(state));
 
   return base::OkStatus();
 }
diff --git a/src/trace_processor/importers/etw/etw_tokenizer.h b/src/trace_processor/importers/etw/etw_tokenizer.h
index 6447ed2f7..028346824 100644
--- a/src/trace_processor/importers/etw/etw_tokenizer.h
+++ b/src/trace_processor/importers/etw/etw_tokenizer.h
@@ -21,6 +21,7 @@
 
 #include "perfetto/base/status.h"
 #include "perfetto/trace_processor/trace_blob_view.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/types/trace_processor_context.h"
 
@@ -29,19 +30,17 @@
 namespace perfetto {
 namespace trace_processor {
 
-class PacketSequenceState;
-
 class EtwTokenizer {
  public:
   explicit EtwTokenizer(TraceProcessorContext* context) : context_(context) {}
 
   base::Status TokenizeEtwBundle(TraceBlobView bundle,
-                                 PacketSequenceState* state);
+                                 RefPtr<PacketSequenceStateGeneration> state);
 
  private:
   base::Status TokenizeEtwEvent(std::optional<uint32_t> fallback_cpu,
                                 TraceBlobView event,
-                                PacketSequenceState* state);
+                                RefPtr<PacketSequenceStateGeneration> state);
 
   TraceProcessorContext* context_;
 };
diff --git a/src/trace_processor/importers/ftrace/ftrace_module_impl.cc b/src/trace_processor/importers/ftrace/ftrace_module_impl.cc
index bd67c644d..0858fa1ba 100644
--- a/src/trace_processor/importers/ftrace/ftrace_module_impl.cc
+++ b/src/trace_processor/importers/ftrace/ftrace_module_impl.cc
@@ -38,14 +38,14 @@ ModuleResult FtraceModuleImpl::TokenizePacket(
     const protos::pbzero::TracePacket::Decoder& decoder,
     TraceBlobView* packet,
     int64_t /*packet_timestamp*/,
-    PacketSequenceState* seq_state,
+    RefPtr<PacketSequenceStateGeneration> seq_state,
     uint32_t field_id) {
   switch (field_id) {
     case TracePacket::kFtraceEventsFieldNumber: {
       auto ftrace_field = decoder.ftrace_events();
       return tokenizer_.TokenizeFtraceBundle(
-          packet->slice(ftrace_field.data, ftrace_field.size), seq_state,
-          decoder.trusted_packet_sequence_id());
+          packet->slice(ftrace_field.data, ftrace_field.size),
+          std::move(seq_state), decoder.trusted_packet_sequence_id());
     }
     case TracePacket::kFtraceStatsFieldNumber: {
       return parser_.ParseFtraceStats(decoder.ftrace_stats(),
diff --git a/src/trace_processor/importers/ftrace/ftrace_module_impl.h b/src/trace_processor/importers/ftrace/ftrace_module_impl.h
index 658bafc84..d6df0f9dc 100644
--- a/src/trace_processor/importers/ftrace/ftrace_module_impl.h
+++ b/src/trace_processor/importers/ftrace/ftrace_module_impl.h
@@ -22,6 +22,7 @@
 #include "src/trace_processor/importers/ftrace/ftrace_module.h"
 #include "src/trace_processor/importers/ftrace/ftrace_parser.h"
 #include "src/trace_processor/importers/ftrace/ftrace_tokenizer.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 
 namespace perfetto {
@@ -37,7 +38,7 @@ class FtraceModuleImpl : public FtraceModule {
       const protos::pbzero::TracePacket::Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
   void ParseFtraceEventData(uint32_t cpu,
diff --git a/src/trace_processor/importers/ftrace/ftrace_tokenizer.cc b/src/trace_processor/importers/ftrace/ftrace_tokenizer.cc
index cf38ac29d..00656fa66 100644
--- a/src/trace_processor/importers/ftrace/ftrace_tokenizer.cc
+++ b/src/trace_processor/importers/ftrace/ftrace_tokenizer.cc
@@ -23,7 +23,6 @@
 #include "perfetto/trace_processor/basic_types.h"
 #include "src/trace_processor/importers/common/machine_tracker.h"
 #include "src/trace_processor/importers/common/metadata_tracker.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
 #include "src/trace_processor/sorter/trace_sorter.h"
 #include "src/trace_processor/storage/metadata.h"
 #include "src/trace_processor/storage/stats.h"
@@ -107,7 +106,7 @@ uint64_t TryFastParseFtraceEventId(const uint8_t* start, const uint8_t* end) {
 PERFETTO_ALWAYS_INLINE
 base::Status FtraceTokenizer::TokenizeFtraceBundle(
     TraceBlobView bundle,
-    PacketSequenceState* state,
+    RefPtr<PacketSequenceStateGeneration> state,
     uint32_t packet_sequence_id) {
   protos::pbzero::FtraceEventBundle::Decoder decoder(bundle.data(),
                                                      bundle.length());
@@ -204,10 +203,11 @@ base::Status FtraceTokenizer::TokenizeFtraceBundle(
 }
 
 PERFETTO_ALWAYS_INLINE
-void FtraceTokenizer::TokenizeFtraceEvent(uint32_t cpu,
-                                          ClockTracker::ClockId clock_id,
-                                          TraceBlobView event,
-                                          PacketSequenceState* state) {
+void FtraceTokenizer::TokenizeFtraceEvent(
+    uint32_t cpu,
+    ClockTracker::ClockId clock_id,
+    TraceBlobView event,
+    RefPtr<PacketSequenceStateGeneration> state) {
   constexpr auto kTimestampFieldNumber =
       protos::pbzero::FtraceEvent::kTimestampFieldNumber;
   constexpr auto kTimestampFieldTag = MakeTagVarInt(kTimestampFieldNumber);
@@ -281,8 +281,7 @@ void FtraceTokenizer::TokenizeFtraceEvent(uint32_t cpu,
   }
 
   context_->sorter->PushFtraceEvent(cpu, *timestamp, std::move(event),
-                                    state->current_generation(),
-                                    context_->machine_id());
+                                    std::move(state), context_->machine_id());
 }
 
 PERFETTO_ALWAYS_INLINE
@@ -427,9 +426,10 @@ void FtraceTokenizer::HandleFtraceClockSnapshot(int64_t ftrace_ts,
                                     boot_ts)});
 }
 
-void FtraceTokenizer::TokenizeFtraceGpuWorkPeriod(uint32_t cpu,
-                                                  TraceBlobView event,
-                                                  PacketSequenceState* state) {
+void FtraceTokenizer::TokenizeFtraceGpuWorkPeriod(
+    uint32_t cpu,
+    TraceBlobView event,
+    RefPtr<PacketSequenceStateGeneration> state) {
   // Special handling of valid gpu_work_period tracepoint events which contain
   // timestamp values for the GPU time period nested inside the event data.
   const uint8_t* data = event.data();
@@ -465,8 +465,7 @@ void FtraceTokenizer::TokenizeFtraceGpuWorkPeriod(uint32_t cpu,
   }
 
   context_->sorter->PushFtraceEvent(cpu, *timestamp, std::move(event),
-                                    state->current_generation(),
-                                    context_->machine_id());
+                                    std::move(state), context_->machine_id());
 }
 
 }  // namespace trace_processor
diff --git a/src/trace_processor/importers/ftrace/ftrace_tokenizer.h b/src/trace_processor/importers/ftrace/ftrace_tokenizer.h
index 915ee9c88..d53504efe 100644
--- a/src/trace_processor/importers/ftrace/ftrace_tokenizer.h
+++ b/src/trace_processor/importers/ftrace/ftrace_tokenizer.h
@@ -21,6 +21,7 @@
 
 #include "perfetto/trace_processor/trace_blob_view.h"
 #include "src/trace_processor/importers/common/clock_tracker.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/types/trace_processor_context.h"
 
@@ -29,22 +30,20 @@
 namespace perfetto {
 namespace trace_processor {
 
-class PacketSequenceState;
-
 class FtraceTokenizer {
  public:
   explicit FtraceTokenizer(TraceProcessorContext* context)
       : context_(context) {}
 
   base::Status TokenizeFtraceBundle(TraceBlobView bundle,
-                                    PacketSequenceState*,
+                                    RefPtr<PacketSequenceStateGeneration>,
                                     uint32_t packet_sequence_id);
 
  private:
   void TokenizeFtraceEvent(uint32_t cpu,
                            ClockTracker::ClockId,
                            TraceBlobView event,
-                           PacketSequenceState* state);
+                           RefPtr<PacketSequenceStateGeneration> state);
   void TokenizeFtraceCompactSched(uint32_t cpu,
                                   ClockTracker::ClockId,
                                   protozero::ConstBytes);
@@ -64,7 +63,7 @@ class FtraceTokenizer {
                                  uint32_t packet_sequence_id);
   void TokenizeFtraceGpuWorkPeriod(uint32_t cpu,
                                    TraceBlobView event,
-                                   PacketSequenceState* state);
+                                   RefPtr<PacketSequenceStateGeneration> state);
 
   void DlogWithLimit(const base::Status& status) {
     static std::atomic<uint32_t> dlog_count(0);
diff --git a/src/trace_processor/importers/proto/BUILD.gn b/src/trace_processor/importers/proto/BUILD.gn
index 9930ea558..18a617af5 100644
--- a/src/trace_processor/importers/proto/BUILD.gn
+++ b/src/trace_processor/importers/proto/BUILD.gn
@@ -72,6 +72,7 @@ source_set("minimal") {
   deps = [
     ":packet_sequence_state_generation_hdr",
     "../../../../gn:default_deps",
+    "../../../../include/perfetto/trace_processor:trace_processor",
     "../../../../protos/perfetto/common:zero",
     "../../../../protos/perfetto/config:zero",
     "../../../../protos/perfetto/trace:zero",
@@ -201,7 +202,9 @@ source_set("proto_importer_module") {
     "proto_importer_module.h",
   ]
   deps = [
+    ":packet_sequence_state_generation_hdr",
     "../../../../gn:default_deps",
+    "../../../../include/perfetto/trace_processor:trace_processor",
     "../../../base",
     "../../types",
     "../common:trace_parser_hdr",
diff --git a/src/trace_processor/importers/proto/android_camera_event_module.cc b/src/trace_processor/importers/proto/android_camera_event_module.cc
index 5c7722a06..60ec66123 100644
--- a/src/trace_processor/importers/proto/android_camera_event_module.cc
+++ b/src/trace_processor/importers/proto/android_camera_event_module.cc
@@ -17,17 +17,18 @@
 #include "src/trace_processor/importers/proto/android_camera_event_module.h"
 
 #include "perfetto/ext/base/string_utils.h"
-#include "protos/perfetto/trace/android/camera_event.pbzero.h"
-#include "protos/perfetto/trace/trace_packet.pbzero.h"
 #include "src/trace_processor/importers/common/async_track_set_tracker.h"
 #include "src/trace_processor/importers/common/machine_tracker.h"
 #include "src/trace_processor/importers/common/parser_types.h"
 #include "src/trace_processor/importers/common/slice_tracker.h"
 #include "src/trace_processor/importers/common/track_tracker.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/sorter/trace_sorter.h"
 #include "src/trace_processor/storage/trace_storage.h"
 
+#include "protos/perfetto/trace/android/camera_event.pbzero.h"
+#include "protos/perfetto/trace/trace_packet.pbzero.h"
+
 namespace perfetto {
 namespace trace_processor {
 
@@ -45,7 +46,7 @@ ModuleResult AndroidCameraEventModule::TokenizePacket(
     const protos::pbzero::TracePacket::Decoder& decoder,
     TraceBlobView* packet,
     int64_t /*packet_timestamp*/,
-    PacketSequenceState* state,
+    RefPtr<PacketSequenceStateGeneration> state,
     uint32_t field_id) {
   if (field_id != TracePacket::kAndroidCameraFrameEventFieldNumber) {
     return ModuleResult::Ignored();
@@ -55,7 +56,7 @@ ModuleResult AndroidCameraEventModule::TokenizePacket(
           decoder.android_camera_frame_event());
   context_->sorter->PushTracePacket(
       android_camera_frame_event.request_processing_started_ns(),
-      state->current_generation(), std::move(*packet), context_->machine_id());
+      std::move(state), std::move(*packet), context_->machine_id());
   return ModuleResult::Handled();
 }
 
diff --git a/src/trace_processor/importers/proto/android_camera_event_module.h b/src/trace_processor/importers/proto/android_camera_event_module.h
index e9df6c469..03f8476b8 100644
--- a/src/trace_processor/importers/proto/android_camera_event_module.h
+++ b/src/trace_processor/importers/proto/android_camera_event_module.h
@@ -23,6 +23,7 @@
 
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 #include "src/trace_processor/importers/common/parser_types.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/tables/sched_tables_py.h"
 #include "src/trace_processor/tables/slice_tables_py.h"
@@ -42,7 +43,7 @@ class AndroidCameraEventModule : public ProtoImporterModule {
       const protos::pbzero::TracePacket::Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
   void ParseTracePacketData(const protos::pbzero::TracePacket::Decoder& decoder,
diff --git a/src/trace_processor/importers/proto/android_probes_module.cc b/src/trace_processor/importers/proto/android_probes_module.cc
index 02dcfe771..31f51dc6a 100644
--- a/src/trace_processor/importers/proto/android_probes_module.cc
+++ b/src/trace_processor/importers/proto/android_probes_module.cc
@@ -23,7 +23,7 @@
 #include "src/trace_processor/importers/common/track_tracker.h"
 #include "src/trace_processor/importers/proto/android_probes_parser.h"
 #include "src/trace_processor/importers/proto/android_probes_tracker.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/sorter/trace_sorter.h"
 
 #include "protos/perfetto/common/android_energy_consumer_descriptor.pbzero.h"
@@ -104,7 +104,7 @@ ModuleResult AndroidProbesModule::TokenizePacket(
     const protos::pbzero::TracePacket_Decoder&,
     TraceBlobView* packet,
     int64_t packet_timestamp,
-    PacketSequenceState* state,
+    RefPtr<PacketSequenceStateGeneration> state,
     uint32_t field_id) {
   protos::pbzero::TracePacket::Decoder decoder(packet->data(),
                                                packet->length());
@@ -193,7 +193,7 @@ ModuleResult AndroidProbesModule::TokenizePacket(
 
     std::vector<uint8_t> vec = data_packet.SerializeAsArray();
     TraceBlob blob = TraceBlob::CopyFrom(vec.data(), vec.size());
-    context_->sorter->PushTracePacket(actual_ts, state->current_generation(),
+    context_->sorter->PushTracePacket(actual_ts, state,
                                       TraceBlobView(std::move(blob)),
                                       context_->machine_id());
   }
diff --git a/src/trace_processor/importers/proto/android_probes_module.h b/src/trace_processor/importers/proto/android_probes_module.h
index 6b3d8cf89..1f45cbb65 100644
--- a/src/trace_processor/importers/proto/android_probes_module.h
+++ b/src/trace_processor/importers/proto/android_probes_module.h
@@ -19,6 +19,7 @@
 
 #include "perfetto/base/build_config.h"
 #include "src/trace_processor/importers/proto/android_probes_parser.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 
 #include "protos/perfetto/config/trace_config.pbzero.h"
@@ -34,7 +35,7 @@ class AndroidProbesModule : public ProtoImporterModule {
   ModuleResult TokenizePacket(const protos::pbzero::TracePacket_Decoder&,
                               TraceBlobView* packet,
                               int64_t packet_timestamp,
-                              PacketSequenceState*,
+                              RefPtr<PacketSequenceStateGeneration>,
                               uint32_t field_id) override;
 
   void ParseTracePacketData(const protos::pbzero::TracePacket_Decoder& decoder,
diff --git a/src/trace_processor/importers/proto/metadata_minimal_module.cc b/src/trace_processor/importers/proto/metadata_minimal_module.cc
index 691127586..dd0127f14 100644
--- a/src/trace_processor/importers/proto/metadata_minimal_module.cc
+++ b/src/trace_processor/importers/proto/metadata_minimal_module.cc
@@ -19,6 +19,7 @@
 #include "perfetto/ext/base/base64.h"
 #include "perfetto/ext/base/string_utils.h"
 #include "src/trace_processor/importers/common/metadata_tracker.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/trace/chrome/chrome_benchmark_metadata.pbzero.h"
@@ -39,7 +40,7 @@ ModuleResult MetadataMinimalModule::TokenizePacket(
     const protos::pbzero::TracePacket::Decoder& decoder,
     TraceBlobView*,
     int64_t,
-    PacketSequenceState*,
+    RefPtr<PacketSequenceStateGeneration>,
     uint32_t field_id) {
   switch (field_id) {
     case TracePacket::kChromeMetadataFieldNumber: {
diff --git a/src/trace_processor/importers/proto/metadata_minimal_module.h b/src/trace_processor/importers/proto/metadata_minimal_module.h
index 3ffdf62f1..80c131bd8 100644
--- a/src/trace_processor/importers/proto/metadata_minimal_module.h
+++ b/src/trace_processor/importers/proto/metadata_minimal_module.h
@@ -18,6 +18,7 @@
 #define SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_METADATA_MINIMAL_MODULE_H_
 
 #include "src/trace_processor/importers/common/trace_parser.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 
 #include "src/trace_processor/storage/trace_storage.h"
@@ -36,7 +37,7 @@ class MetadataMinimalModule : public ProtoImporterModule {
       const protos::pbzero::TracePacket::Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
  private:
diff --git a/src/trace_processor/importers/proto/metadata_module.cc b/src/trace_processor/importers/proto/metadata_module.cc
index 86cc72b71..96bd01768 100644
--- a/src/trace_processor/importers/proto/metadata_module.cc
+++ b/src/trace_processor/importers/proto/metadata_module.cc
@@ -23,6 +23,7 @@
 #include "src/trace_processor/importers/common/slice_tracker.h"
 #include "src/trace_processor/importers/common/track_tracker.h"
 #include "src/trace_processor/importers/proto/config.descriptor.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/util/descriptors.h"
 #include "src/trace_processor/util/protozero_to_text.h"
 
@@ -52,7 +53,7 @@ ModuleResult MetadataModule::TokenizePacket(
     const protos::pbzero::TracePacket::Decoder& decoder,
     TraceBlobView*,
     int64_t,
-    PacketSequenceState*,
+    RefPtr<PacketSequenceStateGeneration>,
     uint32_t field_id) {
   switch (field_id) {
     case TracePacket::kUiStateFieldNumber: {
diff --git a/src/trace_processor/importers/proto/metadata_module.h b/src/trace_processor/importers/proto/metadata_module.h
index 31cbd4259..001493673 100644
--- a/src/trace_processor/importers/proto/metadata_module.h
+++ b/src/trace_processor/importers/proto/metadata_module.h
@@ -18,6 +18,7 @@
 #define SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_METADATA_MODULE_H_
 
 #include "src/trace_processor/importers/common/trace_parser.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
@@ -35,7 +36,7 @@ class MetadataModule : public ProtoImporterModule {
       const protos::pbzero::TracePacket::Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
   void ParseTracePacketData(const protos::pbzero::TracePacket::Decoder& decoder,
diff --git a/src/trace_processor/importers/proto/network_trace_module.cc b/src/trace_processor/importers/proto/network_trace_module.cc
index 01c949463..dda8f7074 100644
--- a/src/trace_processor/importers/proto/network_trace_module.cc
+++ b/src/trace_processor/importers/proto/network_trace_module.cc
@@ -21,7 +21,7 @@
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 #include "src/trace_processor/importers/common/async_track_set_tracker.h"
 #include "src/trace_processor/importers/common/slice_tracker.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/sorter/trace_sorter.h"
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/types/tcp_state.h"
@@ -78,18 +78,17 @@ ModuleResult NetworkTraceModule::TokenizePacket(
     const protos::pbzero::TracePacket::Decoder& decoder,
     TraceBlobView*,
     int64_t ts,
-    PacketSequenceState* state,
+    RefPtr<PacketSequenceStateGeneration> state,
     uint32_t field_id) {
   if (field_id != TracePacket::kNetworkPacketBundleFieldNumber) {
     return ModuleResult::Ignored();
   }
 
-  auto seq_state = state->current_generation();
   NetworkPacketBundle::Decoder evt(decoder.network_packet_bundle());
 
   ConstBytes context = evt.ctx();
   if (evt.has_iid()) {
-    auto* interned = seq_state->LookupInternedMessage<
+    auto* interned = state->LookupInternedMessage<
         protos::pbzero::InternedData::kPacketContextFieldNumber,
         protos::pbzero::NetworkPacketContext>(evt.iid());
     if (!interned) {
@@ -258,11 +257,12 @@ void NetworkTraceModule::ParseNetworkPacketBundle(int64_t ts, ConstBytes blob) {
   });
 }
 
-void NetworkTraceModule::PushPacketBufferForSort(int64_t timestamp,
-                                                 PacketSequenceState* state) {
+void NetworkTraceModule::PushPacketBufferForSort(
+    int64_t timestamp,
+    RefPtr<PacketSequenceStateGeneration> state) {
   std::vector<uint8_t> v = packet_buffer_.SerializeAsArray();
   context_->sorter->PushTracePacket(
-      timestamp, state->current_generation(),
+      timestamp, std::move(state),
       TraceBlobView(TraceBlob::CopyFrom(v.data(), v.size())));
   packet_buffer_.Reset();
 }
diff --git a/src/trace_processor/importers/proto/network_trace_module.h b/src/trace_processor/importers/proto/network_trace_module.h
index 15082c793..984ea822c 100644
--- a/src/trace_processor/importers/proto/network_trace_module.h
+++ b/src/trace_processor/importers/proto/network_trace_module.h
@@ -20,14 +20,16 @@
 #include <cstdint>
 
 #include "perfetto/protozero/scattered_heap_buffer.h"
-#include "protos/perfetto/trace/android/network_trace.pbzero.h"
-#include "protos/perfetto/trace/trace_packet.pbzero.h"
 #include "src/trace_processor/importers/common/args_tracker.h"
 #include "src/trace_processor/importers/common/parser_types.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/types/trace_processor_context.h"
 
+#include "protos/perfetto/trace/android/network_trace.pbzero.h"
+#include "protos/perfetto/trace/trace_packet.pbzero.h"
+
 namespace perfetto {
 namespace trace_processor {
 
@@ -45,7 +47,7 @@ class NetworkTraceModule : public ProtoImporterModule {
       const protos::pbzero::TracePacket::Decoder& decoder,
       TraceBlobView* packet,
       int64_t ts,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
   void ParseTracePacketData(const protos::pbzero::TracePacket::Decoder& decoder,
@@ -65,7 +67,8 @@ class NetworkTraceModule : public ProtoImporterModule {
 
   // Helper to simplify pushing a TracePacket to the sorter. The caller fills in
   // the packet buffer and uses this to push for sorting and reset the buffer.
-  void PushPacketBufferForSort(int64_t timestamp, PacketSequenceState* state);
+  void PushPacketBufferForSort(int64_t timestamp,
+                               RefPtr<PacketSequenceStateGeneration> state);
 
   TraceProcessorContext* context_;
   protozero::HeapBuffered<protos::pbzero::TracePacket> packet_buffer_;
diff --git a/src/trace_processor/importers/proto/packet_sequence_state.h b/src/trace_processor/importers/proto/packet_sequence_state.h
index 2ad0c7b3c..b50af207b 100644
--- a/src/trace_processor/importers/proto/packet_sequence_state.h
+++ b/src/trace_processor/importers/proto/packet_sequence_state.h
@@ -35,8 +35,7 @@ class PacketSequenceState {
  public:
   explicit PacketSequenceState(TraceProcessorContext* context)
       : context_(context) {
-    current_generation_.reset(
-        new PacketSequenceStateGeneration(this, generation_index_++));
+    current_generation_.reset(new PacketSequenceStateGeneration(this));
   }
 
   int64_t IncrementAndGetTrackEventTimeNs(int64_t delta_ns) {
@@ -75,8 +74,7 @@ class PacketSequenceState {
     // sequence. Add a new generation with the updated defaults but the
     // current generation's interned data state.
     current_generation_.reset(new PacketSequenceStateGeneration(
-        this, generation_index_++, current_generation_.get(),
-        std::move(defaults)));
+        this, current_generation_.get(), std::move(defaults)));
   }
 
   void SetThreadDescriptor(int32_t pid,
@@ -101,8 +99,7 @@ class PacketSequenceState {
   // Starts a new generation with clean-slate incremental state and defaults.
   void OnIncrementalStateCleared() {
     packet_loss_ = false;
-    current_generation_.reset(
-        new PacketSequenceStateGeneration(this, generation_index_++));
+    current_generation_.reset(new PacketSequenceStateGeneration(this));
   }
 
   bool IsIncrementalStateValid() const { return !packet_loss_; }
@@ -126,8 +123,6 @@ class PacketSequenceState {
  private:
   TraceProcessorContext* context_;
 
-  size_t generation_index_ = 0;
-
   // If true, incremental state on the sequence is considered invalid until we
   // see the next packet with incremental_state_cleared. We assume that we
   // missed some packets at the beginning of the trace.
diff --git a/src/trace_processor/importers/proto/packet_sequence_state_generation.cc b/src/trace_processor/importers/proto/packet_sequence_state_generation.cc
index 1e890962a..c520e9259 100644
--- a/src/trace_processor/importers/proto/packet_sequence_state_generation.cc
+++ b/src/trace_processor/importers/proto/packet_sequence_state_generation.cc
@@ -25,11 +25,9 @@ namespace trace_processor {
 
 PacketSequenceStateGeneration::PacketSequenceStateGeneration(
     PacketSequenceState* state,
-    size_t generation_index,
     PacketSequenceStateGeneration* prev_gen,
     TraceBlobView defaults)
     : state_(state),
-      generation_index_(generation_index),
       interned_data_(prev_gen->interned_data_),
       trace_packet_defaults_(InternedMessageView(std::move(defaults))),
       trackers_(prev_gen->trackers_) {
@@ -43,6 +41,16 @@ PacketSequenceStateGeneration::PacketSequenceStateGeneration(
 PacketSequenceStateGeneration::InternedDataTracker::~InternedDataTracker() =
     default;
 
+bool PacketSequenceStateGeneration::pid_and_tid_valid() const {
+  return state_->pid_and_tid_valid();
+}
+int32_t PacketSequenceStateGeneration::pid() const {
+  return state_->pid();
+}
+int32_t PacketSequenceStateGeneration::tid() const {
+  return state_->tid();
+}
+
 TraceProcessorContext* PacketSequenceStateGeneration::GetContext() const {
   return state_->context();
 }
@@ -95,5 +103,35 @@ InternedMessageView* PacketSequenceStateGeneration::GetInternedMessageView(
   return nullptr;
 }
 
+int64_t PacketSequenceStateGeneration::IncrementAndGetTrackEventTimeNs(
+    int64_t delta_ns) {
+  return state_->IncrementAndGetTrackEventTimeNs(delta_ns);
+}
+int64_t PacketSequenceStateGeneration::IncrementAndGetTrackEventThreadTimeNs(
+    int64_t delta_ns) {
+  return state_->IncrementAndGetTrackEventThreadTimeNs(delta_ns);
+}
+int64_t
+PacketSequenceStateGeneration::IncrementAndGetTrackEventThreadInstructionCount(
+    int64_t delta) {
+  return state_->IncrementAndGetTrackEventThreadInstructionCount(delta);
+}
+bool PacketSequenceStateGeneration::track_event_timestamps_valid() const {
+  return state_->track_event_timestamps_valid();
+}
+void PacketSequenceStateGeneration::SetThreadDescriptor(
+    int32_t pid,
+    int32_t tid,
+    int64_t timestamp_ns,
+    int64_t thread_timestamp_ns,
+    int64_t thread_instruction_count) {
+  state_->SetThreadDescriptor(pid, tid, timestamp_ns, thread_timestamp_ns,
+                              thread_instruction_count);
+}
+
+bool PacketSequenceStateGeneration::IsIncrementalStateValid() const {
+  return state_->IsIncrementalStateValid();
+}
+
 }  // namespace trace_processor
 }  // namespace perfetto
diff --git a/src/trace_processor/importers/proto/packet_sequence_state_generation.h b/src/trace_processor/importers/proto/packet_sequence_state_generation.h
index f31030924..6b60b596d 100644
--- a/src/trace_processor/importers/proto/packet_sequence_state_generation.h
+++ b/src/trace_processor/importers/proto/packet_sequence_state_generation.h
@@ -19,6 +19,7 @@
 
 #include <array>
 #include <cstddef>
+#include <cstdint>
 #include <memory>
 #include <optional>
 #include <tuple>
@@ -84,7 +85,9 @@ class PacketSequenceStateGeneration : public RefCounted {
       return generation_->GetOrCreate<T>();
     }
 
-    PacketSequenceState* state() const { return generation_->state(); }
+    bool pid_and_tid_valid() const { return generation_->pid_and_tid_valid(); }
+    int32_t pid() const { return generation_->pid(); }
+    int32_t tid() const { return generation_->tid(); }
 
    private:
     friend PacketSequenceStateGeneration;
@@ -103,6 +106,10 @@ class PacketSequenceStateGeneration : public RefCounted {
     PacketSequenceStateGeneration* generation_ = nullptr;
   };
 
+  bool pid_and_tid_valid() const;
+  int32_t pid() const;
+  int32_t tid() const;
+
   // Returns |nullptr| if the message with the given |iid| was not found (also
   // records a stat in this case).
   template <uint32_t FieldId, typename MessageType>
@@ -141,15 +148,28 @@ class PacketSequenceStateGeneration : public RefCounted {
     return nullptr;
   }
 
-  PacketSequenceState* state() const { return state_; }
-  size_t generation_index() const { return generation_index_; }
-
   // Extension point for custom sequence state. To add new per sequence state
   // just subclass ´PacketSequenceStateGeneration´ and get your sequence bound
   // instance by calling this method.
   template <typename T>
   std::remove_cv_t<T>* GetOrCreate();
 
+  // TODO(carlscab): All this should be tracked in a dedicated class
+  // TrackEventSequenceState or something attached to the "incremental state".
+  int64_t IncrementAndGetTrackEventTimeNs(int64_t delta_ns);
+  int64_t IncrementAndGetTrackEventThreadTimeNs(int64_t delta_ns);
+  int64_t IncrementAndGetTrackEventThreadInstructionCount(int64_t delta);
+  bool track_event_timestamps_valid() const;
+  void SetThreadDescriptor(int32_t pid,
+                           int32_t tid,
+                           int64_t timestamp_ns,
+                           int64_t thread_timestamp_ns,
+                           int64_t thread_instruction_count);
+
+  // TODO(carlscab): Nobody other than `ProtoTraceReader` should care about
+  // this. Remove.
+  bool IsIncrementalStateValid() const;
+
  private:
   friend class PacketSequenceState;
 
@@ -175,12 +195,10 @@ class PacketSequenceStateGeneration : public RefCounted {
     }
   }
 
-  PacketSequenceStateGeneration(PacketSequenceState* state,
-                                size_t generation_index)
-      : state_(state), generation_index_(generation_index) {}
+  explicit PacketSequenceStateGeneration(PacketSequenceState* state)
+      : state_(state) {}
 
   PacketSequenceStateGeneration(PacketSequenceState* state,
-                                size_t generation_index,
                                 PacketSequenceStateGeneration* prev_gen,
                                 TraceBlobView defaults);
 
@@ -194,8 +212,9 @@ class PacketSequenceStateGeneration : public RefCounted {
     trace_packet_defaults_ = InternedMessageView(std::move(defaults));
   }
 
+  // TODO(carlscab): This is dangerous given that PacketSequenceStateGeneration
+  // is refcounted and PacketSequenceState is not.
   PacketSequenceState* state_;
-  size_t generation_index_;
   InternedFieldMap interned_data_;
   std::optional<InternedMessageView> trace_packet_defaults_;
   std::array<RefPtr<InternedDataTracker>,
diff --git a/src/trace_processor/importers/proto/profile_module.cc b/src/trace_processor/importers/proto/profile_module.cc
index c194107eb..6064d33f6 100644
--- a/src/trace_processor/importers/proto/profile_module.cc
+++ b/src/trace_processor/importers/proto/profile_module.cc
@@ -27,7 +27,7 @@
 #include "src/trace_processor/importers/common/mapping_tracker.h"
 #include "src/trace_processor/importers/common/process_tracker.h"
 #include "src/trace_processor/importers/common/stack_profile_tracker.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/perf_sample_tracker.h"
 #include "src/trace_processor/importers/proto/profile_packet_sequence_state.h"
 #include "src/trace_processor/importers/proto/profile_packet_utils.h"
@@ -66,14 +66,15 @@ ProfileModule::ProfileModule(TraceProcessorContext* context)
 
 ProfileModule::~ProfileModule() = default;
 
-ModuleResult ProfileModule::TokenizePacket(const TracePacket::Decoder& decoder,
-                                           TraceBlobView* packet,
-                                           int64_t /*packet_timestamp*/,
-                                           PacketSequenceState* state,
-                                           uint32_t field_id) {
+ModuleResult ProfileModule::TokenizePacket(
+    const TracePacket::Decoder& decoder,
+    TraceBlobView* packet,
+    int64_t /*packet_timestamp*/,
+    RefPtr<PacketSequenceStateGeneration> state,
+    uint32_t field_id) {
   switch (field_id) {
     case TracePacket::kStreamingProfilePacketFieldNumber:
-      return TokenizeStreamingProfilePacket(state, packet,
+      return TokenizeStreamingProfilePacket(std::move(state), packet,
                                             decoder.streaming_profile_packet());
   }
   return ModuleResult::Ignored();
@@ -93,7 +94,7 @@ void ProfileModule::ParseTracePacketData(
       ParsePerfSample(ts, data.sequence_state.get(), decoder);
       return;
     case TracePacket::kProfilePacketFieldNumber:
-      ParseProfilePacket(ts, data.sequence_state->state(),
+      ParseProfilePacket(ts, data.sequence_state.get(),
                          decoder.profile_packet());
       return;
     case TracePacket::kModuleSymbolsFieldNumber:
@@ -111,7 +112,7 @@ void ProfileModule::ParseTracePacketData(
 }
 
 ModuleResult ProfileModule::TokenizeStreamingProfilePacket(
-    PacketSequenceState* sequence_state,
+    RefPtr<PacketSequenceStateGeneration> sequence_state,
     TraceBlobView* packet,
     ConstBytes streaming_profile_packet) {
   protos::pbzero::StreamingProfilePacket::Decoder decoder(
@@ -139,8 +140,7 @@ ModuleResult ProfileModule::TokenizeStreamingProfilePacket(
     sequence_state->IncrementAndGetTrackEventTimeNs(*timestamp_it * 1000);
   }
 
-  context_->sorter->PushTracePacket(packet_ts,
-                                    sequence_state->current_generation(),
+  context_->sorter->PushTracePacket(packet_ts, std::move(sequence_state),
                                     std::move(*packet), context_->machine_id());
   return ModuleResult::Handled();
 }
@@ -157,8 +157,8 @@ void ProfileModule::ParseStreamingProfilePacket(
   StackProfileSequenceState& stack_profile_sequence_state =
       *sequence_state->GetOrCreate<StackProfileSequenceState>();
 
-  uint32_t pid = static_cast<uint32_t>(sequence_state->state()->pid());
-  uint32_t tid = static_cast<uint32_t>(sequence_state->state()->tid());
+  uint32_t pid = static_cast<uint32_t>(sequence_state->pid());
+  uint32_t tid = static_cast<uint32_t>(sequence_state->tid());
   const UniqueTid utid = procs->UpdateThread(tid, pid);
   const UniquePid upid = procs->GetOrCreateProcess(pid);
 
@@ -301,12 +301,12 @@ void ProfileModule::ParsePerfSample(
   context_->storage->mutable_perf_sample_table()->Insert(sample_row);
 }
 
-void ProfileModule::ParseProfilePacket(int64_t ts,
-                                       PacketSequenceState* sequence_state,
-                                       ConstBytes blob) {
+void ProfileModule::ParseProfilePacket(
+    int64_t ts,
+    PacketSequenceStateGeneration* sequence_state,
+    ConstBytes blob) {
   ProfilePacketSequenceState& profile_packet_sequence_state =
-      *sequence_state->current_generation()
-           ->GetOrCreate<ProfilePacketSequenceState>();
+      *sequence_state->GetOrCreate<ProfilePacketSequenceState>();
   protos::pbzero::ProfilePacket::Decoder packet(blob.data, blob.size);
   profile_packet_sequence_state.SetProfilePacketIndex(packet.index());
 
diff --git a/src/trace_processor/importers/proto/profile_module.h b/src/trace_processor/importers/proto/profile_module.h
index 883fb7aa2..cabc3bfef 100644
--- a/src/trace_processor/importers/proto/profile_module.h
+++ b/src/trace_processor/importers/proto/profile_module.h
@@ -18,8 +18,8 @@
 #define SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_PROFILE_MODULE_H_
 
 #include "perfetto/protozero/field.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
-#include "src/trace_processor/importers/proto/proto_incremental_state.h"
 
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 
@@ -37,7 +37,7 @@ class ProfileModule : public ProtoImporterModule {
       const protos::pbzero::TracePacket::Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
   void ParseTracePacketData(const protos::pbzero::TracePacket::Decoder& decoder,
@@ -50,7 +50,7 @@ class ProfileModule : public ProtoImporterModule {
  private:
   // chrome stack sampling:
   ModuleResult TokenizeStreamingProfilePacket(
-      PacketSequenceState*,
+      RefPtr<PacketSequenceStateGeneration>,
       TraceBlobView* packet,
       protozero::ConstBytes streaming_profile_packet);
   void ParseStreamingProfilePacket(
@@ -65,7 +65,7 @@ class ProfileModule : public ProtoImporterModule {
 
   // heap profiling:
   void ParseProfilePacket(int64_t ts,
-                          PacketSequenceState*,
+                          PacketSequenceStateGeneration*,
                           protozero::ConstBytes);
   void ParseDeobfuscationMapping(int64_t ts,
                                  PacketSequenceStateGeneration*,
diff --git a/src/trace_processor/importers/proto/profile_packet_sequence_state.cc b/src/trace_processor/importers/proto/profile_packet_sequence_state.cc
index d90976825..b5cdb704b 100644
--- a/src/trace_processor/importers/proto/profile_packet_sequence_state.cc
+++ b/src/trace_processor/importers/proto/profile_packet_sequence_state.cc
@@ -23,7 +23,6 @@
 #include "src/trace_processor/importers/common/mapping_tracker.h"
 #include "src/trace_processor/importers/common/process_tracker.h"
 #include "src/trace_processor/importers/common/stack_profile_tracker.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
 #include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/profile_packet_utils.h"
 #include "src/trace_processor/importers/proto/stack_profile_sequence_state.h"
diff --git a/src/trace_processor/importers/proto/profile_packet_utils.h b/src/trace_processor/importers/proto/profile_packet_utils.h
index a9ea0e68b..68e4e50ad 100644
--- a/src/trace_processor/importers/proto/profile_packet_utils.h
+++ b/src/trace_processor/importers/proto/profile_packet_utils.h
@@ -19,7 +19,6 @@
 #include <optional>
 
 #include "perfetto/ext/base/string_view.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
 #include "src/trace_processor/importers/proto/profile_packet_sequence_state.h"
 
 #include "protos/perfetto/trace/interned_data/interned_data.pbzero.h"
diff --git a/src/trace_processor/importers/proto/proto_importer_module.cc b/src/trace_processor/importers/proto/proto_importer_module.cc
index c79c05dd5..fbfa96037 100644
--- a/src/trace_processor/importers/proto/proto_importer_module.cc
+++ b/src/trace_processor/importers/proto/proto_importer_module.cc
@@ -16,6 +16,8 @@
 
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 
+#include "perfetto/trace_processor/ref_counted.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
@@ -29,7 +31,7 @@ ModuleResult ProtoImporterModule::TokenizePacket(
     const protos::pbzero::TracePacket_Decoder&,
     TraceBlobView* /*packet*/,
     int64_t /*packet_timestamp*/,
-    PacketSequenceState*,
+    RefPtr<PacketSequenceStateGeneration> /*sequence_state*/,
     uint32_t /*field_id*/) {
   return ModuleResult::Ignored();
 }
diff --git a/src/trace_processor/importers/proto/proto_importer_module.h b/src/trace_processor/importers/proto/proto_importer_module.h
index cc2c3a2b3..d770111a5 100644
--- a/src/trace_processor/importers/proto/proto_importer_module.h
+++ b/src/trace_processor/importers/proto/proto_importer_module.h
@@ -21,6 +21,7 @@
 
 #include "perfetto/base/status.h"
 #include "src/trace_processor/importers/common/trace_parser.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 
 namespace perfetto {
 
@@ -33,7 +34,6 @@ class TracePacket_Decoder;
 
 namespace trace_processor {
 
-class PacketSequenceState;
 class TraceBlobView;
 class TraceProcessorContext;
 
@@ -108,7 +108,7 @@ class ProtoImporterModule {
       const protos::pbzero::TracePacket_Decoder&,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState*,
+      RefPtr<PacketSequenceStateGeneration> sequence_state,
       uint32_t field_id);
 
   // Called by ProtoTraceReader during the tokenization stage i.e. before
diff --git a/src/trace_processor/importers/proto/proto_trace_parser_impl.cc b/src/trace_processor/importers/proto/proto_trace_parser_impl.cc
index bbd38c3f0..9d318a1c8 100644
--- a/src/trace_processor/importers/proto/proto_trace_parser_impl.cc
+++ b/src/trace_processor/importers/proto/proto_trace_parser_impl.cc
@@ -38,7 +38,6 @@
 #include "src/trace_processor/importers/common/track_tracker.h"
 #include "src/trace_processor/importers/etw/etw_module.h"
 #include "src/trace_processor/importers/ftrace/ftrace_module.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
 #include "src/trace_processor/importers/proto/track_event_module.h"
 #include "src/trace_processor/storage/metadata.h"
 #include "src/trace_processor/storage/stats.h"
diff --git a/src/trace_processor/importers/proto/proto_trace_reader.cc b/src/trace_processor/importers/proto/proto_trace_reader.cc
index 972acf358..e139cbe0e 100644
--- a/src/trace_processor/importers/proto/proto_trace_reader.cc
+++ b/src/trace_processor/importers/proto/proto_trace_reader.cc
@@ -236,13 +236,13 @@ util::Status ProtoTraceReader::ParsePacket(TraceBlobView packet) {
       for (ProtoImporterModule* global_module :
            context_->modules_for_all_fields) {
         ModuleResult res = global_module->TokenizePacket(
-            decoder, &packet, timestamp, state, field_id);
+            decoder, &packet, timestamp, state->current_generation(), field_id);
         if (!res.ignored())
           return res.ToStatus();
       }
       for (ProtoImporterModule* module : modules[field_id]) {
-        ModuleResult res = module->TokenizePacket(decoder, &packet, timestamp,
-                                                  state, field_id);
+        ModuleResult res = module->TokenizePacket(
+            decoder, &packet, timestamp, state->current_generation(), field_id);
         if (!res.ignored())
           return res.ToStatus();
       }
diff --git a/src/trace_processor/importers/proto/stack_profile_sequence_state.cc b/src/trace_processor/importers/proto/stack_profile_sequence_state.cc
index 2688d26dd..3d133094d 100644
--- a/src/trace_processor/importers/proto/stack_profile_sequence_state.cc
+++ b/src/trace_processor/importers/proto/stack_profile_sequence_state.cc
@@ -63,10 +63,10 @@ StackProfileSequenceState::~StackProfileSequenceState() = default;
 
 VirtualMemoryMapping* StackProfileSequenceState::FindOrInsertMapping(
     uint64_t iid) {
-  if (state()->pid_and_tid_valid()) {
+  if (pid_and_tid_valid()) {
     return FindOrInsertMappingImpl(
         context_->process_tracker->GetOrCreateProcess(
-            static_cast<uint32_t>(state()->pid())),
+            static_cast<uint32_t>(pid())),
         iid);
   }
 
diff --git a/src/trace_processor/importers/proto/statsd_module.cc b/src/trace_processor/importers/proto/statsd_module.cc
index 2f9d3cd0c..e925b5a63 100644
--- a/src/trace_processor/importers/proto/statsd_module.cc
+++ b/src/trace_processor/importers/proto/statsd_module.cc
@@ -23,7 +23,7 @@
 #include "src/trace_processor/importers/common/machine_tracker.h"
 #include "src/trace_processor/importers/common/slice_tracker.h"
 #include "src/trace_processor/importers/common/track_tracker.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/sorter/trace_sorter.h"
 #include "src/trace_processor/storage/stats.h"
 #include "src/trace_processor/storage/trace_storage.h"
@@ -216,11 +216,12 @@ StatsdModule::StatsdModule(TraceProcessorContext* context)
 
 StatsdModule::~StatsdModule() = default;
 
-ModuleResult StatsdModule::TokenizePacket(const TracePacket::Decoder& decoder,
-                                          TraceBlobView* /*packet*/,
-                                          int64_t packet_timestamp,
-                                          PacketSequenceState* state,
-                                          uint32_t field_id) {
+ModuleResult StatsdModule::TokenizePacket(
+    const TracePacket::Decoder& decoder,
+    TraceBlobView* /*packet*/,
+    int64_t packet_timestamp,
+    RefPtr<PacketSequenceStateGeneration> state,
+    uint32_t field_id) {
   if (field_id != TracePacket::kStatsdAtomFieldNumber) {
     return ModuleResult::Ignored();
   }
@@ -246,9 +247,9 @@ ModuleResult StatsdModule::TokenizePacket(const TracePacket::Decoder& decoder,
     std::vector<uint8_t> vec = forged.SerializeAsArray();
     TraceBlob blob = TraceBlob::CopyFrom(vec.data(), vec.size());
 
-    context_->sorter->PushTracePacket(
-        atom_timestamp, state->current_generation(),
-        TraceBlobView(std::move(blob)), context_->machine_id());
+    context_->sorter->PushTracePacket(atom_timestamp, state,
+                                      TraceBlobView(std::move(blob)),
+                                      context_->machine_id());
   }
 
   return ModuleResult::Handled();
diff --git a/src/trace_processor/importers/proto/statsd_module.h b/src/trace_processor/importers/proto/statsd_module.h
index 07a524d55..bb084b8c2 100644
--- a/src/trace_processor/importers/proto/statsd_module.h
+++ b/src/trace_processor/importers/proto/statsd_module.h
@@ -24,6 +24,7 @@
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 #include "src/trace_processor/importers/common/async_track_set_tracker.h"
 #include "src/trace_processor/importers/common/trace_parser.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/tables/sched_tables_py.h"
@@ -66,7 +67,7 @@ class StatsdModule : public ProtoImporterModule {
   ModuleResult TokenizePacket(const protos::pbzero::TracePacket::Decoder&,
                               TraceBlobView* packet,
                               int64_t packet_timestamp,
-                              PacketSequenceState* state,
+                              RefPtr<PacketSequenceStateGeneration> state,
                               uint32_t field_id) override;
 
   void ParseTracePacketData(const protos::pbzero::TracePacket::Decoder& decoder,
diff --git a/src/trace_processor/importers/proto/system_probes_module.cc b/src/trace_processor/importers/proto/system_probes_module.cc
index b451f966a..633a0bc8b 100644
--- a/src/trace_processor/importers/proto/system_probes_module.cc
+++ b/src/trace_processor/importers/proto/system_probes_module.cc
@@ -16,6 +16,7 @@
 
 #include "src/trace_processor/importers/proto/system_probes_module.h"
 #include "perfetto/base/build_config.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/system_probes_parser.h"
 
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
@@ -38,7 +39,7 @@ ModuleResult SystemProbesModule::TokenizePacket(
     const protos::pbzero::TracePacket::Decoder& decoder,
     TraceBlobView*,
     int64_t,
-    PacketSequenceState*,
+    RefPtr<PacketSequenceStateGeneration>,
     uint32_t field_id) {
   switch (field_id) {
     case TracePacket::kSystemInfoFieldNumber:
diff --git a/src/trace_processor/importers/proto/system_probes_module.h b/src/trace_processor/importers/proto/system_probes_module.h
index b781ecce6..96909a54e 100644
--- a/src/trace_processor/importers/proto/system_probes_module.h
+++ b/src/trace_processor/importers/proto/system_probes_module.h
@@ -19,6 +19,7 @@
 
 #include "perfetto/base/build_config.h"
 #include "src/trace_processor/importers/common/trace_parser.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/importers/proto/system_probes_parser.h"
 
@@ -34,7 +35,7 @@ class SystemProbesModule : public ProtoImporterModule {
   ModuleResult TokenizePacket(const protos::pbzero::TracePacket::Decoder&,
                               TraceBlobView* packet,
                               int64_t packet_timestamp,
-                              PacketSequenceState*,
+                              RefPtr<PacketSequenceStateGeneration>,
                               uint32_t field_id) override;
 
   void ParseTracePacketData(const protos::pbzero::TracePacket::Decoder& decoder,
diff --git a/src/trace_processor/importers/proto/track_event_module.cc b/src/trace_processor/importers/proto/track_event_module.cc
index 719704615..aee720c55 100644
--- a/src/trace_processor/importers/proto/track_event_module.cc
+++ b/src/trace_processor/importers/proto/track_event_module.cc
@@ -19,6 +19,7 @@
 #include "perfetto/base/logging.h"
 #include "perfetto/ext/base/string_utils.h"
 #include "src/trace_processor/importers/common/track_tracker.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/track_event_tracker.h"
 #include "src/trace_processor/types/trace_processor_context.h"
 
@@ -48,22 +49,23 @@ ModuleResult TrackEventModule::TokenizePacket(
     const TracePacket::Decoder& decoder,
     TraceBlobView* packet,
     int64_t packet_timestamp,
-    PacketSequenceState* state,
+    RefPtr<PacketSequenceStateGeneration> state,
     uint32_t field_id) {
   switch (field_id) {
     case TracePacket::kTrackEventRangeOfInterestFieldNumber:
-      return tokenizer_.TokenizeRangeOfInterestPacket(state, decoder,
+      return tokenizer_.TokenizeRangeOfInterestPacket(std::move(state), decoder,
                                                       packet_timestamp);
     case TracePacket::kTrackDescriptorFieldNumber:
-      return tokenizer_.TokenizeTrackDescriptorPacket(state, decoder,
+      return tokenizer_.TokenizeTrackDescriptorPacket(std::move(state), decoder,
                                                       packet_timestamp);
     case TracePacket::kTrackEventFieldNumber:
-      tokenizer_.TokenizeTrackEventPacket(state, decoder, packet,
+      tokenizer_.TokenizeTrackEventPacket(std::move(state), decoder, packet,
                                           packet_timestamp);
       return ModuleResult::Handled();
     case TracePacket::kThreadDescriptorFieldNumber:
       // TODO(eseckler): Remove once Chrome has switched to TrackDescriptors.
-      return tokenizer_.TokenizeThreadDescriptorPacket(state, decoder);
+      return tokenizer_.TokenizeThreadDescriptorPacket(std::move(state),
+                                                       decoder);
   }
   return ModuleResult::Ignored();
 }
diff --git a/src/trace_processor/importers/proto/track_event_module.h b/src/trace_processor/importers/proto/track_event_module.h
index 2b68a88b2..4a7ae69e8 100644
--- a/src/trace_processor/importers/proto/track_event_module.h
+++ b/src/trace_processor/importers/proto/track_event_module.h
@@ -17,6 +17,8 @@
 #ifndef SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_TRACK_EVENT_MODULE_H_
 #define SRC_TRACE_PROCESSOR_IMPORTERS_PROTO_TRACK_EVENT_MODULE_H_
 
+#include "perfetto/trace_processor/ref_counted.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/importers/proto/track_event_parser.h"
 #include "src/trace_processor/importers/proto/track_event_tokenizer.h"
@@ -36,7 +38,7 @@ class TrackEventModule : public ProtoImporterModule {
       const protos::pbzero::TracePacket::Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
   void OnIncrementalStateCleared(uint32_t) override;
diff --git a/src/trace_processor/importers/proto/track_event_parser.cc b/src/trace_processor/importers/proto/track_event_parser.cc
index cb8084474..09ef47433 100644
--- a/src/trace_processor/importers/proto/track_event_parser.cc
+++ b/src/trace_processor/importers/proto/track_event_parser.cc
@@ -33,7 +33,6 @@
 #include "src/trace_processor/importers/common/track_tracker.h"
 #include "src/trace_processor/importers/json/json_utils.h"
 #include "src/trace_processor/importers/proto/packet_analyzer.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
 #include "src/trace_processor/importers/proto/profile_packet_utils.h"
 #include "src/trace_processor/importers/proto/stack_profile_sequence_state.h"
 #include "src/trace_processor/importers/proto/track_event_tracker.h"
@@ -477,11 +476,9 @@ class TrackEventParser::EventImporter {
             storage_->process_track_table().id().IndexOf(track_id_);
         if (process_track_row) {
           upid_ = storage_->process_track_table().upid()[*process_track_row];
-          if (sequence_state_->state()->pid_and_tid_valid()) {
-            uint32_t pid =
-                static_cast<uint32_t>(sequence_state_->state()->pid());
-            uint32_t tid =
-                static_cast<uint32_t>(sequence_state_->state()->tid());
+          if (sequence_state_->pid_and_tid_valid()) {
+            uint32_t pid = static_cast<uint32_t>(sequence_state_->pid());
+            uint32_t tid = static_cast<uint32_t>(sequence_state_->tid());
             UniqueTid utid_candidate = procs->UpdateThread(tid, pid);
             if (storage_->thread_table().upid()[utid_candidate] == upid_)
               legacy_passthrough_utid_ = utid_candidate;
@@ -495,17 +492,15 @@ class TrackEventParser::EventImporter {
               tracks->mutable_name()->Set(*track_index, name_id_);
           }
 
-          if (sequence_state_->state()->pid_and_tid_valid()) {
-            uint32_t pid =
-                static_cast<uint32_t>(sequence_state_->state()->pid());
-            uint32_t tid =
-                static_cast<uint32_t>(sequence_state_->state()->tid());
+          if (sequence_state_->pid_and_tid_valid()) {
+            uint32_t pid = static_cast<uint32_t>(sequence_state_->pid());
+            uint32_t tid = static_cast<uint32_t>(sequence_state_->tid());
             legacy_passthrough_utid_ = procs->UpdateThread(tid, pid);
           }
         }
       }
     } else {
-      bool pid_tid_state_valid = sequence_state_->state()->pid_and_tid_valid();
+      bool pid_tid_state_valid = sequence_state_->pid_and_tid_valid();
 
       // We have a 0-value |track_uuid|. Nevertheless, we should only fall back
       // if we have either no |track_uuid| specified at all or |track_uuid| was
@@ -524,8 +519,8 @@ class TrackEventParser::EventImporter {
           legacy_event_.has_tid_override() && pid_tid_state_valid;
 
       if (fallback_to_legacy_pid_tid_tracks) {
-        uint32_t pid = static_cast<uint32_t>(sequence_state_->state()->pid());
-        uint32_t tid = static_cast<uint32_t>(sequence_state_->state()->tid());
+        uint32_t pid = static_cast<uint32_t>(sequence_state_->pid());
+        uint32_t tid = static_cast<uint32_t>(sequence_state_->tid());
         if (legacy_event_.has_pid_override()) {
           pid = static_cast<uint32_t>(legacy_event_.pid_override());
           tid = static_cast<uint32_t>(-1);
diff --git a/src/trace_processor/importers/proto/track_event_tokenizer.cc b/src/trace_processor/importers/proto/track_event_tokenizer.cc
index 7f2074280..d9ed33607 100644
--- a/src/trace_processor/importers/proto/track_event_tokenizer.cc
+++ b/src/trace_processor/importers/proto/track_event_tokenizer.cc
@@ -17,13 +17,14 @@
 #include "src/trace_processor/importers/proto/track_event_tokenizer.h"
 
 #include "perfetto/base/logging.h"
+#include "perfetto/trace_processor/ref_counted.h"
 #include "perfetto/trace_processor/trace_blob_view.h"
 #include "src/trace_processor/importers/common/clock_tracker.h"
 #include "src/trace_processor/importers/common/machine_tracker.h"
 #include "src/trace_processor/importers/common/metadata_tracker.h"
 #include "src/trace_processor/importers/common/process_tracker.h"
 #include "src/trace_processor/importers/common/track_tracker.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_trace_reader.h"
 #include "src/trace_processor/importers/proto/track_event_tracker.h"
 #include "src/trace_processor/sorter/trace_sorter.h"
@@ -58,7 +59,7 @@ TrackEventTokenizer::TrackEventTokenizer(TraceProcessorContext* context,
           context_->storage->InternString("thread_instruction_count")) {}
 
 ModuleResult TrackEventTokenizer::TokenizeRangeOfInterestPacket(
-    PacketSequenceState* /*state*/,
+    RefPtr<PacketSequenceStateGeneration> /*state*/,
     const protos::pbzero::TracePacket::Decoder& packet,
     int64_t /*packet_timestamp*/) {
   protos::pbzero::TrackEventRangeOfInterest::Decoder range_of_interest(
@@ -75,7 +76,7 @@ ModuleResult TrackEventTokenizer::TokenizeRangeOfInterestPacket(
 }
 
 ModuleResult TrackEventTokenizer::TokenizeTrackDescriptorPacket(
-    PacketSequenceState* state,
+    RefPtr<PacketSequenceStateGeneration> state,
     const protos::pbzero::TracePacket::Decoder& packet,
     int64_t packet_timestamp) {
   auto track_descriptor_field = packet.track_descriptor();
@@ -109,7 +110,7 @@ ModuleResult TrackEventTokenizer::TokenizeTrackDescriptorPacket(
     }
 
     if (state->IsIncrementalStateValid()) {
-      TokenizeThreadDescriptor(state, thread);
+      TokenizeThreadDescriptor(*state, thread);
     }
 
     track_event_tracker_->ReserveDescriptorThreadTrack(
@@ -181,7 +182,7 @@ ModuleResult TrackEventTokenizer::TokenizeTrackDescriptorPacket(
 }
 
 ModuleResult TrackEventTokenizer::TokenizeThreadDescriptorPacket(
-    PacketSequenceState* state,
+    RefPtr<PacketSequenceStateGeneration> state,
     const protos::pbzero::TracePacket::Decoder& packet) {
   if (PERFETTO_UNLIKELY(!packet.has_trusted_packet_sequence_id())) {
     PERFETTO_ELOG("ThreadDescriptor packet without trusted_packet_sequence_id");
@@ -200,25 +201,25 @@ ModuleResult TrackEventTokenizer::TokenizeThreadDescriptorPacket(
   }
 
   protos::pbzero::ThreadDescriptor::Decoder thread(packet.thread_descriptor());
-  TokenizeThreadDescriptor(state, thread);
+  TokenizeThreadDescriptor(*state, thread);
 
   // Let ProtoTraceReader forward the packet to the parser.
   return ModuleResult::Ignored();
 }
 
 void TrackEventTokenizer::TokenizeThreadDescriptor(
-    PacketSequenceState* state,
+    PacketSequenceStateGeneration& state,
     const protos::pbzero::ThreadDescriptor::Decoder& thread) {
   // TODO(eseckler): Remove support for legacy thread descriptor-based default
   // tracks and delta timestamps.
-  state->SetThreadDescriptor(thread.pid(), thread.tid(),
-                             thread.reference_timestamp_us() * 1000,
-                             thread.reference_thread_time_us() * 1000,
-                             thread.reference_thread_instruction_count());
+  state.SetThreadDescriptor(thread.pid(), thread.tid(),
+                            thread.reference_timestamp_us() * 1000,
+                            thread.reference_thread_time_us() * 1000,
+                            thread.reference_thread_instruction_count());
 }
 
 void TrackEventTokenizer::TokenizeTrackEventPacket(
-    PacketSequenceState* state,
+    RefPtr<PacketSequenceStateGeneration> state,
     const protos::pbzero::TracePacket::Decoder& packet,
     TraceBlobView* packet_blob,
     int64_t packet_timestamp) {
@@ -232,10 +233,10 @@ void TrackEventTokenizer::TokenizeTrackEventPacket(
   protos::pbzero::TrackEvent::Decoder event(field.data, field.size);
 
   protos::pbzero::TrackEventDefaults::Decoder* defaults =
-      state->current_generation()->GetTrackEventDefaults();
+      state->GetTrackEventDefaults();
 
   int64_t timestamp;
-  TrackEventData data(std::move(*packet_blob), state->current_generation());
+  TrackEventData data(std::move(*packet_blob), state);
 
   // TODO(eseckler): Remove handling of timestamps relative to ThreadDescriptors
   // once all producers have switched to clock-domain timestamps (e.g.
diff --git a/src/trace_processor/importers/proto/track_event_tokenizer.h b/src/trace_processor/importers/proto/track_event_tokenizer.h
index 010bc87dd..6ce334713 100644
--- a/src/trace_processor/importers/proto/track_event_tokenizer.h
+++ b/src/trace_processor/importers/proto/track_event_tokenizer.h
@@ -20,6 +20,7 @@
 #include <stdint.h>
 
 #include "perfetto/protozero/proto_decoder.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/storage/trace_storage.h"
 
@@ -36,7 +37,6 @@ class TracePacket_Decoder;
 
 namespace trace_processor {
 
-class PacketSequenceState;
 class TraceProcessorContext;
 class TraceBlobView;
 class TrackEventTracker;
@@ -47,24 +47,24 @@ class TrackEventTokenizer {
   explicit TrackEventTokenizer(TraceProcessorContext*, TrackEventTracker*);
 
   ModuleResult TokenizeRangeOfInterestPacket(
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       const protos::pbzero::TracePacket_Decoder&,
       int64_t packet_timestamp);
   ModuleResult TokenizeTrackDescriptorPacket(
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       const protos::pbzero::TracePacket_Decoder&,
       int64_t packet_timestamp);
   ModuleResult TokenizeThreadDescriptorPacket(
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       const protos::pbzero::TracePacket_Decoder&);
-  void TokenizeTrackEventPacket(PacketSequenceState* state,
+  void TokenizeTrackEventPacket(RefPtr<PacketSequenceStateGeneration> state,
                                 const protos::pbzero::TracePacket_Decoder&,
                                 TraceBlobView* packet,
                                 int64_t packet_timestamp);
 
  private:
   void TokenizeThreadDescriptor(
-      PacketSequenceState* state,
+      PacketSequenceStateGeneration& state,
       const protos::pbzero::ThreadDescriptor_Decoder&);
   template <typename T>
   base::Status AddExtraCounterValues(
diff --git a/src/trace_processor/importers/proto/translation_table_module.cc b/src/trace_processor/importers/proto/translation_table_module.cc
index b86c3d3a5..8684df031 100644
--- a/src/trace_processor/importers/proto/translation_table_module.cc
+++ b/src/trace_processor/importers/proto/translation_table_module.cc
@@ -17,6 +17,7 @@
 
 #include "src/trace_processor/importers/common/args_translation_table.h"
 #include "src/trace_processor/importers/common/slice_translation_table.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 #include "protos/perfetto/trace/translation/translation_table.pbzero.h"
@@ -37,7 +38,7 @@ ModuleResult TranslationTableModule::TokenizePacket(
     const protos::pbzero::TracePacket_Decoder& decoder,
     TraceBlobView* /*packet*/,
     int64_t /*packet_timestamp*/,
-    PacketSequenceState* /*state*/,
+    RefPtr<PacketSequenceStateGeneration> /*state*/,
     uint32_t field_id) {
   if (field_id != TracePacket::kTranslationTableFieldNumber) {
     return ModuleResult::Ignored();
diff --git a/src/trace_processor/importers/proto/translation_table_module.h b/src/trace_processor/importers/proto/translation_table_module.h
index a318e6015..2060f4139 100644
--- a/src/trace_processor/importers/proto/translation_table_module.h
+++ b/src/trace_processor/importers/proto/translation_table_module.h
@@ -21,6 +21,7 @@
 #include <optional>
 
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/types/trace_processor_context.h"
 
@@ -37,7 +38,7 @@ class TranslationTableModule : public ProtoImporterModule {
       const protos::pbzero::TracePacket_Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
  private:
diff --git a/src/trace_processor/importers/proto/v8_module.cc b/src/trace_processor/importers/proto/v8_module.cc
index f90a71c40..89360b6b4 100644
--- a/src/trace_processor/importers/proto/v8_module.cc
+++ b/src/trace_processor/importers/proto/v8_module.cc
@@ -24,7 +24,6 @@
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 #include "src/trace_processor/importers/common/parser_types.h"
 #include "src/trace_processor/importers/common/process_tracker.h"
-#include "src/trace_processor/importers/proto/packet_sequence_state.h"
 #include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/v8_sequence_state.h"
 #include "src/trace_processor/importers/proto/v8_tracker.h"
@@ -58,11 +57,12 @@ V8Module::V8Module(TraceProcessorContext* context)
 
 V8Module::~V8Module() = default;
 
-ModuleResult V8Module::TokenizePacket(const TracePacket::Decoder&,
-                                      TraceBlobView* /*packet*/,
-                                      int64_t /*packet_timestamp*/,
-                                      PacketSequenceState* /*state*/,
-                                      uint32_t /*field_id*/) {
+ModuleResult V8Module::TokenizePacket(
+    const TracePacket::Decoder&,
+    TraceBlobView* /*packet*/,
+    int64_t /*packet_timestamp*/,
+    RefPtr<PacketSequenceStateGeneration> /*state*/,
+    uint32_t /*field_id*/) {
   return ModuleResult::Ignored();
 }
 
diff --git a/src/trace_processor/importers/proto/v8_module.h b/src/trace_processor/importers/proto/v8_module.h
index ea49b3208..e5825f717 100644
--- a/src/trace_processor/importers/proto/v8_module.h
+++ b/src/trace_processor/importers/proto/v8_module.h
@@ -22,6 +22,7 @@
 
 #include "perfetto/ext/base/flat_hash_map.h"
 #include "perfetto/protozero/field.h"
+#include "src/trace_processor/importers/proto/packet_sequence_state_generation.h"
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/tables/v8_tables_py.h"
@@ -53,7 +54,7 @@ class V8Module : public ProtoImporterModule {
       const protos::pbzero::TracePacket_Decoder& decoder,
       TraceBlobView* packet,
       int64_t packet_timestamp,
-      PacketSequenceState* state,
+      RefPtr<PacketSequenceStateGeneration> state,
       uint32_t field_id) override;
 
   void ParseTracePacketData(const protos::pbzero::TracePacket_Decoder&,
diff --git a/src/trace_processor/perfetto_sql/stdlib/android/battery_stats.sql b/src/trace_processor/perfetto_sql/stdlib/android/battery_stats.sql
index ff6902916..9431896f1 100644
--- a/src/trace_processor/perfetto_sql/stdlib/android/battery_stats.sql
+++ b/src/trace_processor/perfetto_sql/stdlib/android/battery_stats.sql
@@ -119,8 +119,10 @@ SELECT
 CREATE PERFETTO VIEW android_battery_stats_state(
   -- Timestamp in nanoseconds.
   ts INT,
-  -- The duration the state was active.
+  -- The duration the state was active, may be negative for incomplete slices.
   dur INT,
+  -- The same as `dur`, but extends to trace end for incomplete slices.
+  safe_dur INT,
   -- The name of the counter track.
   track_name STRING,
   -- The counter value as a number.
@@ -131,6 +133,7 @@ CREATE PERFETTO VIEW android_battery_stats_state(
 SELECT
   ts,
   IFNULL(LEAD(ts) OVER (PARTITION BY name ORDER BY ts) - ts, -1) AS dur,
+  LEAD(ts, 1, TRACE_END()) OVER (PARTITION BY name ORDER BY ts) - ts AS safe_dur,
   name AS track_name,
   CAST(value AS INT64) AS value,
   android_battery_stats_counter_to_string(name, value) AS value_name
@@ -160,8 +163,10 @@ WHERE counter_track.name GLOB 'battery_stats.*';
 CREATE PERFETTO VIEW android_battery_stats_event_slices(
   -- Timestamp in nanoseconds.
   ts INT,
-  -- The duration the state was active.
+  -- The duration the state was active, may be negative for incomplete slices.
   dur INT,
+  -- The same as `dur`, but extends to trace end for incomplete slices.
+  safe_dur INT,
   -- The name of the counter track.
   track_name STRING,
   -- String value.
@@ -208,6 +213,7 @@ WITH
 SELECT
   ts,
   IFNULL(end_ts-ts, -1) AS dur,
+  IFNULL(end_ts, TRACE_END()) - ts AS safe_dur,
   track_name,
   str_split(key, '"', 1) AS str_value,
   CAST(str_split(key, ':', 0) AS INT64) AS int_value
diff --git a/src/trace_processor/rpc/httpd.cc b/src/trace_processor/rpc/httpd.cc
index e63bdb499..21ca077ae 100644
--- a/src/trace_processor/rpc/httpd.cc
+++ b/src/trace_processor/rpc/httpd.cc
@@ -68,28 +68,28 @@ class Httpd : public base::HttpRequestHandler {
   base::HttpServer http_srv_;
 };
 
-base::HttpServerConnection* g_cur_conn;
-
 base::StringView Vec2Sv(const std::vector<uint8_t>& v) {
   return {reinterpret_cast<const char*>(v.data()), v.size()};
 }
 
 // Used both by websockets and /rpc chunked HTTP endpoints.
-void SendRpcChunk(const void* data, uint32_t len) {
+void SendRpcChunk(base::HttpServerConnection* conn,
+                  const void* data,
+                  uint32_t len) {
   if (data == nullptr) {
     // Unrecoverable RPC error case.
-    if (!g_cur_conn->is_websocket())
-      g_cur_conn->SendResponseBody("0\r\n\r\n", 5);
-    g_cur_conn->Close();
+    if (!conn->is_websocket())
+      conn->SendResponseBody("0\r\n\r\n", 5);
+    conn->Close();
     return;
   }
-  if (g_cur_conn->is_websocket()) {
-    g_cur_conn->SendWebsocketMessage(data, len);
+  if (conn->is_websocket()) {
+    conn->SendWebsocketMessage(data, len);
   } else {
     base::StackString<32> chunk_hdr("%x\r\n", len);
-    g_cur_conn->SendResponseBody(chunk_hdr.c_str(), chunk_hdr.len());
-    g_cur_conn->SendResponseBody(data, len);
-    g_cur_conn->SendResponseBody("\r\n", 2);
+    conn->SendResponseBody(chunk_hdr.c_str(), chunk_hdr.len());
+    conn->SendResponseBody(data, len);
+    conn->SendResponseBody("\r\n", 2);
   }
 }
 
@@ -165,13 +165,13 @@ void Httpd::OnHttpRequest(const base::HttpRequest& req) {
     // Start the chunked reply.
     conn.SendResponseHeaders("200 OK", chunked_headers,
                              base::HttpServerConnection::kOmitContentLength);
-    PERFETTO_CHECK(g_cur_conn == nullptr);
-    g_cur_conn = req.conn;
-    global_trace_processor_rpc_.SetRpcResponseFunction(SendRpcChunk);
+    global_trace_processor_rpc_.SetRpcResponseFunction(
+        [&](const void* data, uint32_t len) {
+          SendRpcChunk(&conn, data, len);
+        });
     // OnRpcRequest() will call SendRpcChunk() one or more times.
     global_trace_processor_rpc_.OnRpcRequest(req.body.data(), req.body.size());
     global_trace_processor_rpc_.SetRpcResponseFunction(nullptr);
-    g_cur_conn = nullptr;
 
     // Terminate chunked stream.
     conn.SendResponseBody("0\r\n\r\n", 5);
@@ -249,13 +249,13 @@ void Httpd::OnHttpRequest(const base::HttpRequest& req) {
 }
 
 void Httpd::OnWebsocketMessage(const base::WebsocketMessage& msg) {
-  PERFETTO_CHECK(g_cur_conn == nullptr);
-  g_cur_conn = msg.conn;
-  global_trace_processor_rpc_.SetRpcResponseFunction(SendRpcChunk);
+  global_trace_processor_rpc_.SetRpcResponseFunction(
+      [&](const void* data, uint32_t len) {
+        SendRpcChunk(msg.conn, data, len);
+      });
   // OnRpcRequest() will call SendRpcChunk() one or more times.
   global_trace_processor_rpc_.OnRpcRequest(msg.data.data(), msg.data.size());
   global_trace_processor_rpc_.SetRpcResponseFunction(nullptr);
-  g_cur_conn = nullptr;
 }
 
 }  // namespace
diff --git a/src/trace_processor/rpc/rpc.cc b/src/trace_processor/rpc/rpc.cc
index c715784f5..b676e455d 100644
--- a/src/trace_processor/rpc/rpc.cc
+++ b/src/trace_processor/rpc/rpc.cc
@@ -504,8 +504,7 @@ std::vector<uint8_t> Rpc::GetStatus() {
   protozero::HeapBuffered<protos::pbzero::StatusResult> status;
   status->set_loaded_trace_name(trace_processor_->GetCurrentTraceName());
   status->set_human_readable_version(base::GetVersionString());
-  const char* version_code = base::GetVersionCode();
-  if (version_code) {
+  if (const char* version_code = base::GetVersionCode(); version_code) {
     status->set_version_code(version_code);
   }
   status->set_api_version(protos::pbzero::TRACE_PROCESSOR_CURRENT_API_VERSION);
diff --git a/src/trace_processor/rpc/rpc.h b/src/trace_processor/rpc/rpc.h
index 61cf2e636..d3b2d41ff 100644
--- a/src/trace_processor/rpc/rpc.h
+++ b/src/trace_processor/rpc/rpc.h
@@ -22,6 +22,7 @@
 #include <functional>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "perfetto/base/status.h"
@@ -81,8 +82,11 @@ class Rpc {
   // with Wasm, where size_t = uint32_t.
   // (nullptr, 0) has the semantic of "close the channel" and is issued when an
   // unrecoverable wire-protocol framing error is detected.
-  using RpcResponseFunction = void (*)(const void* /*data*/, uint32_t /*len*/);
-  void SetRpcResponseFunction(RpcResponseFunction f) { rpc_response_fn_ = f; }
+  using RpcResponseFunction =
+      std::function<void(const void* /*data*/, uint32_t /*len*/)>;
+  void SetRpcResponseFunction(RpcResponseFunction f) {
+    rpc_response_fn_ = std::move(f);
+  }
 
   // 2. TraceProcessor legacy RPC endpoints.
   // The methods below are exposed for the old RPC interfaces, where each RPC
diff --git a/src/trace_processor/rpc/wasm_bridge.cc b/src/trace_processor/rpc/wasm_bridge.cc
index a9ba607c5..ca27abc0b 100644
--- a/src/trace_processor/rpc/wasm_bridge.cc
+++ b/src/trace_processor/rpc/wasm_bridge.cc
@@ -22,6 +22,8 @@
 namespace perfetto::trace_processor {
 
 namespace {
+using RpcResponseFn = void(const void*, uint32_t);
+
 Rpc* g_trace_processor_rpc;
 
 // The buffer used to pass the request arguments. The caller (JS) decides how
@@ -35,9 +37,9 @@ uint8_t* g_req_buf;
 extern "C" {
 
 // Returns the address of the allocated request buffer.
-uint8_t* EMSCRIPTEN_KEEPALIVE trace_processor_rpc_init(Rpc::RpcResponseFunction,
-                                                       uint32_t);
-uint8_t* trace_processor_rpc_init(Rpc::RpcResponseFunction resp_function,
+uint8_t* EMSCRIPTEN_KEEPALIVE
+trace_processor_rpc_init(RpcResponseFn* RpcResponseFn, uint32_t);
+uint8_t* trace_processor_rpc_init(RpcResponseFn* resp_function,
                                   uint32_t req_buffer_size) {
   g_trace_processor_rpc = new Rpc();
 
diff --git a/src/trace_redaction/BUILD.gn b/src/trace_redaction/BUILD.gn
index 8ff31a347..29c4a75e0 100644
--- a/src/trace_redaction/BUILD.gn
+++ b/src/trace_redaction/BUILD.gn
@@ -67,6 +67,8 @@ source_set("trace_redaction") {
     "redact_sched_switch.h",
     "redact_task_newtask.cc",
     "redact_task_newtask.h",
+    "remap_scheduling_events.cc",
+    "remap_scheduling_events.h",
     "scrub_ftrace_events.cc",
     "scrub_ftrace_events.h",
     "scrub_process_stats.cc",
@@ -105,6 +107,7 @@ source_set("integrationtests") {
     "filter_task_rename_integrationtest.cc",
     "prune_package_list_integrationtest.cc",
     "redact_sched_switch_integrationtest.cc",
+    "remap_scheduling_events_integrationtest.cc",
     "scrub_ftrace_events_integrationtest.cc",
     "scrub_process_stats_integrationtest.cc",
     "scrub_process_trees_integrationtest.cc",
@@ -143,6 +146,7 @@ perfetto_unittest_source_set("unittests") {
     "redact_process_free_unittest.cc",
     "redact_sched_switch_unittest.cc",
     "redact_task_newtask_unittest.cc",
+    "remap_scheduling_events_unittest.cc",
     "suspend_resume_unittest.cc",
   ]
   deps = [
diff --git a/src/trace_redaction/main.cc b/src/trace_redaction/main.cc
index 04cfa1eb0..051a424ac 100644
--- a/src/trace_redaction/main.cc
+++ b/src/trace_redaction/main.cc
@@ -32,6 +32,7 @@
 #include "src/trace_redaction/redact_process_free.h"
 #include "src/trace_redaction/redact_sched_switch.h"
 #include "src/trace_redaction/redact_task_newtask.h"
+#include "src/trace_redaction/remap_scheduling_events.h"
 #include "src/trace_redaction/scrub_ftrace_events.h"
 #include "src/trace_redaction/scrub_process_stats.h"
 #include "src/trace_redaction/scrub_process_trees.h"
@@ -88,6 +89,23 @@ static base::Status Main(std::string_view input,
   redact_ftrace_events
       ->emplace_back<RedactProcessFree::kFieldId, RedactProcessFree>();
 
+  // This set of transformations will change pids. This will break the
+  // connections between pids and the timeline (the synth threads are not in the
+  // timeline). If a transformation uses the timeline, it must be before this
+  // transformation.
+  auto* redact_sched_events = redactor.emplace_transform<RedactFtraceEvent>();
+  redact_sched_events->emplace_back<ThreadMergeRemapFtraceEventPid::kFieldId,
+                                    ThreadMergeRemapFtraceEventPid>();
+  redact_sched_events->emplace_back<ThreadMergeRemapSchedSwitchPid::kFieldId,
+                                    ThreadMergeRemapSchedSwitchPid>();
+  redact_sched_events->emplace_back<ThreadMergeRemapSchedWakingPid::kFieldId,
+                                    ThreadMergeRemapSchedWakingPid>();
+  redact_sched_events->emplace_back<
+      ThreadMergeDropField::kTaskNewtaskFieldNumber, ThreadMergeDropField>();
+  redact_sched_events
+      ->emplace_back<ThreadMergeDropField::kSchedProcessFreeFieldNumber,
+                     ThreadMergeDropField>();
+
   Context context;
   context.package_name = package_name;
 
diff --git a/src/trace_redaction/remap_scheduling_events.cc b/src/trace_redaction/remap_scheduling_events.cc
new file mode 100644
index 000000000..4817f1972
--- /dev/null
+++ b/src/trace_redaction/remap_scheduling_events.cc
@@ -0,0 +1,261 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+#include "src/trace_redaction/remap_scheduling_events.h"
+
+#include "src/trace_redaction/proto_util.h"
+
+#include "protos/perfetto/trace/ftrace/ftrace_event.pbzero.h"
+#include "protos/perfetto/trace/ftrace/ftrace_event_bundle.pbzero.h"
+#include "protos/perfetto/trace/ftrace/sched.pbzero.h"
+
+namespace perfetto::trace_redaction {
+
+namespace {
+int32_t RemapPid(const Context& context,
+                 uint64_t timestamp,
+                 uint32_t cpu,
+                 int32_t pid) {
+  PERFETTO_DCHECK(context.package_uid.value());
+  PERFETTO_DCHECK(cpu < context.synthetic_threads->tids.size());
+
+  auto slice = context.timeline->Search(timestamp, pid);
+
+  auto expected_uid = NormalizeUid(slice.uid);
+  auto actual_uid = NormalizeUid(context.package_uid.value());
+
+  return !pid || expected_uid == actual_uid
+             ? pid
+             : context.synthetic_threads->tids[cpu];
+}
+}  // namespace
+
+base::Status ThreadMergeRemapFtraceEventPid::Redact(
+    const Context& context,
+    const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+    protozero::ProtoDecoder& event,
+    protos::pbzero::FtraceEvent* event_message) const {
+  if (!context.package_uid.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: missing package uid");
+  }
+
+  if (!context.synthetic_threads.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: missing synthetic threads");
+  }
+
+  // This should never happen. A bundle should have a cpu.
+  if (!bundle.has_cpu()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: Invalid ftrace event, missing cpu.");
+  }
+
+  if (bundle.cpu() >= context.synthetic_threads->tids.size()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: synthetic thread count");
+  }
+
+  auto timestamp =
+      event.FindField(protos::pbzero::FtraceEvent::kTimestampFieldNumber);
+
+  // This should never happen. An event should have a timestamp.
+  if (!timestamp.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: Invalid ftrace event, missing "
+        "timestamp.");
+  }
+
+  // This handler should only be called for the pid field.
+  auto pid = event.FindField(protos::pbzero::FtraceEvent::kPidFieldNumber);
+  PERFETTO_DCHECK(pid.valid());
+
+  // The event's pid is technically a uint, but we need it as a int.
+  auto new_pid =
+      RemapPid(context, timestamp.as_uint64(), bundle.cpu(), pid.as_int32());
+  event_message->set_pid(static_cast<uint32_t>(new_pid));
+
+  return base::OkStatus();
+}
+
+base::Status ThreadMergeRemapSchedSwitchPid::Redact(
+    const Context& context,
+    const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+    protozero::ProtoDecoder& event,
+    protos::pbzero::FtraceEvent* event_message) const {
+  if (!context.package_uid.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: missing package uid");
+  }
+
+  if (!context.synthetic_threads.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: missing synthetic threads");
+  }
+
+  // This should never happen. A bundle should have a cpu.
+  if (!bundle.has_cpu()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: Invalid ftrace event, missing cpu.");
+  }
+
+  if (bundle.cpu() >= context.synthetic_threads->tids.size()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: synthetic thread count");
+  }
+
+  auto timestamp =
+      event.FindField(protos::pbzero::FtraceEvent::kTimestampFieldNumber);
+
+  // This should never happen. An event should have a timestamp.
+  if (!timestamp.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: Invalid ftrace event, missing "
+        "timestamp.");
+  }
+
+  // This handler should only be called for the sched switch field.
+  auto sched_switch =
+      event.FindField(protos::pbzero::FtraceEvent::kSchedSwitchFieldNumber);
+  PERFETTO_DCHECK(sched_switch.valid());
+
+  protozero::ProtoDecoder sched_switch_decoder(sched_switch.as_bytes());
+
+  auto old_prev_pid_field = sched_switch_decoder.FindField(
+      protos::pbzero::SchedSwitchFtraceEvent::kPrevPidFieldNumber);
+  auto old_next_pid_field = sched_switch_decoder.FindField(
+      protos::pbzero::SchedSwitchFtraceEvent::kNextPidFieldNumber);
+
+  if (!old_prev_pid_field.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: Invalid sched switch event, missing "
+        "prev pid");
+  }
+
+  if (!old_next_pid_field.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: Invalid sched switch event, missing "
+        "next pid");
+  }
+
+  auto new_prev_pid_field =
+      RemapPid(context, timestamp.as_uint64(), bundle.cpu(),
+               old_prev_pid_field.as_int32());
+  auto new_next_pid_field =
+      RemapPid(context, timestamp.as_uint64(), bundle.cpu(),
+               old_next_pid_field.as_int32());
+
+  auto* sched_switch_message = event_message->set_sched_switch();
+
+  for (auto f = sched_switch_decoder.ReadField(); f.valid();
+       f = sched_switch_decoder.ReadField()) {
+    switch (f.id()) {
+      case protos::pbzero::SchedSwitchFtraceEvent::kPrevPidFieldNumber:
+        sched_switch_message->set_prev_pid(new_prev_pid_field);
+        break;
+
+      case protos::pbzero::SchedSwitchFtraceEvent::kNextPidFieldNumber:
+        sched_switch_message->set_next_pid(new_next_pid_field);
+        break;
+
+      default:
+        proto_util::AppendField(f, sched_switch_message);
+        break;
+    }
+  }
+
+  return base::OkStatus();
+}
+
+base::Status ThreadMergeRemapSchedWakingPid::Redact(
+    const Context& context,
+    const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+    protozero::ProtoDecoder& event,
+    protos::pbzero::FtraceEvent* event_message) const {
+  if (!context.package_uid.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: missing package uid");
+  }
+
+  if (!context.synthetic_threads.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: missing synthetic threads");
+  }
+
+  // This should never happen. A bundle should have a cpu.
+  if (!bundle.has_cpu()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: Invalid ftrace event, missing cpu.");
+  }
+
+  if (bundle.cpu() >= context.synthetic_threads->tids.size()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: synthetic thread count");
+  }
+
+  auto timestamp =
+      event.FindField(protos::pbzero::FtraceEvent::kTimestampFieldNumber);
+
+  // This should never happen. An event should have a timestamp.
+  if (!timestamp.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: Invalid ftrace event, missing "
+        "timestamp.");
+  }
+
+  // This handler should only be called for the sched waking field.
+  auto sched_waking =
+      event.FindField(protos::pbzero::FtraceEvent::kSchedWakingFieldNumber);
+  PERFETTO_DCHECK(sched_waking.valid());
+
+  protozero::ProtoDecoder sched_waking_decoder(sched_waking.as_bytes());
+
+  auto old_pid = sched_waking_decoder.FindField(
+      protos::pbzero::SchedWakingFtraceEvent::kPidFieldNumber);
+
+  if (!old_pid.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: Invalid sched waking event, missing "
+        "pid");
+  }
+
+  auto new_pid_field = RemapPid(context, timestamp.as_uint64(), bundle.cpu(),
+                                old_pid.as_int32());
+
+  auto* sched_waking_message = event_message->set_sched_waking();
+
+  for (auto f = sched_waking_decoder.ReadField(); f.valid();
+       f = sched_waking_decoder.ReadField()) {
+    if (f.id() == protos::pbzero::SchedWakingFtraceEvent::kPidFieldNumber) {
+      sched_waking_message->set_pid(new_pid_field);
+    } else {
+      proto_util::AppendField(f, sched_waking_message);
+    }
+  }
+
+  return base::OkStatus();
+}
+
+// By doing nothing, the field gets dropped.
+base::Status ThreadMergeDropField::Redact(
+    const Context&,
+    const protos::pbzero::FtraceEventBundle::Decoder&,
+    protozero::ProtoDecoder&,
+    protos::pbzero::FtraceEvent*) const {
+  return base::OkStatus();
+}
+
+}  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/remap_scheduling_events.h b/src/trace_redaction/remap_scheduling_events.h
new file mode 100644
index 000000000..016053489
--- /dev/null
+++ b/src/trace_redaction/remap_scheduling_events.h
@@ -0,0 +1,137 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+#ifndef SRC_TRACE_REDACTION_REMAP_SCHEDULING_EVENTS_H_
+#define SRC_TRACE_REDACTION_REMAP_SCHEDULING_EVENTS_H_
+
+#include "perfetto/protozero/proto_decoder.h"
+#include "src/trace_redaction/redact_ftrace_event.h"
+#include "src/trace_redaction/trace_redaction_framework.h"
+
+#include "protos/perfetto/trace/ftrace/ftrace_event.pbzero.h"
+
+namespace perfetto::trace_redaction {
+
+// Reads the Ftrace event's pid and replaces it with a synthetic thread id (if
+// necessary).
+class ThreadMergeRemapFtraceEventPid : public FtraceEventRedaction {
+ public:
+  static constexpr auto kFieldId = protos::pbzero::FtraceEvent::kPidFieldNumber;
+
+  base::Status Redact(
+      const Context& context,
+      const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+      protozero::ProtoDecoder& event,
+      protos::pbzero::FtraceEvent* event_message) const override;
+};
+
+// Reads the sched switch pid and replaces it with a synthetic thread id (if
+// necessary).
+//
+//  event {
+//    timestamp: 6702093743539938
+//    pid: 0
+//    sched_switch {
+//      prev_comm: "swapper/7"
+//      prev_pid: 0
+//      prev_prio: 120
+//      prev_state: 0
+//      next_comm: "FMOD stream thr"
+//      next_pid: 7174
+//      next_prio: 104
+//    }
+//  }
+class ThreadMergeRemapSchedSwitchPid : public FtraceEventRedaction {
+ public:
+  static constexpr auto kFieldId =
+      protos::pbzero::FtraceEvent::kSchedSwitchFieldNumber;
+
+  base::Status Redact(
+      const Context& context,
+      const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+      protozero::ProtoDecoder& event,
+      protos::pbzero::FtraceEvent* event_message) const override;
+};
+
+// Reads the sched waking pid and replaces it with a synthetic thread id (if
+// necessary).
+//
+//  event {
+//    timestamp: 6702093743527386
+//    pid: 0
+//    sched_waking {
+//      comm: "FMOD stream thr"
+//      pid: 7174
+//      prio: 104
+//      success: 1
+//      target_cpu: 7
+//    }
+//  }
+class ThreadMergeRemapSchedWakingPid : public FtraceEventRedaction {
+ public:
+  static constexpr auto kFieldId =
+      protos::pbzero::FtraceEvent::kSchedWakingFieldNumber;
+
+  base::Status Redact(
+      const Context& context,
+      const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+      protozero::ProtoDecoder& event,
+      protos::pbzero::FtraceEvent* event_message) const override;
+};
+
+// Drop "new task" events because it's safe to assume that the threads always
+// exist.
+//
+//  event {
+//    timestamp: 6702094133317685
+//    pid: 6167
+//    task_newtask {
+//      pid: 7972                 <-- Pid being started
+//      comm: "adbd"
+//      clone_flags: 4001536
+//      oom_score_adj: -1000
+//    }
+//  }
+//
+// Drop "process free" events because it's safe to assume that the threads
+// always exist.
+//
+//  event {
+//    timestamp: 6702094703942898
+//    pid: 10
+//    sched_process_free {
+//      comm: "shell svc 7973"
+//      pid: 7974                 <-- Pid being freed
+//      prio: 120
+//    }
+//  }
+class ThreadMergeDropField : public FtraceEventRedaction {
+ public:
+  static constexpr auto kTaskNewtaskFieldNumber =
+      protos::pbzero::FtraceEvent::kTaskNewtaskFieldNumber;
+  static constexpr auto kSchedProcessFreeFieldNumber =
+      protos::pbzero::FtraceEvent::kSchedProcessFreeFieldNumber;
+
+  base::Status Redact(
+      const Context& context,
+      const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+      protozero::ProtoDecoder& event,
+      protos::pbzero::FtraceEvent* event_message) const override;
+};
+
+}  // namespace perfetto::trace_redaction
+
+#endif  // SRC_TRACE_REDACTION_REMAP_SCHEDULING_EVENTS_H_
diff --git a/src/trace_redaction/remap_scheduling_events_integrationtest.cc b/src/trace_redaction/remap_scheduling_events_integrationtest.cc
new file mode 100644
index 000000000..2f68c95bf
--- /dev/null
+++ b/src/trace_redaction/remap_scheduling_events_integrationtest.cc
@@ -0,0 +1,292 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+#include "src/base/test/status_matchers.h"
+#include "src/trace_redaction/collect_system_info.h"
+#include "src/trace_redaction/collect_timeline_events.h"
+#include "src/trace_redaction/find_package_uid.h"
+#include "src/trace_redaction/redact_ftrace_event.h"
+#include "src/trace_redaction/remap_scheduling_events.h"
+#include "src/trace_redaction/trace_redaction_integration_fixture.h"
+#include "test/gtest_and_gmock.h"
+
+#include "protos/perfetto/trace/ftrace/sched.pbzero.h"
+#include "protos/perfetto/trace/ftrace/task.pbzero.h"
+
+namespace perfetto::trace_redaction {
+
+// Runs ThreadMergeRemapFtraceEventPid, ThreadMergeRemapSchedSwitchPid,
+// ThreadMergeRemapSchedWakingPid, and ThreadMergeDropField to replace pids with
+// synthetic pids (for all threads outside of the target package);
+class RemapSchedulingEventsIntegrationTest
+    : public testing::Test,
+      protected TraceRedactionIntegrationFixure {
+ public:
+  static constexpr auto kPackageName =
+      "com.Unity.com.unity.multiplayer.samples.coop";
+  static constexpr uint64_t kPackageId = 10252;
+  static constexpr int32_t kPid = 7105;
+
+  // Threads belonging to pid 7105. Collected using trace processors.
+  static constexpr auto kTids = {
+      0,  // pid 0 will always be included because CPU idle uses it.
+      7105, 7111, 7112, 7113, 7114, 7115, 7116, 7117, 7118, 7119, 7120,
+      7124, 7125, 7127, 7129, 7130, 7131, 7132, 7133, 7134, 7135, 7136,
+      7137, 7139, 7141, 7142, 7143, 7144, 7145, 7146, 7147, 7148, 7149,
+      7150, 7151, 7152, 7153, 7154, 7155, 7156, 7157, 7158, 7159, 7160,
+      7161, 7162, 7163, 7164, 7165, 7166, 7167, 7171, 7172, 7174, 7178,
+      7180, 7184, 7200, 7945, 7946, 7947, 7948, 7950, 7969,
+  };
+
+ protected:
+  void SetUp() override {
+    trace_redactor()->emplace_collect<FindPackageUid>();
+
+    // In order to remap threads, we need to have synth threads.
+    trace_redactor()->emplace_collect<CollectSystemInfo>();
+    trace_redactor()->emplace_build<BuildSyntheticThreads>();
+
+    // Timeline information is needed to know if a pid belongs to a package.
+    trace_redactor()->emplace_collect<CollectTimelineEvents>();
+
+    auto* redactions = trace_redactor()->emplace_transform<RedactFtraceEvent>();
+    redactions->emplace_back<ThreadMergeRemapFtraceEventPid::kFieldId,
+                             ThreadMergeRemapFtraceEventPid>();
+    redactions->emplace_back<ThreadMergeRemapSchedSwitchPid::kFieldId,
+                             ThreadMergeRemapSchedSwitchPid>();
+    redactions->emplace_back<ThreadMergeRemapSchedWakingPid::kFieldId,
+                             ThreadMergeRemapSchedWakingPid>();
+    redactions->emplace_back<ThreadMergeDropField::kSchedProcessFreeFieldNumber,
+                             ThreadMergeDropField>();
+    redactions->emplace_back<ThreadMergeDropField::kTaskNewtaskFieldNumber,
+                             ThreadMergeDropField>();
+
+    context()->package_name = kPackageName;
+  }
+
+  struct Index {
+    // List of FtraceEvent
+    std::vector<protozero::ConstBytes> events;
+
+    // List of SchedSwitchFtraceEvent
+    std::vector<protozero::ConstBytes> events_sched_switch;
+
+    // List of SchedWakingFtraceEvent
+    std::vector<protozero::ConstBytes> events_sched_waking;
+
+    // List of SchedProcessFreeFtraceEvent
+    std::vector<protozero::ConstBytes> events_sched_process_free;
+
+    // List of TaskNewtaskFtraceEvent
+    std::vector<protozero::ConstBytes> events_task_newtask;
+  };
+
+  void UpdateFtraceIndex(protozero::ConstBytes bytes, Index* index) {
+    protos::pbzero::FtraceEventBundle::Decoder bundle(bytes);
+
+    for (auto event = bundle.event(); event; ++event) {
+      index->events.push_back(event->as_bytes());
+
+      // protos::pbzero::FtraceEvent
+      protozero::ProtoDecoder ftrace_event(event->as_bytes());
+
+      auto sched_switch = ftrace_event.FindField(
+          protos::pbzero::FtraceEvent::kSchedSwitchFieldNumber);
+      if (sched_switch.valid()) {
+        index->events_sched_switch.push_back(sched_switch.as_bytes());
+      }
+
+      auto sched_waking = ftrace_event.FindField(
+          protos::pbzero::FtraceEvent::kSchedWakingFieldNumber);
+      if (sched_waking.valid()) {
+        index->events_sched_waking.push_back(sched_waking.as_bytes());
+      }
+
+      auto sched_process_free = ftrace_event.FindField(
+          protos::pbzero::FtraceEvent::kSchedProcessFreeFieldNumber);
+      if (sched_process_free.valid()) {
+        index->events_sched_process_free.push_back(
+            sched_process_free.as_bytes());
+      }
+
+      auto task_newtask = ftrace_event.FindField(
+          protos::pbzero::FtraceEvent::kTaskNewtaskFieldNumber);
+      if (task_newtask.valid()) {
+        index->events_task_newtask.push_back(task_newtask.as_bytes());
+      }
+    }
+  }
+
+  // Bytes should be TracePacket
+  Index CreateFtraceIndex(const std::string& bytes) {
+    Index index;
+
+    protozero::ProtoDecoder packet_decoder(bytes);
+
+    for (auto packet = packet_decoder.ReadField(); packet.valid();
+         packet = packet_decoder.ReadField()) {
+      auto events = packet_decoder.FindField(
+          protos::pbzero::TracePacket::kFtraceEventsFieldNumber);
+
+      if (events.valid()) {
+        UpdateFtraceIndex(events.as_bytes(), &index);
+      }
+    }
+
+    return index;
+  }
+
+  base::StatusOr<std::string> LoadAndRedactTrace() {
+    auto source = LoadOriginal();
+
+    if (!source.ok()) {
+      return source.status();
+    }
+
+    auto redact = Redact();
+
+    if (!redact.ok()) {
+      return redact;
+    }
+
+    // Double-check the package id with the one from trace processor. If this
+    // was wrong and this check was missing, finding the problem would be much
+    // harder.
+    if (!context()->package_uid.has_value()) {
+      return base::ErrStatus("Missing package uid.");
+    }
+
+    if (context()->package_uid.value() != kPackageId) {
+      return base::ErrStatus("Unexpected package uid found.");
+    }
+
+    auto redacted = LoadRedacted();
+
+    if (redacted.ok()) {
+      return redacted;
+    }
+
+    // System info is used to initialize the synth threads. If these are wrong,
+    // then the synth threads will be wrong.
+    if (!context()->system_info.has_value()) {
+      return base::ErrStatus("Missing system info.");
+    }
+
+    if (context()->system_info->last_cpu() != 7u) {
+      return base::ErrStatus("Unexpected cpu count.");
+    }
+
+    // The synth threads should have been initialized. They will be used here to
+    // verify which threads exist in the redacted trace.
+    if (!context()->synthetic_threads.has_value()) {
+      return base::ErrStatus("Missing synthetic threads.");
+    }
+
+    if (context()->synthetic_threads->tids.size() != 8u) {
+      return base::ErrStatus("Unexpected synthentic thread count.");
+    }
+
+    return redacted;
+  }
+
+  // Should be called after redaction since it requires data from the context.
+  std::unordered_set<int32_t> CopyAllowedTids(const Context& context) const {
+    std::unordered_set<int32_t> tids(kTids.begin(), kTids.end());
+
+    tids.insert(context.synthetic_threads->tgid);
+    tids.insert(context.synthetic_threads->tids.begin(),
+                context.synthetic_threads->tids.end());
+
+    return tids;
+  }
+
+ private:
+  std::unordered_set<int32_t> allowed_tids_;
+};
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FilterFtraceEventPid) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events) {
+    protos::pbzero::FtraceEvent::Decoder decoder(event);
+    auto pid = static_cast<int32_t>(decoder.pid());
+    ASSERT_TRUE(allowlist.count(pid));
+  }
+}
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FiltersSchedSwitch) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events_sched_switch) {
+    protos::pbzero::SchedSwitchFtraceEvent::Decoder decoder(event);
+    ASSERT_TRUE(allowlist.count(decoder.prev_pid()));
+    ASSERT_TRUE(allowlist.count(decoder.next_pid()));
+  }
+}
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FiltersSchedWaking) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events_sched_waking) {
+    protos::pbzero::SchedWakingFtraceEvent::Decoder decoder(event);
+    ASSERT_TRUE(allowlist.count(decoder.pid()));
+  }
+}
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FiltersProcessFree) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events_sched_process_free) {
+    protos::pbzero::SchedProcessFreeFtraceEvent::Decoder decoder(event);
+    ASSERT_TRUE(allowlist.count(decoder.pid()));
+  }
+}
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FiltersNewTask) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events_task_newtask) {
+    protos::pbzero::TaskNewtaskFtraceEvent::Decoder decoder(event);
+    ASSERT_TRUE(allowlist.count(decoder.pid()));
+  }
+}
+
+}  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/remap_scheduling_events_unittest.cc b/src/trace_redaction/remap_scheduling_events_unittest.cc
new file mode 100644
index 000000000..5cdc75373
--- /dev/null
+++ b/src/trace_redaction/remap_scheduling_events_unittest.cc
@@ -0,0 +1,441 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+#include "src/trace_redaction/remap_scheduling_events.h"
+
+#include "perfetto/protozero/scattered_heap_buffer.h"
+#include "src/base/test/status_matchers.h"
+#include "src/trace_redaction/trace_redaction_framework.h"
+#include "test/gtest_and_gmock.h"
+
+#include "protos/perfetto/trace/ftrace/ftrace_event.gen.h"
+#include "protos/perfetto/trace/ftrace/ftrace_event_bundle.gen.h"
+#include "protos/perfetto/trace/ftrace/sched.gen.h"
+
+namespace perfetto::trace_redaction {
+
+template <class T>
+class ThreadMergeTest {
+ protected:
+  struct Process {
+    uint64_t uid;
+    int32_t ppid;
+    int32_t pid;
+  };
+
+  base::Status Redact(protos::pbzero::FtraceEvent* event_message) {
+    T redact;
+
+    auto bundle_str = bundle_.SerializeAsString();
+    protos::pbzero::FtraceEventBundle::Decoder bundle_decoder(bundle_str);
+
+    auto event_str = bundle_.event().back().SerializeAsString();
+    protos::pbzero::FtraceEvent::Decoder event_decoder(event_str);
+
+    return redact.Redact(context_, bundle_decoder, event_decoder,
+                         event_message);
+  }
+
+  Context context_;
+  protos::gen::FtraceEventBundle bundle_;
+};
+
+// All ftrace events have a timestamp and a pid. This test focuses on the
+// event's pid value. When that pid doesn't belong to the target package, it
+// should be replaced with a synthetic thread id.
+//
+//  event {
+//    timestamp: 6702093743539938
+//    pid: 0
+//    sched_switch { ... }
+//  }
+class ThreadMergeRemapFtraceEventPidTest
+    : public testing::Test,
+      protected ThreadMergeTest<ThreadMergeRemapFtraceEventPid> {
+ protected:
+  static constexpr uint32_t kCpu = 3;
+
+  static constexpr auto kTimestamp = 123456789;
+
+  // This process will be connected to the target package.
+  static constexpr Process kProcess = {12, 5, 7};
+
+  // This process will not be connected to the target package.
+  static constexpr Process kOtherProcess = {120, 50, 70};
+
+  void SetUp() override {
+    bundle_.add_event();
+
+    context_.package_uid = kProcess.uid;
+
+    context_.timeline = std::make_unique<ProcessThreadTimeline>();
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kProcess.pid, kProcess.ppid, kProcess.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kOtherProcess.pid, kOtherProcess.ppid, kOtherProcess.uid));
+    context_.timeline->Sort();
+
+    // Because kCpu is 3, it means that there are four CPUs (id 0, id 1, ...).
+    context_.synthetic_threads.emplace();
+    context_.synthetic_threads->tids.assign({100, 101, 102, 103});
+  }
+};
+
+// This should never happen, a bundle should always have a cpu. If it doesn't
+// have a CPU, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapFtraceEventPidTest, MissingCpuReturnsError) {
+  // Do not call set_cpu(uint32_t value). There should be no cpu for this case.
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+// This should never happen, an event should always have a timestamp. If it
+// doesn't have a timestamp, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapFtraceEventPidTest, MissingTimestampReturnsError) {
+  bundle_.set_cpu(kCpu);
+  // Do not call set_timestamp(uint64_t value). There should be no timestamp for
+  // this case.
+  bundle_.mutable_event()->back().set_pid(kProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+TEST_F(ThreadMergeRemapFtraceEventPidTest, NoopWhenPidIsInPackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.pid()), kProcess.pid);
+}
+
+TEST_F(ThreadMergeRemapFtraceEventPidTest, ChangesPidWhenPidIsOutsidePackage) {
+  bundle_.set_cpu(kCpu);  // The CPU is used to select the pid.
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kOtherProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.pid()),
+            context_.synthetic_threads->tids[kCpu]);
+}
+
+// When creating a sched_switch event, the event pid and the previous pid should
+// be the same pid.
+//
+//  event {
+//    timestamp: 6702093743539938
+//    pid: 0
+//    sched_switch {
+//      prev_comm: "swapper/7"
+//      prev_pid: 0
+//      prev_prio: 120
+//      prev_state: 0
+//      next_comm: "FMOD stream thr"
+//      next_pid: 7174
+//      next_prio: 104
+//    }
+//  }
+class ThreadMergeRemapSchedSwitchPidTest
+    : public testing::Test,
+      protected ThreadMergeTest<ThreadMergeRemapSchedSwitchPid> {
+ protected:
+  static constexpr uint32_t kCpu = 3;
+
+  static constexpr auto kTimestamp = 123456789;
+
+  // This process will be connected to the target package.
+  static constexpr Process kPrevProcess = {12, 5, 7};
+  static constexpr Process kNextProcess = {12, 5, 8};
+
+  // This process will not be connected to the target package.
+  static constexpr Process kOtherProcess = {120, 50, 70};
+
+  void SetUp() override {
+    bundle_.add_event();
+
+    context_.package_uid = kPrevProcess.uid;
+
+    context_.timeline = std::make_unique<ProcessThreadTimeline>();
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kPrevProcess.pid, kPrevProcess.ppid, kPrevProcess.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kNextProcess.pid, kNextProcess.ppid, kNextProcess.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kOtherProcess.pid, kOtherProcess.ppid, kOtherProcess.uid));
+
+    context_.timeline->Sort();
+
+    // Because kCpu is 3, it means that there are four CPUs (id 0, id 1, ...).
+    context_.synthetic_threads.emplace();
+    context_.synthetic_threads->tids.assign({100, 101, 102, 103});
+  }
+};
+
+// This should never happen, a bundle should always have a cpu. If it doesn't
+// have a CPU, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapSchedSwitchPidTest, MissingCpuReturnsError) {
+  // Do not call set_cpu(uint32_t value). There should be no cpu for this case.
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+// This should never happen, an event should always have a timestamp. If it
+// doesn't have a timestamp, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapSchedSwitchPidTest, MissingTimestampReturnsError) {
+  bundle_.set_cpu(kCpu);
+  // Do not call set_timestamp(uint64_t value). There should be no timestamp for
+  // this case.
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+TEST_F(ThreadMergeRemapSchedSwitchPidTest, NoopWhenPidIsInPackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  auto* sched_switch = bundle_.mutable_event()->back().mutable_sched_switch();
+  sched_switch->set_prev_pid(kPrevProcess.pid);
+  sched_switch->set_next_pid(kNextProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_switch());
+
+  ASSERT_TRUE(event.sched_switch().has_prev_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().prev_pid()),
+            kPrevProcess.pid);
+
+  ASSERT_TRUE(event.sched_switch().has_next_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().next_pid()),
+            kNextProcess.pid);
+}
+
+TEST_F(ThreadMergeRemapSchedSwitchPidTest,
+       ChangesPrevPidWhenPidIsOutsidePackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  auto* sched_switch = bundle_.mutable_event()->back().mutable_sched_switch();
+  sched_switch->set_prev_pid(kOtherProcess.pid);
+  sched_switch->set_next_pid(kNextProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_switch());
+
+  ASSERT_TRUE(event.sched_switch().has_prev_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().prev_pid()),
+            context_.synthetic_threads->tids[kCpu]);
+
+  ASSERT_TRUE(event.sched_switch().has_next_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().next_pid()),
+            kNextProcess.pid);
+}
+
+TEST_F(ThreadMergeRemapSchedSwitchPidTest,
+       ChangesNextPidWhenPidIsOutsidePackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  auto* sched_switch = bundle_.mutable_event()->back().mutable_sched_switch();
+  sched_switch->set_prev_pid(kPrevProcess.pid);
+  sched_switch->set_next_pid(kOtherProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_switch());
+
+  ASSERT_TRUE(event.sched_switch().has_prev_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().prev_pid()),
+            kPrevProcess.pid);
+
+  ASSERT_TRUE(event.sched_switch().has_next_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().next_pid()),
+            context_.synthetic_threads->tids[kCpu]);
+}
+
+//  event {
+//    timestamp: 6702093743527386
+//    pid: 0
+//    sched_waking {
+//      comm: "FMOD stream thr"
+//      pid: 7174
+//      prio: 104
+//      success: 1
+//      target_cpu: 7
+//    }
+//  }
+class ThreadMergeRemapSchedWakingPidTest
+    : public testing::Test,
+      protected ThreadMergeTest<ThreadMergeRemapSchedWakingPid> {
+ protected:
+  static constexpr uint32_t kCpu = 3;
+
+  static constexpr auto kTimestamp = 123456789;
+
+  // This process will be connected to the target package.
+  static constexpr Process kWakerProcess = {12, 5, 7};
+  static constexpr Process kWakeTarget = {12, 5, 8};
+
+  // This process will not be connected to the target package.
+  static constexpr Process kOtherProcess = {120, 50, 70};
+
+  void SetUp() override {
+    bundle_.add_event();
+
+    context_.package_uid = kWakerProcess.uid;
+
+    context_.timeline = std::make_unique<ProcessThreadTimeline>();
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kWakerProcess.pid, kWakerProcess.ppid, kWakerProcess.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kWakeTarget.pid, kWakeTarget.ppid, kWakeTarget.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kOtherProcess.pid, kOtherProcess.ppid, kOtherProcess.uid));
+
+    context_.timeline->Sort();
+
+    // Because kCpu is 3, it means that there are four CPUs (id 0, id 1, ...).
+    context_.synthetic_threads.emplace();
+    context_.synthetic_threads->tids.assign({100, 101, 102, 103});
+  }
+};
+
+// This should never happen, a bundle should always have a cpu. If it doesn't
+// have a CPU, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapSchedWakingPidTest, MissingCpuReturnsError) {
+  // Do not call set_cpu(uint32_t value). There should be no cpu for this case.
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kWakerProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+// This should never happen, an event should always have a timestamp. If it
+// doesn't have a timestamp, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapSchedWakingPidTest, MissingTimestampReturnsError) {
+  bundle_.set_cpu(kCpu);
+  // Do not call set_timestamp(uint64_t value). There should be no timestamp for
+  // this case.
+  bundle_.mutable_event()->back().set_pid(kWakerProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+TEST_F(ThreadMergeRemapSchedWakingPidTest, NoopWhenPidIsInPackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kWakerProcess.pid);
+
+  auto* sched_waking = bundle_.mutable_event()->back().mutable_sched_waking();
+  sched_waking->set_pid(kWakeTarget.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_waking());
+
+  ASSERT_TRUE(event.sched_waking().has_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_waking().pid()), kWakeTarget.pid);
+}
+
+TEST_F(ThreadMergeRemapSchedWakingPidTest, ChangesPidWhenPidIsOutsidePackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kWakerProcess.pid);
+
+  auto* sched_switch = bundle_.mutable_event()->back().mutable_sched_waking();
+  sched_switch->set_pid(kOtherProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_waking());
+
+  ASSERT_TRUE(event.sched_waking().has_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_waking().pid()),
+            context_.synthetic_threads->tids[kCpu]);
+}
+
+}  // namespace perfetto::trace_redaction
diff --git a/test/data/ui-screenshots/ui-chrome_missing_track_names_load.png.sha256 b/test/data/ui-screenshots/ui-chrome_missing_track_names_load.png.sha256
index 81c36d88e..ea62427a6 100644
--- a/test/data/ui-screenshots/ui-chrome_missing_track_names_load.png.sha256
+++ b/test/data/ui-screenshots/ui-chrome_missing_track_names_load.png.sha256
@@ -1 +1 @@
-ff45b08442a6f59a4768f3ad697cffb43d19889f1f6523d5ed4ef83aff944864
\ No newline at end of file
+8b878dd1d284fda4558de3f184d078293546c09a1438e14183161eb6040cee26
\ No newline at end of file
diff --git a/test/data/ui-screenshots/ui-chrome_rendering_desktop_expand_browser_proc.png.sha256 b/test/data/ui-screenshots/ui-chrome_rendering_desktop_expand_browser_proc.png.sha256
index 4d7b25a06..e6a988a37 100644
--- a/test/data/ui-screenshots/ui-chrome_rendering_desktop_expand_browser_proc.png.sha256
+++ b/test/data/ui-screenshots/ui-chrome_rendering_desktop_expand_browser_proc.png.sha256
@@ -1 +1 @@
-161c05f6bae9b0e2142e73e83f3bcf7a9c689c58411db9e345a63ba882605557
\ No newline at end of file
+d3e7cdfbb83dd4e1b674c5d7b90554fd915dab8747f3981a2b6dbb1fff955f61
\ No newline at end of file
diff --git a/test/data/ui-screenshots/ui-chrome_rendering_desktop_load.png.sha256 b/test/data/ui-screenshots/ui-chrome_rendering_desktop_load.png.sha256
index 2d4876158..1be3a540d 100644
--- a/test/data/ui-screenshots/ui-chrome_rendering_desktop_load.png.sha256
+++ b/test/data/ui-screenshots/ui-chrome_rendering_desktop_load.png.sha256
@@ -1 +1 @@
-c0bf63de1c6c738fb994d64e5f2f7c5c0e4315f8627fc5fd68f5b906acbb814d
\ No newline at end of file
+e1d34e03cd7540d476eb93e159a785ea18a80083b3410898a967a56f9d7af5f8
\ No newline at end of file
diff --git a/test/trace_processor/diff_tests/stdlib/android/android_battery_stats_event_slices.out b/test/trace_processor/diff_tests/stdlib/android/android_battery_stats_event_slices.out
index 82cd36c91..fc2247611 100644
--- a/test/trace_processor/diff_tests/stdlib/android/android_battery_stats_event_slices.out
+++ b/test/trace_processor/diff_tests/stdlib/android/android_battery_stats_event_slices.out
@@ -1,4 +1,4 @@
-"ts","dur","track_name","str_value","int_value"
-1000,8000,"battery_stats.top","mail",123
-3000,-1,"battery_stats.job","mail_job",456
-1000,3000,"battery_stats.job","video_job",789
+"ts","dur","safe_dur","track_name","str_value","int_value"
+1000,8000,8000,"battery_stats.top","mail",123
+3000,-1,6000,"battery_stats.job","mail_job",456
+1000,3000,3000,"battery_stats.job","video_job",789
diff --git a/test/trace_processor/diff_tests/stdlib/android/tests.py b/test/trace_processor/diff_tests/stdlib/android/tests.py
index 9542f22cc..59a8877a9 100644
--- a/test/trace_processor/diff_tests/stdlib/android/tests.py
+++ b/test/trace_processor/diff_tests/stdlib/android/tests.py
@@ -106,10 +106,10 @@ class AndroidStdlib(TestSuite):
         ORDER BY ts, track_name;
         """,
         out=Csv("""
-        "ts","dur","track_name","value","value_name"
-        1000,-1,"battery_stats.audio",1,"active"
-        1000,3000,"battery_stats.data_conn",13,"4G (LTE)"
-        4000,-1,"battery_stats.data_conn",20,"5G (NR)"
+        "ts","dur","safe_dur","track_name","value","value_name"
+        1000,-1,3000,"battery_stats.audio",1,"active"
+        1000,3000,3000,"battery_stats.data_conn",13,"4G (LTE)"
+        4000,-1,0,"battery_stats.data_conn",20,"5G (NR)"
         """))
 
   def test_anrs(self):
diff --git a/test/trace_processor/diff_tests/syntax/table_tests.py b/test/trace_processor/diff_tests/syntax/table_tests.py
index 1a3bbb342..7de63e9d4 100644
--- a/test/trace_processor/diff_tests/syntax/table_tests.py
+++ b/test/trace_processor/diff_tests/syntax/table_tests.py
@@ -156,16 +156,16 @@ class PerfettoTable(TestSuite):
         trace=DataPath('example_android_trace_30s.pb'),
         query="""
         WITH trivial_count AS (
-          SELECT DISTINCT name AS c FROM slice
+          SELECT DISTINCT name FROM slice
         ),
         few_results AS (
-          SELECT DISTINCT depth AS c FROM slice
+          SELECT DISTINCT depth FROM slice
         ),
         simple_nullable AS (
-          SELECT DISTINCT parent_id AS c FROM slice
+          SELECT DISTINCT parent_id FROM slice
         ),
         selector AS (
-          SELECT DISTINCT cpu AS c FROM ftrace_event
+          SELECT DISTINCT cpu FROM ftrace_event
         )
         SELECT
           (SELECT COUNT(*) FROM trivial_count) AS name,
@@ -251,3 +251,49 @@ class PerfettoTable(TestSuite):
         8,80
         9,90
         """))
+
+  def test_max(self):
+    return DiffTestBlueprint(
+        trace=DataPath('example_android_trace_30s.pb'),
+        query="""
+        CREATE PERFETTO MACRO max(col ColumnName)
+        RETURNS TableOrSubquery AS (
+          SELECT id, $col
+          FROM slice
+          ORDER BY $col DESC
+          LIMIT 1
+        );
+
+        SELECT
+          (SELECT id FROM max!(id)) AS id,
+          (SELECT id FROM max!(dur)) AS numeric,
+          (SELECT id FROM max!(name)) AS string,
+          (SELECT id FROM max!(parent_id)) AS nullable;
+        """,
+        out=Csv("""
+        "id","numeric","string","nullable"
+        20745,2698,148,20729
+        """))
+
+  def test_min(self):
+    return DiffTestBlueprint(
+        trace=DataPath('example_android_trace_30s.pb'),
+        query="""
+        CREATE PERFETTO MACRO min(col ColumnName)
+        RETURNS TableOrSubquery AS (
+          SELECT id, $col
+          FROM slice
+          ORDER BY $col ASC
+          LIMIT 1
+        );
+
+        SELECT
+          (SELECT id FROM min!(id)) AS id,
+          (SELECT id FROM min!(dur)) AS numeric,
+          (SELECT id FROM min!(name)) AS string,
+          (SELECT id FROM min!(parent_id)) AS nullable;
+        """,
+        out=Csv("""
+        "id","numeric","string","nullable"
+        0,3111,460,0
+        """))
diff --git a/tools/diff_test_trace_processor.py b/tools/diff_test_trace_processor.py
index 531676316..369cde438 100755
--- a/tools/diff_test_trace_processor.py
+++ b/tools/diff_test_trace_processor.py
@@ -41,6 +41,7 @@ def main():
   parser.add_argument('--metrics-descriptor', nargs='+', type=str)
   parser.add_argument('--chrome-track-event-descriptor', type=str, default=None)
   parser.add_argument('--test-extensions', type=str, default=None)
+  parser.add_argument('--winscope-extensions', type=str, default=None)
   parser.add_argument('--perf-file', type=str)
   parser.add_argument(
       '--override-sql-module', type=str, action='append', default=[])
@@ -72,6 +73,9 @@ def main():
   if args.test_extensions is None:
     args.test_extensions = os.path.join(out_path, 'gen', 'protos', 'perfetto',
                                         'trace', 'test_extensions.descriptor')
+  if args.winscope_extensions is None:
+    args.winscope_extensions = os.path.join(out_path, 'gen', 'protos', 'perfetto', 'trace',
+                                            'android', 'winscope.descriptor')
 
   test_runner = DiffTestsRunner(args.name_filter, args.trace_processor,
                                 args.trace_descriptor, args.no_colors,
@@ -79,9 +83,8 @@ def main():
   sys.stderr.write(f"[==========] Running {len(test_runner.tests)} tests.\n")
 
   results = test_runner.run_all_tests(args.metrics_descriptor,
-                                      args.chrome_track_event_descriptor,
-                                      args.test_extensions, args.keep_input,
-                                      args.rebase)
+                                      args.chrome_track_event_descriptor, args.test_extensions,
+                                      args.winscope_extensions, args.keep_input, args.rebase)
   sys.stderr.write(results.str(args.no_colors, len(test_runner.tests)))
 
   if args.rebase:
diff --git a/tools/gen_amalgamated_python_tools b/tools/gen_amalgamated_python_tools
index 4b21069d7..e7ea96793 100755
--- a/tools/gen_amalgamated_python_tools
+++ b/tools/gen_amalgamated_python_tools
@@ -27,6 +27,7 @@ AMALGAMATION_MAP = {
     'python/tools/trace_processor.py': 'tools/trace_processor',
     'python/tools/cpu_profile.py': 'tools/cpu_profile',
     'python/tools/heap_profile.py': 'tools/heap_profile',
+    'python/tools/install_test_reporter_app.py': 'tools/install_test_reporter_app',
 }
 
 
diff --git a/tools/install_test_reporter_app b/tools/install_test_reporter_app
index a2582ef95..3a2d04608 100755
--- a/tools/install_test_reporter_app
+++ b/tools/install_test_reporter_app
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-# Copyright (C) 2023 The Android Open Source Project
+# Copyright (C) 2024 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.
@@ -13,6 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+# DO NOT EDIT. Auto-generated by tools/gen_amalgamated_python_tools
+# !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
 import argparse
 import os
 import re
@@ -20,10 +24,144 @@ import sys
 import tempfile
 import time
 
-ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-sys.path.append(os.path.join(ROOT_DIR))
 
-from python.perfetto.prebuilts.perfetto_prebuilts import *
+# ----- Amalgamator: begin of python/perfetto/prebuilts/perfetto_prebuilts.py
+# 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.
+"""
+Functions to fetch pre-pinned Perfetto prebuilts.
+
+This function is used 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.
+
+The manifest argument looks as follows:
+TRACECONV_MANIFEST = [
+  {
+    'arch': 'mac-amd64',
+    'file_name': 'traceconv',
+    'file_size': 7087080,
+    'url': https://commondatastorage.googleapis.com/.../trace_to_text',
+    'sha256': 7d957c005b0dc130f5bd855d6cec27e060d38841b320d04840afc569f9087490',
+    'platform': 'darwin',
+    'machine': 'x86_64'
+  },
+  ...
+]
+
+The intended usage is:
+
+  from perfetto.prebuilts.manifests.traceconv import TRACECONV_MANIFEST
+  bin_path = get_perfetto_prebuilt(TRACECONV_MANIFEST)
+  subprocess.call(bin_path, ...)
+"""
+
+import hashlib
+import os
+import platform
+import subprocess
+import sys
+import threading
+
+DOWNLOAD_LOCK = threading.Lock()
+
+
+def download_or_get_cached(file_name, url, sha256):
+  """ Downloads a prebuilt or returns a cached version
+
+  The first time this is invoked, it downloads the |url| and caches it into
+  ~/.local/share/perfetto/prebuilts/$tool_name. On subsequent invocations it
+  just runs the cached version.
+  """
+  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
+
+  try:
+    # In BatchTraceProcessor, many threads can be trying to execute the below
+    # code in parallel. For this reason, protect the whole operation with a
+    # lock.
+    DOWNLOAD_LOCK.acquire()
+
+    # 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.replace(tmp_path, bin_path)
+      with open(sha256_path, 'w') as f:
+        f.write(sha256)
+  finally:
+    DOWNLOAD_LOCK.release()
+  return bin_path
+
+
+def get_perfetto_prebuilt(manifest, soft_fail=False, arch=None):
+  """ Downloads the prebuilt, if necessary, and returns its path on disk. """
+  plat = sys.platform.lower()
+  machine = platform.machine().lower()
+  manifest_entry = None
+  for entry in manifest:
+    # If the caller overrides the arch, just match that (for Android prebuilts).
+    if arch:
+      if entry.get('arch') == arch:
+        manifest_entry = entry
+        break
+      continue
+    # Otherwise guess the local machine arch.
+    if entry.get('platform') == plat and machine in entry.get('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'])
+
+
+def run_perfetto_prebuilt(manifest):
+  bin_path = get_perfetto_prebuilt(manifest)
+  if sys.platform.lower() == 'win32':
+    sys.exit(subprocess.check_call([bin_path, *sys.argv[1:]]))
+  os.execv(bin_path, [bin_path] + sys.argv[1:])
+
+# ----- Amalgamator: end of python/perfetto/prebuilts/perfetto_prebuilts.py
 
 PERMSISION_REGEX = re.compile(r'''uses-permission: name='(.*)'.*''')
 NAME_REGEX = re.compile(r'''package: name='(.*?)' .*''')
diff --git a/tools/serialize_test_trace.py b/tools/serialize_test_trace.py
index 6c7fa85d0..f20710143 100755
--- a/tools/serialize_test_trace.py
+++ b/tools/serialize_test_trace.py
@@ -45,8 +45,12 @@ def main():
     trace_descriptor_path = os.path.join(trace_protos_path, 'trace.descriptor')
     test_extensions_descriptor_path = os.path.join(
         trace_protos_path, 'test_extensions.descriptor')
+    winscope_extensions_descriptor_path = os.path.join(
+        trace_protos_path, 'android', 'winscope.descriptor')
     extension_descriptors = [
-        chrome_extension_descriptor_path, test_extensions_descriptor_path
+        chrome_extension_descriptor_path,
+        test_extensions_descriptor_path,
+        winscope_extensions_descriptor_path
     ]
   elif args.descriptor and not args.out:
     trace_descriptor_path = args.descriptor
diff --git a/ui/release/channels.json b/ui/release/channels.json
index c87709e1d..6e49a9282 100644
--- a/ui/release/channels.json
+++ b/ui/release/channels.json
@@ -2,7 +2,7 @@
   "channels": [
     {
       "name": "stable",
-      "rev": "b88536ad5a5569af9559fe0a3f295f12cff37751"
+      "rev": "257a029903d42e76d14a1039edc7de6cf6a10e25"
     },
     {
       "name": "canary",
diff --git a/ui/src/common/actions_unittest.ts b/ui/src/common/actions_unittest.ts
index b31d2ae9c..078440292 100644
--- a/ui/src/common/actions_unittest.ts
+++ b/ui/src/common/actions_unittest.ts
@@ -17,7 +17,6 @@ import {produce} from 'immer';
 import {assertExists} from '../base/logging';
 import {Time} from '../base/time';
 import {PrimaryTrackSortKey} from '../public';
-import {SLICE_TRACK_KIND} from '../tracks/chrome_slices';
 import {HEAP_PROFILE_TRACK_KIND} from '../tracks/heap_profile';
 import {PROCESS_SCHEDULING_TRACK_KIND} from '../tracks/process_summary/process_scheduling_track';
 import {THREAD_STATE_TRACK_KIND} from '../tracks/thread_state';
@@ -32,6 +31,7 @@ import {
   TraceUrlSource,
   TrackSortKey,
 } from './state';
+import {SLICE_TRACK_KIND} from '../tracks/chrome_slices/chrome_slice_track';
 
 function fakeTrack(
   state: State,
diff --git a/ui/src/common/track_helper.ts b/ui/src/common/track_helper.ts
index 0b58e01f1..84808dfcb 100644
--- a/ui/src/common/track_helper.ts
+++ b/ui/src/common/track_helper.ts
@@ -80,6 +80,10 @@ export class TimelineFetcher<Data> implements Disposable {
     return this.data_;
   }
 
+  invalidate() {
+    this.data_ = undefined;
+  }
+
   dispose() {
     this.data_ = undefined;
   }
diff --git a/ui/src/controller/aggregation/slice_aggregation_controller.ts b/ui/src/controller/aggregation/slice_aggregation_controller.ts
index 62085946c..3e18dd040 100644
--- a/ui/src/controller/aggregation/slice_aggregation_controller.ts
+++ b/ui/src/controller/aggregation/slice_aggregation_controller.ts
@@ -17,7 +17,7 @@ import {Area, Sorting} from '../../common/state';
 import {globals} from '../../frontend/globals';
 import {Engine} from '../../trace_processor/engine';
 import {ASYNC_SLICE_TRACK_KIND} from '../../tracks/async_slices';
-import {SLICE_TRACK_KIND} from '../../tracks/chrome_slices';
+import {SLICE_TRACK_KIND} from '../../tracks/chrome_slices/chrome_slice_track';
 
 import {AggregationController} from './aggregation_controller';
 
diff --git a/ui/src/controller/flow_events_controller.ts b/ui/src/controller/flow_events_controller.ts
index d53142c81..183786287 100644
--- a/ui/src/controller/flow_events_controller.ts
+++ b/ui/src/controller/flow_events_controller.ts
@@ -20,7 +20,7 @@ import {publishConnectedFlows, publishSelectedFlows} from '../frontend/publish';
 import {asSliceSqlId} from '../frontend/sql_types';
 import {Engine} from '../trace_processor/engine';
 import {LONG, NUM, STR_NULL} from '../trace_processor/query_result';
-import {SLICE_TRACK_KIND} from '../tracks/chrome_slices';
+import {SLICE_TRACK_KIND} from '../tracks/chrome_slices/chrome_slice_track';
 import {ACTUAL_FRAMES_SLICE_TRACK_KIND} from '../tracks/frames';
 
 import {Controller} from './controller';
diff --git a/ui/src/controller/selection_controller.ts b/ui/src/controller/selection_controller.ts
index d30b458b3..ce9b15f18 100644
--- a/ui/src/controller/selection_controller.ts
+++ b/ui/src/controller/selection_controller.ts
@@ -37,7 +37,7 @@ import {
   STR_NULL,
   timeFromSql,
 } from '../trace_processor/query_result';
-import {SLICE_TRACK_KIND} from '../tracks/chrome_slices';
+import {SLICE_TRACK_KIND} from '../tracks/chrome_slices/chrome_slice_track';
 
 import {Controller} from './controller';
 
diff --git a/ui/src/controller/track_decider.ts b/ui/src/controller/track_decider.ts
index e7e66f9f2..58f77c003 100644
--- a/ui/src/controller/track_decider.ts
+++ b/ui/src/controller/track_decider.ts
@@ -35,7 +35,6 @@ import {
   getScrollJankTracks,
 } from '../tracks/chrome_scroll_jank';
 import {decideTracks as scrollJankDecideTracks} from '../tracks/chrome_scroll_jank/chrome_tasks_scroll_jank_track';
-import {SLICE_TRACK_KIND} from '../tracks/chrome_slices';
 import {COUNTER_TRACK_KIND} from '../tracks/counter';
 import {
   ACTUAL_FRAMES_SLICE_TRACK_KIND,
@@ -43,6 +42,7 @@ import {
 } from '../tracks/frames';
 import {decideTracks as screenshotDecideTracks} from '../tracks/screenshots';
 import {THREAD_STATE_TRACK_KIND} from '../tracks/thread_state';
+import {SLICE_TRACK_KIND} from '../tracks/chrome_slices/chrome_slice_track';
 
 const MEM_DMA_COUNTER_NAME = 'mem.dma_heap';
 const MEM_DMA = 'mem.dma_buffer';
@@ -805,7 +805,7 @@ class TrackDecider {
         t.uid as uid,
         iif(g.cnt = 1, g.package_name, 'UID ' || g.uid) as packageName
       from _uid_track_track_summary_by_uid_and_name t
-      join grouped_packages g using (uid)
+      left join grouped_packages g using (uid)
     `);
 
     const it = result.iter({
@@ -823,7 +823,7 @@ class TrackDecider {
       }
       const rawName = it.name;
       const uid = it.uid === null ? undefined : it.uid;
-      const userName = it.packageName === null ? `UID: ${uid}` : it.packageName;
+      const userName = it.packageName === null ? `UID ${uid}` : it.packageName;
 
       const groupUuid = `uid-track-group${rawName}`;
       if (groupMap.get(rawName) === undefined) {
diff --git a/ui/src/frontend/app.ts b/ui/src/frontend/app.ts
index 3dea6fd21..ddc682625 100644
--- a/ui/src/frontend/app.ts
+++ b/ui/src/frontend/app.ts
@@ -20,7 +20,7 @@ import {findRef} from '../base/dom_utils';
 import {FuzzyFinder} from '../base/fuzzy';
 import {assertExists} from '../base/logging';
 import {undoCommonChatAppReplacements} from '../base/string_utils';
-import {duration, Span, Time, time, TimeSpan} from '../base/time';
+import {duration, Span, time, TimeSpan} from '../base/time';
 import {Actions} from '../common/actions';
 import {getLegacySelection} from '../common/state';
 import {runQuery} from '../common/queries';
@@ -62,6 +62,7 @@ import {
   lockSliceSpan,
   moveByFocusedFlow,
 } from './keyboard_event_handler';
+import {exists} from '../base/utils';
 
 function renderPermalink(): m.Children {
   const permalink = globals.state.permalink;
@@ -1009,7 +1010,7 @@ export class App implements m.ClassComponent {
 // there is no current selection.
 function getTimeSpanOfSelectionOrVisibleWindow(): Span<time, duration> {
   const range = globals.findTimeRangeOfSelection();
-  if (range.end !== Time.INVALID && range.start !== Time.INVALID) {
+  if (exists(range)) {
     return new TimeSpan(range.start, range.end);
   } else {
     return globals.stateVisibleTime();
diff --git a/ui/src/frontend/chrome_slice_details_tab.ts b/ui/src/frontend/chrome_slice_details_tab.ts
index f42d4b60a..49196fa21 100644
--- a/ui/src/frontend/chrome_slice_details_tab.ts
+++ b/ui/src/frontend/chrome_slice_details_tab.ts
@@ -222,7 +222,7 @@ async function getSliceDetails(
   id: number,
   table: string,
 ): Promise<SliceDetails | undefined> {
-  if (table === 'annotation') {
+  if (table === 'annotation_slice') {
     return getAnnotationSlice(engine, id);
   } else {
     return getSlice(engine, asSliceSqlId(id));
diff --git a/ui/src/frontend/globals.ts b/ui/src/frontend/globals.ts
index 68e16a559..50b2286cb 100644
--- a/ui/src/frontend/globals.ts
+++ b/ui/src/frontend/globals.ts
@@ -55,6 +55,7 @@ import {ServiceWorkerController} from './service_worker_controller';
 import {SliceSqlId} from './sql_types';
 import {PxSpan, TimeScale} from './time_scale';
 import {SelectionManager, LegacySelection} from '../core/selection_manager';
+import {exists} from '../base/utils';
 
 const INSTANT_FOCUS_DURATION = 1n;
 const INCOMPLETE_SLICE_DURATION = 30_000n;
@@ -786,46 +787,25 @@ class Globals {
     return Time.sub(ts, this.timestampOffset());
   }
 
-  findTimeRangeOfSelection(): {start: time; end: time} {
+  findTimeRangeOfSelection(): {start: time; end: time} | undefined {
     const selection = getLegacySelection(this.state);
-    let start = Time.INVALID;
-    let end = Time.INVALID;
     if (selection === null) {
-      return {start, end};
-    } else if (
-      selection.kind === 'SLICE' ||
-      selection.kind === 'CHROME_SLICE'
-    ) {
+      return undefined;
+    }
+
+    if (selection.kind === 'SLICE' || selection.kind === 'CHROME_SLICE') {
       const slice = this.sliceDetails;
-      if (slice.ts && slice.dur !== undefined && slice.dur > 0) {
-        start = slice.ts;
-        end = Time.add(start, slice.dur);
-      } else if (slice.ts) {
-        start = slice.ts;
-        // This will handle either:
-        // a)slice.dur === -1 -> unfinished slice
-        // b)slice.dur === 0  -> instant event
-        end =
-          slice.dur === -1n
-            ? Time.add(start, INCOMPLETE_SLICE_DURATION)
-            : Time.add(start, INSTANT_FOCUS_DURATION);
-      }
+      return findTimeRangeOfSlice(slice);
     } else if (selection.kind === 'THREAD_STATE') {
       const threadState = this.threadStateDetails;
-      // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions
-      if (threadState.ts && threadState.dur) {
-        start = threadState.ts;
-        end = Time.add(start, threadState.dur);
-      }
+      return findTimeRangeOfSlice(threadState);
     } else if (selection.kind === 'COUNTER') {
-      start = selection.leftTs;
-      end = selection.rightTs;
+      return {start: selection.leftTs, end: selection.rightTs};
     } else if (selection.kind === 'AREA') {
       const selectedArea = this.state.areas[selection.areaId];
       // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions
       if (selectedArea) {
-        start = selectedArea.start;
-        end = selectedArea.end;
+        return {start: selectedArea.start, end: selectedArea.end};
       }
     } else if (selection.kind === 'NOTE') {
       const selectedNote = this.state.notes[selection.id];
@@ -833,21 +813,21 @@ class Globals {
       // above in the AREA case.
       // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions
       if (selectedNote && selectedNote.noteType === 'DEFAULT') {
-        start = selectedNote.timestamp;
-        end = Time.add(selectedNote.timestamp, INSTANT_FOCUS_DURATION);
+        return {
+          start: selectedNote.timestamp,
+          end: Time.add(selectedNote.timestamp, INSTANT_FOCUS_DURATION),
+        };
       }
     } else if (selection.kind === 'LOG') {
       // TODO(hjd): Make focus selection work for logs.
     } else if (selection.kind === 'GENERIC_SLICE') {
-      start = selection.start;
-      if (selection.duration > 0) {
-        end = Time.add(start, selection.duration);
-      } else {
-        end = Time.add(start, INSTANT_FOCUS_DURATION);
-      }
+      return findTimeRangeOfSlice({
+        ts: selection.start,
+        dur: selection.duration,
+      });
     }
 
-    return {start, end};
+    return undefined;
   }
 
   panToTimestamp(ts: time): void {
@@ -855,4 +835,34 @@ class Globals {
   }
 }
 
+interface SliceLike {
+  ts: time;
+  dur: duration;
+}
+
+// Returns the start and end points of a slice-like object If slice is instant
+// or incomplete, dummy time will be returned which instead.
+function findTimeRangeOfSlice(slice: Partial<SliceLike>): {
+  start: time;
+  end: time;
+} {
+  if (exists(slice.ts) && exists(slice.dur)) {
+    if (slice.dur === -1n) {
+      return {
+        start: slice.ts,
+        end: Time.add(slice.ts, INCOMPLETE_SLICE_DURATION),
+      };
+    } else if (slice.dur === 0n) {
+      return {
+        start: slice.ts,
+        end: Time.add(slice.ts, INSTANT_FOCUS_DURATION),
+      };
+    } else {
+      return {start: slice.ts, end: Time.add(slice.ts, slice.dur)};
+    }
+  } else {
+    return {start: Time.INVALID, end: Time.INVALID};
+  }
+}
+
 export const globals = new Globals();
diff --git a/ui/src/frontend/keyboard_event_handler.ts b/ui/src/frontend/keyboard_event_handler.ts
index 59f1d95c1..8ca508c65 100644
--- a/ui/src/frontend/keyboard_event_handler.ts
+++ b/ui/src/frontend/keyboard_event_handler.ts
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import {exists} from '../base/utils';
 import {Actions} from '../common/actions';
 import {Area, getLegacySelection} from '../common/state';
 
@@ -119,7 +120,7 @@ export function moveByFocusedFlow(direction: Direction): void {
 export function lockSliceSpan(persistent = false) {
   const range = globals.findTimeRangeOfSelection();
   const currentSelection = getLegacySelection(globals.state);
-  if (range.start !== -1n && range.end !== -1n && currentSelection !== null) {
+  if (exists(range) && currentSelection !== null) {
     const tracks = currentSelection.trackKey ? [currentSelection.trackKey] : [];
     const area: Area = {start: range.start, end: range.end, tracks};
     globals.dispatch(Actions.markArea({area, persistent}));
@@ -131,7 +132,7 @@ export function findCurrentSelection() {
   if (selection === null) return;
 
   const range = globals.findTimeRangeOfSelection();
-  if (range.start !== -1n && range.end !== -1n) {
+  if (exists(range)) {
     focusHorizontalRange(range.start, range.end);
   }
 
diff --git a/ui/src/frontend/named_slice_track.ts b/ui/src/frontend/named_slice_track.ts
index 5eff6af74..9d0eb75f7 100644
--- a/ui/src/frontend/named_slice_track.ts
+++ b/ui/src/frontend/named_slice_track.ts
@@ -81,9 +81,6 @@ export abstract class NamedSliceTrack<
         kind: 'CHROME_SLICE',
         id: args.slice.id,
         trackKey: this.trackKey,
-        // |table| here can be either 'slice' or 'annotation'. The
-        // AnnotationSliceTrack overrides the onSliceClick and sets this to
-        // 'annotation'
         table: 'slice',
       },
       {
diff --git a/ui/src/frontend/rpc_http_dialog.ts b/ui/src/frontend/rpc_http_dialog.ts
index 909502331..13db6b6ca 100644
--- a/ui/src/frontend/rpc_http_dialog.ts
+++ b/ui/src/frontend/rpc_http_dialog.ts
@@ -130,7 +130,7 @@ Trace processor RPC API: ${tpStatus.apiVersion}
 //    |  |                                       |No             |Yes        |
 //    |  |                                       |  +---------------------+  |
 //    |  |                                       |  | Dialog: Preloaded?  |  |
-//    |  |                                       +--+ YES, use loaded trace  |
+//    |  |                                       |  + YES, use loaded trace  |
 //    |  |                                 +--------| YES, but reset state|  |
 //    |  |  +---------------------------------------| NO, Use builtin Wasm|  |
 //    |  |  |                              |     |  +---------------------+  |
diff --git a/ui/src/frontend/slice_track.ts b/ui/src/frontend/slice_track.ts
deleted file mode 100644
index e12cd7377..000000000
--- a/ui/src/frontend/slice_track.ts
+++ /dev/null
@@ -1,404 +0,0 @@
-// Copyright (C) 2023 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.
-
-import {duration, Time, time} from '../base/time';
-import {Actions} from '../common/actions';
-import {cropText, drawIncompleteSlice} from '../common/canvas_utils';
-import {getColorForSlice} from '../core/colorizer';
-import {HighPrecisionTime} from '../common/high_precision_time';
-import {TrackData} from '../common/track_data';
-import {TimelineFetcher} from '../common/track_helper';
-import {SliceRect, Track} from '../public';
-import {getLegacySelection} from '../common/state';
-
-import {CROP_INCOMPLETE_SLICE_FLAG} from './base_slice_track';
-import {checkerboardExcept} from './checkerboard';
-import {globals} from './globals';
-import {PanelSize} from './panel';
-
-export const SLICE_TRACK_KIND = 'ChromeSliceTrack';
-const SLICE_HEIGHT = 18;
-const TRACK_PADDING = 2;
-const CHEVRON_WIDTH_PX = 10;
-const HALF_CHEVRON_WIDTH_PX = CHEVRON_WIDTH_PX / 2;
-const INCOMPLETE_SLICE_WIDTH_PX = 20;
-
-export interface SliceData extends TrackData {
-  // Slices are stored in a columnar fashion.
-  strings: string[];
-  sliceIds: Float64Array;
-  starts: BigInt64Array;
-  ends: BigInt64Array;
-  depths: Uint16Array;
-  titles: Uint16Array; // Index into strings.
-  colors?: Uint16Array; // Index into strings.
-  isInstant: Uint16Array;
-  isIncomplete: Uint16Array;
-  cpuTimeRatio?: Float64Array;
-}
-
-// Track base class which handles rendering slices in a generic way.
-// This is the old way of rendering slices - i.e. "track v1" format  - and
-// exists as a patch to allow old tracks to be converted to controller-less
-// tracks before they are ported to v2.
-// Slice tracks should extend this class and implement the abstract methods,
-// notably onBoundsChange().
-// Note: This class is deprecated and should not be used for new tracks. Use
-// |BaseSliceTrack| instead.
-export abstract class SliceTrackLEGACY implements Track {
-  private fetcher = new TimelineFetcher(this.onBoundsChange.bind(this));
-
-  constructor(
-    private maxDepth: number,
-    protected trackKey: string,
-    private tableName: string,
-    private namespace?: string,
-  ) {}
-
-  async onUpdate(): Promise<void> {
-    await this.fetcher.requestDataForCurrentTime();
-  }
-
-  async onDestroy(): Promise<void> {
-    this.fetcher.dispose();
-  }
-
-  abstract onBoundsChange(
-    start: time,
-    end: time,
-    resolution: duration,
-  ): Promise<SliceData>;
-
-  protected namespaceTable(tableName: string = this.tableName): string {
-    if (this.namespace) {
-      return this.namespace + '_' + tableName;
-    } else {
-      return tableName;
-    }
-  }
-
-  private hoveredTitleId = -1;
-
-  // Font used to render the slice name on the current track.
-  protected getFont() {
-    return '12px Roboto Condensed';
-  }
-
-  render(ctx: CanvasRenderingContext2D, size: PanelSize): void {
-    // TODO: fonts and colors should come from the CSS and not hardcoded here.
-    const data = this.fetcher.data;
-    if (data === undefined) return; // Can't possibly draw anything.
-
-    const {visibleTimeSpan, visibleTimeScale} = globals.timeline;
-
-    // If the cached trace slices don't fully cover the visible time range,
-    // show a gray rectangle with a "Loading..." label.
-    checkerboardExcept(
-      ctx,
-      this.getHeight(),
-      0,
-      size.width,
-      visibleTimeScale.timeToPx(data.start),
-      visibleTimeScale.timeToPx(data.end),
-    );
-
-    ctx.textAlign = 'center';
-
-    // measuretext is expensive so we only use it once.
-    const charWidth = ctx.measureText('ACBDLqsdfg').width / 10;
-
-    // The draw of the rect on the selected slice must happen after the other
-    // drawings, otherwise it would result under another rect.
-    let drawRectOnSelected = () => {};
-
-    for (let i = 0; i < data.starts.length; i++) {
-      const tStart = Time.fromRaw(data.starts[i]);
-      let tEnd = Time.fromRaw(data.ends[i]);
-      const depth = data.depths[i];
-      const titleId = data.titles[i];
-      const sliceId = data.sliceIds[i];
-      const isInstant = data.isInstant[i];
-      const isIncomplete = data.isIncomplete[i];
-      const title = data.strings[titleId];
-      const colorOverride = data.colors && data.strings[data.colors[i]];
-      if (isIncomplete) {
-        // incomplete slice
-        // TODO(stevegolton): This isn't exactly equivalent, ideally we should
-        // choose tEnd once we've converted to screen space coords.
-        tEnd = this.getEndTimeIfInComplete(tStart);
-      }
-
-      if (!visibleTimeSpan.intersects(tStart, tEnd)) {
-        continue;
-      }
-
-      const pxEnd = size.width;
-      const left = Math.max(visibleTimeScale.timeToPx(tStart), 0);
-      const right = Math.min(visibleTimeScale.timeToPx(tEnd), pxEnd);
-
-      const rect = {
-        left,
-        width: Math.max(right - left, 1),
-        top: TRACK_PADDING + depth * SLICE_HEIGHT,
-        height: SLICE_HEIGHT,
-      };
-
-      const currentSelection = getLegacySelection(globals.state);
-      const isSelected =
-        currentSelection &&
-        currentSelection.kind === 'CHROME_SLICE' &&
-        currentSelection.id !== undefined &&
-        currentSelection.id === sliceId;
-
-      const highlighted =
-        titleId === this.hoveredTitleId ||
-        globals.state.highlightedSliceId === sliceId;
-
-      const hasFocus = highlighted || isSelected;
-      const colorScheme = getColorForSlice(title);
-      const colorObj = hasFocus ? colorScheme.variant : colorScheme.base;
-      const textColor = hasFocus
-        ? colorScheme.textVariant
-        : colorScheme.textBase;
-
-      let color: string;
-      if (colorOverride === undefined) {
-        color = colorObj.cssString;
-      } else {
-        color = colorOverride;
-      }
-      ctx.fillStyle = color;
-
-      // We draw instant events as upward facing chevrons starting at A:
-      //     A
-      //    ###
-      //   ##C##
-      //  ##   ##
-      // D       B
-      // Then B, C, D and back to A:
-      if (isInstant) {
-        if (isSelected) {
-          drawRectOnSelected = () => {
-            ctx.save();
-            ctx.translate(rect.left, rect.top);
-
-            // Draw a rectangle around the selected slice
-            ctx.strokeStyle = colorObj.setHSL({s: 100, l: 10}).cssString;
-            ctx.beginPath();
-            ctx.lineWidth = 3;
-            ctx.strokeRect(
-              -HALF_CHEVRON_WIDTH_PX,
-              0,
-              CHEVRON_WIDTH_PX,
-              SLICE_HEIGHT,
-            );
-            ctx.closePath();
-
-            // Draw inner chevron as interior
-            ctx.fillStyle = color;
-            this.drawChevron(ctx);
-
-            ctx.restore();
-          };
-        } else {
-          ctx.save();
-          ctx.translate(rect.left, rect.top);
-          this.drawChevron(ctx);
-          ctx.restore();
-        }
-        continue;
-      }
-
-      if (isIncomplete && rect.width > SLICE_HEIGHT / 4) {
-        drawIncompleteSlice(
-          ctx,
-          rect.left,
-          rect.top,
-          rect.width,
-          SLICE_HEIGHT,
-          !CROP_INCOMPLETE_SLICE_FLAG.get(),
-        );
-      } else if (
-        data.cpuTimeRatio !== undefined &&
-        data.cpuTimeRatio[i] < 1 - 1e-9
-      ) {
-        // We draw two rectangles, representing the ratio between wall time and
-        // time spent on cpu.
-        const cpuTimeRatio = data.cpuTimeRatio![i];
-        const firstPartWidth = rect.width * cpuTimeRatio;
-        const secondPartWidth = rect.width * (1 - cpuTimeRatio);
-        ctx.fillRect(rect.left, rect.top, rect.width, SLICE_HEIGHT);
-        ctx.fillStyle = '#FFFFFF50';
-        ctx.fillRect(
-          rect.left + firstPartWidth,
-          rect.top,
-          secondPartWidth,
-          SLICE_HEIGHT,
-        );
-      } else {
-        ctx.fillRect(rect.left, rect.top, rect.width, SLICE_HEIGHT);
-      }
-
-      // Selected case
-      if (isSelected) {
-        drawRectOnSelected = () => {
-          ctx.strokeStyle = colorObj.setHSL({s: 100, l: 10}).cssString;
-          ctx.beginPath();
-          ctx.lineWidth = 3;
-          ctx.strokeRect(
-            rect.left,
-            rect.top - 1.5,
-            rect.width,
-            SLICE_HEIGHT + 3,
-          );
-          ctx.closePath();
-        };
-      }
-
-      // Don't render text when we have less than 5px to play with.
-      if (rect.width >= 5) {
-        ctx.fillStyle = textColor.cssString;
-        const displayText = cropText(title, charWidth, rect.width);
-        const rectXCenter = rect.left + rect.width / 2;
-        ctx.textBaseline = 'middle';
-        ctx.font = this.getFont();
-        ctx.fillText(displayText, rectXCenter, rect.top + SLICE_HEIGHT / 2);
-      }
-    }
-    drawRectOnSelected();
-  }
-
-  drawChevron(ctx: CanvasRenderingContext2D) {
-    // Draw a chevron at a fixed location and size. Should be used with
-    // ctx.translate and ctx.scale to alter location and size.
-    ctx.beginPath();
-    ctx.moveTo(0, 0);
-    ctx.lineTo(HALF_CHEVRON_WIDTH_PX, SLICE_HEIGHT);
-    ctx.lineTo(0, SLICE_HEIGHT - HALF_CHEVRON_WIDTH_PX);
-    ctx.lineTo(-HALF_CHEVRON_WIDTH_PX, SLICE_HEIGHT);
-    ctx.lineTo(0, 0);
-    ctx.fill();
-  }
-
-  getSliceIndex({x, y}: {x: number; y: number}): number | void {
-    const data = this.fetcher.data;
-    if (data === undefined) return;
-    const {visibleTimeScale: timeScale} = globals.timeline;
-    if (y < TRACK_PADDING) return;
-    const instantWidthTime = timeScale.pxDeltaToDuration(HALF_CHEVRON_WIDTH_PX);
-    const t = timeScale.pxToHpTime(x);
-    const depth = Math.floor((y - TRACK_PADDING) / SLICE_HEIGHT);
-
-    for (let i = 0; i < data.starts.length; i++) {
-      if (depth !== data.depths[i]) {
-        continue;
-      }
-      const start = Time.fromRaw(data.starts[i]);
-      const tStart = HighPrecisionTime.fromTime(start);
-      if (data.isInstant[i]) {
-        if (tStart.sub(t).abs().lt(instantWidthTime)) {
-          return i;
-        }
-      } else {
-        const end = Time.fromRaw(data.ends[i]);
-        let tEnd = HighPrecisionTime.fromTime(end);
-        if (data.isIncomplete[i]) {
-          const endTime = this.getEndTimeIfInComplete(start);
-          tEnd = HighPrecisionTime.fromTime(endTime);
-        }
-        if (tStart.lte(t) && t.lte(tEnd)) {
-          return i;
-        }
-      }
-    }
-  }
-
-  getEndTimeIfInComplete(start: time): time {
-    const {visibleTimeScale, visibleWindowTime} = globals.timeline;
-
-    let end = visibleWindowTime.end.toTime('ceil');
-    if (CROP_INCOMPLETE_SLICE_FLAG.get()) {
-      const widthTime = visibleTimeScale
-        .pxDeltaToDuration(INCOMPLETE_SLICE_WIDTH_PX)
-        .toTime();
-      end = Time.add(start, widthTime);
-    }
-
-    return end;
-  }
-
-  onMouseMove({x, y}: {x: number; y: number}) {
-    this.hoveredTitleId = -1;
-    globals.dispatch(Actions.setHighlightedSliceId({sliceId: -1}));
-    const sliceIndex = this.getSliceIndex({x, y});
-    if (sliceIndex === undefined) return;
-    const data = this.fetcher.data;
-    if (data === undefined) return;
-    this.hoveredTitleId = data.titles[sliceIndex];
-    const sliceId = data.sliceIds[sliceIndex];
-    globals.dispatch(Actions.setHighlightedSliceId({sliceId}));
-  }
-
-  onMouseOut() {
-    this.hoveredTitleId = -1;
-    globals.dispatch(Actions.setHighlightedSliceId({sliceId: -1}));
-  }
-
-  onMouseClick({x, y}: {x: number; y: number}): boolean {
-    const sliceIndex = this.getSliceIndex({x, y});
-    if (sliceIndex === undefined) return false;
-    const data = this.fetcher.data;
-    if (data === undefined) return false;
-    const sliceId = data.sliceIds[sliceIndex];
-    if (sliceId !== undefined && sliceId !== -1) {
-      globals.setLegacySelection(
-        {
-          kind: 'CHROME_SLICE',
-          id: sliceId,
-          trackKey: this.trackKey,
-          table: this.namespace,
-        },
-        {
-          clearSearch: true,
-          pendingScrollId: undefined,
-          switchToCurrentSelectionTab: true,
-        },
-      );
-      return true;
-    }
-    return false;
-  }
-
-  getHeight() {
-    return SLICE_HEIGHT * (this.maxDepth + 1) + 2 * TRACK_PADDING;
-  }
-
-  getSliceRect(tStart: time, tEnd: time, depth: number): SliceRect | undefined {
-    const {windowSpan, visibleTimeScale, visibleTimeSpan} = globals.timeline;
-
-    const pxEnd = windowSpan.end;
-    const left = Math.max(visibleTimeScale.timeToPx(tStart), 0);
-    const right = Math.min(visibleTimeScale.timeToPx(tEnd), pxEnd);
-
-    const visible = visibleTimeSpan.intersects(tStart, tEnd);
-
-    return {
-      left,
-      width: Math.max(right - left, 1),
-      top: TRACK_PADDING + depth * SLICE_HEIGHT,
-      height: SLICE_HEIGHT,
-      visible,
-    };
-  }
-}
diff --git a/ui/src/plugins/dev.perfetto.AndroidLongBatteryTracing/index.ts b/ui/src/plugins/dev.perfetto.AndroidLongBatteryTracing/index.ts
index 414a6620e..2a22affc5 100644
--- a/ui/src/plugins/dev.perfetto.AndroidLongBatteryTracing/index.ts
+++ b/ui/src/plugins/dev.perfetto.AndroidLongBatteryTracing/index.ts
@@ -51,6 +51,33 @@ const DEFAULT_NETWORK = `
       end as name
   from diff where keep is null or keep`;
 
+const RADIO_TRANSPORT_TYPE = `
+  create or replace perfetto view radio_transport_data_conn as
+  select ts, safe_dur AS dur, value_name as data_conn, value AS data_conn_val
+  from android_battery_stats_state
+  where track_name = "battery_stats.data_conn";
+
+  create or replace perfetto view radio_transport_nr_state as
+  select ts, safe_dur AS dur, value AS nr_state_val
+  from android_battery_stats_state
+  where track_name = "battery_stats.nr_state";
+
+  drop table if exists radio_transport_join;
+  create virtual table radio_transport_join
+  using span_left_join(radio_transport_data_conn, radio_transport_nr_state);
+
+  create or replace perfetto view radio_transport as
+  select
+    ts, dur,
+    case data_conn_val
+      -- On LTE with NR connected is 5G NSA.
+      when 13 then iif(nr_state_val = 3, '5G (NSA)', data_conn)
+      -- On NR with NR state present, is 5G SA.
+      when 20 then iif(nr_state_val is null, '5G (SA or NSA)', '5G (SA)')
+      else data_conn
+    end as name
+  from radio_transport_join;`;
+
 const TETHERING = `
   with base as (
       select
@@ -1144,7 +1171,7 @@ class AndroidLongBatteryTracing implements Plugin {
     this.addSliceTrack(
       ctx,
       name,
-      `SELECT ts, dur, value_name AS name
+      `SELECT ts, safe_dur AS dur, value_name AS name
     FROM android_battery_stats_state
     WHERE track_name = "${track}"`,
       groupName,
@@ -1165,7 +1192,7 @@ class AndroidLongBatteryTracing implements Plugin {
     this.addSliceTrack(
       ctx,
       name,
-      `SELECT ts, dur, str_value AS name
+      `SELECT ts, safe_dur AS dur, str_value AS name
     FROM android_battery_stats_event_slices
     WHERE track_name = "${track}"`,
       groupName,
@@ -1201,7 +1228,7 @@ class AndroidLongBatteryTracing implements Plugin {
       'Device State: Long wakelocks',
       `SELECT
             ts - 60000000000 as ts,
-            dur + 60000000000 as dur,
+            safe_dur + 60000000000 as dur,
             str_value AS name,
             ifnull(
             (select package_name from package_list where uid = int_value % 100000),
@@ -1236,6 +1263,7 @@ class AndroidLongBatteryTracing implements Plugin {
 
     const e = ctx.engine;
     await e.query(NETWORK_SUMMARY);
+    await e.query(RADIO_TRANSPORT_TYPE);
 
     this.addSliceTrack(ctx, 'Default network', DEFAULT_NETWORK, groupName);
 
@@ -1295,12 +1323,11 @@ class AndroidLongBatteryTracing implements Plugin {
       groupName,
       features,
     );
-    this.addBatteryStatsState(
+    this.addSliceTrack(
       ctx,
       'Cellular connection',
-      'battery_stats.data_conn',
+      `select ts, dur, name from radio_transport`,
       groupName,
-      features,
     );
     this.addBatteryStatsState(
       ctx,
diff --git a/ui/src/tracks/annotation/index.ts b/ui/src/tracks/annotation/index.ts
index d36e9a20d..bda66f3fd 100644
--- a/ui/src/tracks/annotation/index.ts
+++ b/ui/src/tracks/annotation/index.ts
@@ -13,8 +13,11 @@
 // limitations under the License.
 
 import {Plugin, PluginContextTrace, PluginDescriptor} from '../../public';
+import {
+  ChromeSliceTrack,
+  SLICE_TRACK_KIND,
+} from '../chrome_slices/chrome_slice_track';
 import {NUM, NUM_NULL, STR} from '../../trace_processor/query_result';
-import {ChromeSliceTrack, SLICE_TRACK_KIND} from '../chrome_slices/';
 import {COUNTER_TRACK_KIND, TraceProcessorCounterTrack} from '../counter';
 
 class AnnotationPlugin implements Plugin {
@@ -49,7 +52,15 @@ class AnnotationPlugin implements Plugin {
           metric: true,
         },
         trackFactory: ({trackKey}) => {
-          return new ChromeSliceTrack(engine, 0, trackKey, id, 'annotation');
+          return new ChromeSliceTrack(
+            {
+              engine: ctx.engine,
+              trackKey,
+            },
+            id,
+            0,
+            'annotation_slice',
+          );
         },
       });
     }
diff --git a/ui/src/tracks/async_slices/index.ts b/ui/src/tracks/async_slices/index.ts
index 49520a031..644d22f24 100644
--- a/ui/src/tracks/async_slices/index.ts
+++ b/ui/src/tracks/async_slices/index.ts
@@ -151,7 +151,7 @@ class AsyncSlicePlugin implements Plugin {
         __max_layout_depth(t.track_count, t.track_ids) as maxDepth,
         iif(g.cnt = 1, g.package_name, 'UID ' || g.uid) as packageName
       from _uid_track_track_summary_by_uid_and_name t
-      join grouped_packages g using (uid)
+      left join grouped_packages g using (uid)
     `);
 
     const it = result.iter({
@@ -165,8 +165,8 @@ class AsyncSlicePlugin implements Plugin {
     for (; it.valid(); it.next()) {
       const kind = ASYNC_SLICE_TRACK_KIND;
       const rawName = it.name === null ? undefined : it.name;
-      const userName = it.packageName === null ? undefined : it.packageName;
       const uid = it.uid === null ? undefined : it.uid;
+      const userName = it.packageName === null ? `UID ${uid}` : it.packageName;
       const rawTrackIds = it.trackIds;
       const trackIds = rawTrackIds.split(',').map((v) => Number(v));
       const maxDepth = it.maxDepth;
diff --git a/ui/src/tracks/chrome_slices/chrome_slice_track.ts b/ui/src/tracks/chrome_slices/chrome_slice_track.ts
new file mode 100644
index 000000000..c9fda8bff
--- /dev/null
+++ b/ui/src/tracks/chrome_slices/chrome_slice_track.ts
@@ -0,0 +1,109 @@
+// Copyright (C) 2024 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.
+
+import {BigintMath as BIMath} from '../../base/bigint_math';
+import {clamp} from '../../base/math_utils';
+import {OnSliceClickArgs} from '../../frontend/base_slice_track';
+import {globals} from '../../frontend/globals';
+import {
+  NAMED_ROW,
+  NamedSliceTrack,
+  NamedSliceTrackTypes,
+} from '../../frontend/named_slice_track';
+import {SLICE_LAYOUT_FIT_CONTENT_DEFAULTS} from '../../frontend/slice_layout';
+import {NewTrackArgs} from '../../frontend/track';
+import {LONG_NULL} from '../../trace_processor/query_result';
+
+export const SLICE_TRACK_KIND = 'ChromeSliceTrack';
+
+export const CHROME_SLICE_ROW = {
+  // Base columns (tsq, ts, dur, id, depth).
+  ...NAMED_ROW,
+
+  // Chrome-specific columns.
+  threadDur: LONG_NULL,
+};
+export type ChromeSliceRow = typeof CHROME_SLICE_ROW;
+
+export interface ChromeSliceTrackTypes extends NamedSliceTrackTypes {
+  row: ChromeSliceRow;
+}
+
+export class ChromeSliceTrack extends NamedSliceTrack<ChromeSliceTrackTypes> {
+  constructor(
+    args: NewTrackArgs,
+    private trackId: number,
+    maxDepth: number,
+    private tableName: string = 'slice',
+  ) {
+    super(args);
+    this.sliceLayout = {
+      ...SLICE_LAYOUT_FIT_CONTENT_DEFAULTS,
+      depthGuess: maxDepth,
+    };
+  }
+
+  // This is used by the base class to call iter().
+  getRowSpec() {
+    return CHROME_SLICE_ROW;
+  }
+
+  getSqlSource(): string {
+    return `select
+      ts,
+      dur,
+      id,
+      depth,
+      ifnull(name, '') as name,
+      thread_dur as threadDur
+    from ${this.tableName}
+    where track_id = ${this.trackId}`;
+  }
+
+  // Converts a SQL result row to an "Impl" Slice.
+  rowToSlice(
+    row: ChromeSliceTrackTypes['row'],
+  ): ChromeSliceTrackTypes['slice'] {
+    const namedSlice = super.rowToSlice(row);
+
+    if (row.dur > 0n && row.threadDur !== null) {
+      const fillRatio = clamp(BIMath.ratio(row.threadDur, row.dur), 0, 1);
+      return {...namedSlice, fillRatio};
+    } else {
+      return namedSlice;
+    }
+  }
+
+  onUpdatedSlices(slices: ChromeSliceTrackTypes['slice'][]) {
+    for (const slice of slices) {
+      slice.isHighlighted = slice === this.hoveredSlice;
+    }
+  }
+
+  onSliceClick(args: OnSliceClickArgs<ChromeSliceTrackTypes['slice']>) {
+    globals.setLegacySelection(
+      {
+        kind: 'CHROME_SLICE',
+        id: args.slice.id,
+        trackKey: this.trackKey,
+        table: this.tableName,
+      },
+      {
+        clearSearch: true,
+        pendingScrollId: undefined,
+        switchToCurrentSelectionTab: true,
+      },
+    );
+  }
+}
diff --git a/ui/src/tracks/chrome_slices/generic_slice_track.ts b/ui/src/tracks/chrome_slices/generic_slice_track.ts
deleted file mode 100644
index 1e817a894..000000000
--- a/ui/src/tracks/chrome_slices/generic_slice_track.ts
+++ /dev/null
@@ -1,30 +0,0 @@
-// 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.
-
-import {
-  NamedSliceTrack,
-  NamedSliceTrackTypes,
-} from '../../frontend/named_slice_track';
-import {NewTrackArgs} from '../../frontend/track';
-
-export class GenericSliceTrack extends NamedSliceTrack<NamedSliceTrackTypes> {
-  constructor(args: NewTrackArgs, private sqlTrackId: number) {
-    super(args);
-  }
-
-  getSqlSource(): string {
-    return `select ts, dur, id, depth, ifnull(name, '') as name
-    from slice where track_id = ${this.sqlTrackId}`;
-  }
-}
diff --git a/ui/src/tracks/chrome_slices/index.ts b/ui/src/tracks/chrome_slices/index.ts
index cae255728..830b7e4e9 100644
--- a/ui/src/tracks/chrome_slices/index.ts
+++ b/ui/src/tracks/chrome_slices/index.ts
@@ -12,213 +12,17 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import {BigintMath as BIMath} from '../../base/bigint_math';
-import {clamp} from '../../base/math_utils';
-import {Duration, duration, time} from '../../base/time';
 import {uuidv4} from '../../base/uuid';
 import {ChromeSliceDetailsTab} from '../../frontend/chrome_slice_details_tab';
 import {
-  NAMED_ROW,
-  NamedSliceTrack,
-  NamedSliceTrackTypes,
-} from '../../frontend/named_slice_track';
-import {SLICE_LAYOUT_FIT_CONTENT_DEFAULTS} from '../../frontend/slice_layout';
-import {SliceData, SliceTrackLEGACY} from '../../frontend/slice_track';
-import {NewTrackArgs} from '../../frontend/track';
-import {
   BottomTabToSCSAdapter,
-  EngineProxy,
   Plugin,
   PluginContextTrace,
   PluginDescriptor,
 } from '../../public';
 import {getTrackName} from '../../public/utils';
-import {
-  LONG,
-  LONG_NULL,
-  NUM,
-  NUM_NULL,
-  STR,
-  STR_NULL,
-} from '../../trace_processor/query_result';
-
-export const SLICE_TRACK_KIND = 'ChromeSliceTrack';
-
-export class ChromeSliceTrack extends SliceTrackLEGACY {
-  private maxDurNs: duration = 0n;
-
-  constructor(
-    protected engine: EngineProxy,
-    maxDepth: number,
-    trackKey: string,
-    private trackId: number,
-    namespace?: string,
-  ) {
-    super(maxDepth, trackKey, 'slice', namespace);
-  }
-
-  async onBoundsChange(
-    start: time,
-    end: time,
-    resolution: duration,
-  ): Promise<SliceData> {
-    const tableName = this.namespaceTable('slice');
-
-    if (this.maxDurNs === Duration.ZERO) {
-      const query = `
-          SELECT max(iif(dur = -1, (SELECT end_ts FROM trace_bounds) - ts, dur))
-          AS maxDur FROM ${tableName} WHERE track_id = ${this.trackId}`;
-      const queryRes = await this.engine.query(query);
-      this.maxDurNs = queryRes.firstRow({maxDur: LONG_NULL}).maxDur ?? 0n;
-    }
-
-    const query = `
-      SELECT
-        (ts + ${resolution / 2n}) / ${resolution} * ${resolution} as tsq,
-        ts,
-        max(iif(dur = -1, (SELECT end_ts FROM trace_bounds) - ts, dur)) as dur,
-        depth,
-        id as sliceId,
-        ifnull(name, '[null]') as name,
-        dur = 0 as isInstant,
-        dur = -1 as isIncomplete,
-        thread_dur as threadDur
-      FROM ${tableName}
-      WHERE track_id = ${this.trackId} AND
-        ts >= (${start - this.maxDurNs}) AND
-        ts <= ${end}
-      GROUP BY depth, tsq`;
-    const queryRes = await this.engine.query(query);
-
-    const numRows = queryRes.numRows();
-    const slices: SliceData = {
-      start,
-      end,
-      resolution,
-      length: numRows,
-      strings: [],
-      sliceIds: new Float64Array(numRows),
-      starts: new BigInt64Array(numRows),
-      ends: new BigInt64Array(numRows),
-      depths: new Uint16Array(numRows),
-      titles: new Uint16Array(numRows),
-      isInstant: new Uint16Array(numRows),
-      isIncomplete: new Uint16Array(numRows),
-      cpuTimeRatio: new Float64Array(numRows),
-    };
-
-    const stringIndexes = new Map<string, number>();
-    function internString(str: string) {
-      let idx = stringIndexes.get(str);
-      if (idx !== undefined) return idx;
-      idx = slices.strings.length;
-      slices.strings.push(str);
-      stringIndexes.set(str, idx);
-      return idx;
-    }
-
-    const it = queryRes.iter({
-      tsq: LONG,
-      ts: LONG,
-      dur: LONG,
-      depth: NUM,
-      sliceId: NUM,
-      name: STR,
-      isInstant: NUM,
-      isIncomplete: NUM,
-      threadDur: LONG_NULL,
-    });
-    for (let row = 0; it.valid(); it.next(), row++) {
-      const startQ = it.tsq;
-      const start = it.ts;
-      const dur = it.dur;
-      const end = start + dur;
-      const minEnd = startQ + resolution;
-      const endQ = BIMath.max(BIMath.quant(end, resolution), minEnd);
-
-      slices.starts[row] = startQ;
-      slices.ends[row] = endQ;
-      slices.depths[row] = it.depth;
-      slices.sliceIds[row] = it.sliceId;
-      slices.titles[row] = internString(it.name);
-      slices.isInstant[row] = it.isInstant;
-      slices.isIncomplete[row] = it.isIncomplete;
-
-      let cpuTimeRatio = 1;
-      if (!it.isInstant && !it.isIncomplete && it.threadDur !== null) {
-        // Rounding the CPU time ratio to two decimal places and ensuring
-        // it is less than or equal to one, incase the thread duration exceeds
-        // the total duration.
-        cpuTimeRatio = Math.min(
-          Math.round(BIMath.ratio(it.threadDur, it.dur) * 100) / 100,
-          1,
-        );
-      }
-      slices.cpuTimeRatio![row] = cpuTimeRatio;
-    }
-    return slices;
-  }
-}
-
-export const CHROME_SLICE_ROW = {
-  // Base columns (tsq, ts, dur, id, depth).
-  ...NAMED_ROW,
-
-  // Chrome-specific columns.
-  threadDur: LONG_NULL,
-};
-export type ChromeSliceRow = typeof CHROME_SLICE_ROW;
-
-export interface ChromeSliceTrackTypes extends NamedSliceTrackTypes {
-  row: ChromeSliceRow;
-}
-
-export class ChromeSliceTrackV2 extends NamedSliceTrack<ChromeSliceTrackTypes> {
-  constructor(args: NewTrackArgs, private trackId: number, maxDepth: number) {
-    super(args);
-    this.sliceLayout = {
-      ...SLICE_LAYOUT_FIT_CONTENT_DEFAULTS,
-      depthGuess: maxDepth,
-    };
-  }
-
-  // This is used by the base class to call iter().
-  getRowSpec() {
-    return CHROME_SLICE_ROW;
-  }
-
-  getSqlSource(): string {
-    return `select
-      ts,
-      dur,
-      id,
-      depth,
-      ifnull(name, '') as name,
-      thread_dur as threadDur
-    from slice
-    where track_id = ${this.trackId}`;
-  }
-
-  // Converts a SQL result row to an "Impl" Slice.
-  rowToSlice(
-    row: ChromeSliceTrackTypes['row'],
-  ): ChromeSliceTrackTypes['slice'] {
-    const namedSlice = super.rowToSlice(row);
-
-    if (row.dur > 0n && row.threadDur !== null) {
-      const fillRatio = clamp(BIMath.ratio(row.threadDur, row.dur), 0, 1);
-      return {...namedSlice, fillRatio};
-    } else {
-      return namedSlice;
-    }
-  }
-
-  onUpdatedSlices(slices: ChromeSliceTrackTypes['slice'][]) {
-    for (const slice of slices) {
-      slice.isHighlighted = slice === this.hoveredSlice;
-    }
-  }
-}
+import {NUM, NUM_NULL, STR_NULL} from '../../trace_processor/query_result';
+import {ChromeSliceTrack, SLICE_TRACK_KIND} from './chrome_slice_track';
 
 class ChromeSlicesPlugin implements Plugin {
   async onTraceLoad(ctx: PluginContextTrace): Promise<void> {
@@ -281,7 +85,7 @@ class ChromeSlicesPlugin implements Plugin {
             engine: ctx.engine,
             trackKey,
           };
-          return new ChromeSliceTrackV2(newTrackArgs, trackId, maxDepth);
+          return new ChromeSliceTrack(newTrackArgs, trackId, maxDepth);
         },
       });
     }
diff --git a/ui/src/tracks/ftrace/ftrace_track.ts b/ui/src/tracks/ftrace/ftrace_track.ts
index 7b9def953..6950eca4f 100644
--- a/ui/src/tracks/ftrace/ftrace_track.ts
+++ b/ui/src/tracks/ftrace/ftrace_track.ts
@@ -24,6 +24,7 @@ import {EngineProxy, Track} from '../../public';
 import {LONG, STR} from '../../trace_processor/query_result';
 import {FtraceFilter} from './common';
 import {Store} from '../../public';
+import {Monitor} from '../../base/monitor';
 
 const MARGIN = 2;
 const RECT_HEIGHT = 18;
@@ -43,14 +44,20 @@ export class FtraceRawTrack implements Track {
   private engine: EngineProxy;
   private cpu: number;
   private store: Store<FtraceFilter>;
+  private readonly monitor: Monitor;
 
   constructor(engine: EngineProxy, cpu: number, store: Store<FtraceFilter>) {
     this.engine = engine;
     this.cpu = cpu;
     this.store = store;
+
+    this.monitor = new Monitor([() => store.state]);
   }
 
   async onUpdate(): Promise<void> {
+    this.monitor.ifStateChanged(() => {
+      this.fetcher.invalidate();
+    });
     await this.fetcher.requestDataForCurrentTime();
   }
 
diff --git a/ui/src/tracks/visualised_args/index.ts b/ui/src/tracks/visualised_args/index.ts
index 9e34de3f1..459a7ad4a 100644
--- a/ui/src/tracks/visualised_args/index.ts
+++ b/ui/src/tracks/visualised_args/index.ts
@@ -12,98 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import m from 'mithril';
-import {v4 as uuidv4} from 'uuid';
-
-import {Actions} from '../../common/actions';
-import {globals} from '../../frontend/globals';
-import {
-  EngineProxy,
-  Plugin,
-  PluginContextTrace,
-  PluginDescriptor,
-  TrackContext,
-} from '../../public';
-import {ChromeSliceTrack} from '../chrome_slices';
+import {Plugin, PluginContextTrace, PluginDescriptor} from '../../public';
 import {
   VISUALISED_ARGS_SLICE_TRACK_URI,
   VisualisedArgsState,
 } from '../../frontend/visualized_args_tracks';
-import {Button} from '../../widgets/button';
-import {Icons} from '../../base/semantic_icons';
-
-export class VisualisedArgsTrack extends ChromeSliceTrack {
-  private helperViewName: string;
-
-  constructor(
-    engine: EngineProxy,
-    maxDepth: number,
-    trackKey: string,
-    trackId: number,
-    private argName: string,
-  ) {
-    const uuid = uuidv4();
-    const namespace = `__arg_visualisation_helper_${argName}_${uuid}`;
-    const escapedNamespace = namespace.replace(/[^a-zA-Z]/g, '_');
-    super(engine, maxDepth, trackKey, trackId, escapedNamespace);
-    this.helperViewName = `${escapedNamespace}_slice`;
-  }
-
-  async onCreate(_ctx: TrackContext): Promise<void> {
-    // Create the helper view - just one which is relevant to this slice
-    await this.engine.query(`
-          create view ${this.helperViewName} as
-          with slice_with_arg as (
-            select
-              slice.id,
-              slice.track_id,
-              slice.ts,
-              slice.dur,
-              slice.thread_dur,
-              NULL as cat,
-              args.display_value as name
-            from slice
-            join args using (arg_set_id)
-            where args.key='${this.argName}'
-          )
-          select
-            *,
-            (select count()
-            from ancestor_slice(s1.id) s2
-            join slice_with_arg s3 on s2.id=s3.id
-            ) as depth
-          from slice_with_arg s1
-          order by id;
-      `);
-  }
-
-  async onDestroy(): Promise<void> {
-    if (this.engine.isAlive) {
-      await this.engine.query(`drop view ${this.helperViewName}`);
-    }
-  }
-
-  getFont() {
-    return 'italic 11px Roboto';
-  }
-
-  getTrackShellButtons(): m.Children {
-    return m(Button, {
-      onclick: () => {
-        // This behavior differs to the original behavior a little.
-        // Originally, hitting the close button on a single track removed ALL
-        // tracks with this argName, whereas this one only closes the single
-        // track.
-        // This will be easily fixable once we transition to using dynamic
-        // tracks instead of this "initial state" approach to add these tracks.
-        globals.dispatch(Actions.removeTracks({trackKeys: [this.trackKey]}));
-      },
-      icon: Icons.Close,
-      title: 'Close',
-      compact: true,
-    });
-  }
-}
+import {VisualisedArgsTrack} from './visualized_args_track';
 
 class VisualisedArgsPlugin implements Plugin {
   async onTraceLoad(ctx: PluginContextTrace): Promise<void> {
@@ -117,10 +31,12 @@ class VisualisedArgsPlugin implements Plugin {
         // worse than the situation we had before with track config.
         const params = trackCtx.params as VisualisedArgsState;
         return new VisualisedArgsTrack(
-          ctx.engine,
-          params.maxDepth,
-          trackCtx.trackKey,
+          {
+            engine: ctx.engine,
+            trackKey: trackCtx.trackKey,
+          },
           params.trackId,
+          params.maxDepth,
           params.argName,
         );
       },
diff --git a/ui/src/tracks/visualised_args/visualized_args_track.ts b/ui/src/tracks/visualised_args/visualized_args_track.ts
new file mode 100644
index 000000000..8de85d34f
--- /dev/null
+++ b/ui/src/tracks/visualised_args/visualized_args_track.ts
@@ -0,0 +1,93 @@
+// Copyright (C) 2024 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.
+
+import m from 'mithril';
+
+import {Actions} from '../../common/actions';
+import {globals} from '../../frontend/globals';
+import {Button} from '../../widgets/button';
+import {Icons} from '../../base/semantic_icons';
+import {ChromeSliceTrack} from '../chrome_slices/chrome_slice_track';
+import {uuidv4Sql} from '../../base/uuid';
+import {NewTrackArgs} from '../../frontend/track';
+import {Disposable, DisposableCallback} from '../../base/disposable';
+
+// Similar to a SliceTrack, but creates a view
+export class VisualisedArgsTrack extends ChromeSliceTrack {
+  private viewName: string;
+
+  constructor(
+    args: NewTrackArgs,
+    trackId: number,
+    maxDepth: number,
+    private argName: string,
+  ) {
+    const uuid = uuidv4Sql();
+    const escapedArgName = argName.replace(/[^a-zA-Z]/g, '_');
+    const viewName = `__arg_visualisation_helper_${escapedArgName}_${uuid}_slice`;
+    super(args, trackId, maxDepth, viewName);
+    this.viewName = viewName;
+  }
+
+  async onInit(): Promise<Disposable> {
+    // Create the helper view - just one which is relevant to this slice
+    await this.engine.query(`
+        create view ${this.viewName} as
+        with slice_with_arg as (
+          select
+            slice.id,
+            slice.track_id,
+            slice.ts,
+            slice.dur,
+            slice.thread_dur,
+            NULL as cat,
+            args.display_value as name
+          from slice
+          join args using (arg_set_id)
+          where args.key='${this.argName}'
+        )
+        select
+          *,
+          (select count()
+          from ancestor_slice(s1.id) s2
+          join slice_with_arg s3 on s2.id=s3.id
+          ) as depth
+        from slice_with_arg s1
+        order by id;
+    `);
+
+    return new DisposableCallback(() => {
+      if (this.engine.isAlive) {
+        this.engine.query(`drop view ${this.viewName}`);
+      }
+    });
+  }
+
+  getTrackShellButtons(): m.Children {
+    return m(Button, {
+      onclick: () => {
+        // This behavior differs to the original behavior a little.
+        // Originally, hitting the close button on a single track removed ALL
+        // tracks with this argName, whereas this one only closes the single
+        // track.
+        // This will be easily fixable once we transition to using dynamic
+        // tracks instead of this "initial state" approach to add these tracks.
+        globals.dispatch(Actions.removeTracks({trackKeys: [this.trackKey]}));
+      },
+      icon: Icons.Close,
+      title: 'Close',
+      compact: true,
+    });
+  }
+}