Merge "Add some aggregate values about screen to the battery metrics."
diff --git a/Android.bp b/Android.bp
index 12c79b2..c9661b9 100644
--- a/Android.bp
+++ b/Android.bp
@@ -6566,7 +6566,6 @@
     "src/trace_processor/args_tracker.cc",
     "src/trace_processor/clock_tracker.cc",
     "src/trace_processor/default_modules.cc",
-    "src/trace_processor/destructible.cc",
     "src/trace_processor/event_tracker.cc",
     "src/trace_processor/forwarding_trace_parser.cc",
     "src/trace_processor/global_args_tracker.cc",
@@ -6632,6 +6631,7 @@
 filegroup {
   name: "perfetto_src_trace_processor_types_types",
   srcs: [
+    "src/trace_processor/types/destructible.cc",
     "src/trace_processor/types/gfp_flags.cc",
     "src/trace_processor/types/task_state.cc",
     "src/trace_processor/types/variadic.cc",
diff --git a/BUILD b/BUILD
index 23c24a6..ba890a3 100644
--- a/BUILD
+++ b/BUILD
@@ -837,10 +837,13 @@
 filegroup(
     name = "src_trace_processor_types_types",
     srcs = [
+        "src/trace_processor/types/destructible.cc",
+        "src/trace_processor/types/destructible.h",
         "src/trace_processor/types/gfp_flags.cc",
         "src/trace_processor/types/gfp_flags.h",
         "src/trace_processor/types/task_state.cc",
         "src/trace_processor/types/task_state.h",
+        "src/trace_processor/types/trace_processor_context.h",
         "src/trace_processor/types/variadic.cc",
         "src/trace_processor/types/variadic.h",
     ],
@@ -971,8 +974,6 @@
         "src/trace_processor/clock_tracker.h",
         "src/trace_processor/default_modules.cc",
         "src/trace_processor/default_modules.h",
-        "src/trace_processor/destructible.cc",
-        "src/trace_processor/destructible.h",
         "src/trace_processor/event_tracker.cc",
         "src/trace_processor/event_tracker.h",
         "src/trace_processor/forwarding_trace_parser.cc",
@@ -1029,7 +1030,6 @@
         "src/trace_processor/trace_blob_view.h",
         "src/trace_processor/trace_parser.h",
         "src/trace_processor/trace_processor_context.cc",
-        "src/trace_processor/trace_processor_context.h",
         "src/trace_processor/trace_processor_storage.cc",
         "src/trace_processor/trace_processor_storage_impl.cc",
         "src/trace_processor/trace_processor_storage_impl.h",
diff --git a/include/perfetto/ext/base/subprocess.h b/include/perfetto/ext/base/subprocess.h
index be813b3..3210fbe 100644
--- a/include/perfetto/ext/base/subprocess.h
+++ b/include/perfetto/ext/base/subprocess.h
@@ -17,6 +17,19 @@
 #ifndef INCLUDE_PERFETTO_EXT_BASE_SUBPROCESS_H_
 #define INCLUDE_PERFETTO_EXT_BASE_SUBPROCESS_H_
 
+#include "perfetto/base/build_config.h"
+
+// This is a #if as opposite to a GN condition, because GN conditions aren't propagated when
+// translating to Bazel or other build systems, as they get resolved at translation time. Without
+// this, the Bazel build breaks on Windows.
+#if PERFETTO_BUILDFLAG(PERFETTO_OS_LINUX) ||   \
+    PERFETTO_BUILDFLAG(PERFETTO_OS_ANDROID) || \
+    PERFETTO_BUILDFLAG(PERFETTO_OS_MACOSX)
+#define PERFETTO_HAS_SUBPROCESS() 1
+#else
+#define PERFETTO_HAS_SUBPROCESS() 0
+#endif
+
 #include <functional>
 #include <initializer_list>
 #include <string>
diff --git a/protos/perfetto/ipc/consumer_port.proto b/protos/perfetto/ipc/consumer_port.proto
index 545f802..b8b929c 100644
--- a/protos/perfetto/ipc/consumer_port.proto
+++ b/protos/perfetto/ipc/consumer_port.proto
@@ -230,6 +230,10 @@
 message QueryServiceStateRequest {}
 
 message QueryServiceStateResponse {
+  // In order to avoid hitting IPC message size limitations, the service will
+  // return >1 replies for each query, chunking the TracingServiceState. The
+  // receiver is expected to merge replies together and parse that when the
+  // last reply is received (i.e. when IPC's |has_more| == false).
   optional TracingServiceState service_state = 1;
 }
 
diff --git a/protos/perfetto/trace/perfetto_trace.proto b/protos/perfetto/trace/perfetto_trace.proto
index 2629c27..ae9cfbb 100644
--- a/protos/perfetto/trace/perfetto_trace.proto
+++ b/protos/perfetto/trace/perfetto_trace.proto
@@ -5854,8 +5854,8 @@
   // Legacy way to encode types. Names here are emitted by old perfetto_hprof,
   // which do not include the class_name. Handle like a HeapGraphType with
   // no location_id.
-  // TODO(fmayer): Remove this. This was was not used in any publicly available
-  // release.
+  // TODO(b/153552977): Remove this. This was was not used in any publicly
+  // available release.
   repeated InternedString type_names = 3;
 
   // Field names for references in managed heap graph.
diff --git a/protos/perfetto/trace/profiling/heap_graph.proto b/protos/perfetto/trace/profiling/heap_graph.proto
index 947033a..eeb5bec 100644
--- a/protos/perfetto/trace/profiling/heap_graph.proto
+++ b/protos/perfetto/trace/profiling/heap_graph.proto
@@ -117,8 +117,8 @@
   // Legacy way to encode types. Names here are emitted by old perfetto_hprof,
   // which do not include the class_name. Handle like a HeapGraphType with
   // no location_id.
-  // TODO(fmayer): Remove this. This was was not used in any publicly available
-  // release.
+  // TODO(b/153552977): Remove this. This was was not used in any publicly
+  // available release.
   repeated InternedString type_names = 3;
 
   // Field names for references in managed heap graph.
diff --git a/src/base/BUILD.gn b/src/base/BUILD.gn
index 8e51a11..1749b7b 100644
--- a/src/base/BUILD.gn
+++ b/src/base/BUILD.gn
@@ -33,6 +33,7 @@
     "string_splitter.cc",
     "string_utils.cc",
     "string_view.cc",
+    "subprocess.cc",
     "thread_checker.cc",
     "time.cc",
     "uuid.cc",
@@ -52,10 +53,6 @@
     ]
   }
 
-  if (!is_win && !is_nacl && !is_fuchsia) {
-    sources += [ "subprocess.cc" ]
-  }
-
   if (enable_stack_trace) {
     deps += [ ":debug_crash_stack_trace" ]
   }
@@ -130,6 +127,7 @@
     "string_utils_unittest.cc",
     "string_view_unittest.cc",
     "string_writer_unittest.cc",
+    "subprocess_unittest.cc",
     "time_unittest.cc",
     "uuid_unittest.cc",
     "weak_ptr_unittest.cc",
@@ -146,9 +144,6 @@
       "utils_unittest.cc",
     ]
   }
-  if (!is_win && !is_nacl && !is_fuchsia) {
-    sources += [ "subprocess_unittest.cc" ]
-  }
   if (perfetto_build_standalone || perfetto_build_with_android) {
     # This causes some problems on the chromium waterfall.
     sources += [ "unix_socket_unittest.cc" ]
diff --git a/src/base/subprocess.cc b/src/base/subprocess.cc
index 9f1bc0e..2e6e74a 100644
--- a/src/base/subprocess.cc
+++ b/src/base/subprocess.cc
@@ -16,6 +16,8 @@
 
 #include "perfetto/ext/base/subprocess.h"
 
+#if PERFETTO_HAS_SUBPROCESS()
+
 #include <poll.h>
 #include <signal.h>
 #include <stdio.h>
@@ -427,3 +429,5 @@
 
 }  // namespace base
 }  // namespace perfetto
+
+#endif  // PERFETTO_HAS_SUBPROCESS()
diff --git a/src/base/subprocess_unittest.cc b/src/base/subprocess_unittest.cc
index f430f45..f31b360 100644
--- a/src/base/subprocess_unittest.cc
+++ b/src/base/subprocess_unittest.cc
@@ -16,6 +16,7 @@
 
 #include "perfetto/ext/base/subprocess.h"
 
+#if PERFETTO_HAS_SUBPROCESS()
 #include <thread>
 
 #include <signal.h>
@@ -264,3 +265,5 @@
 }  // namespace
 }  // namespace base
 }  // namespace perfetto
+
+#endif  // PERFETTO_HAS_SUBPROCESS()
diff --git a/src/trace_processor/BUILD.gn b/src/trace_processor/BUILD.gn
index b69977a..dd77e13 100644
--- a/src/trace_processor/BUILD.gn
+++ b/src/trace_processor/BUILD.gn
@@ -85,8 +85,6 @@
     "clock_tracker.h",
     "default_modules.cc",
     "default_modules.h",
-    "destructible.cc",
-    "destructible.h",
     "event_tracker.cc",
     "event_tracker.h",
     "forwarding_trace_parser.cc",
@@ -143,7 +141,6 @@
     "trace_blob_view.h",
     "trace_parser.h",
     "trace_processor_context.cc",
-    "trace_processor_context.h",
     "trace_processor_storage.cc",
     "trace_processor_storage_impl.cc",
     "trace_processor_storage_impl.h",
@@ -285,6 +282,7 @@
     "../../gn:default_deps",
     "../base",
     "storage",
+    "types",
   ]
   public_deps = [ "../../include/perfetto/ext/trace_processor:export_json" ]
   if (enable_perfetto_trace_processor_json) {
@@ -407,6 +405,7 @@
     "db:unittests",
     "storage",
     "tables:unittests",
+    "types",
     "types:unittests",
   ]
 
diff --git a/src/trace_processor/additional_modules.h b/src/trace_processor/additional_modules.h
index d68720d..3228494 100644
--- a/src/trace_processor/additional_modules.h
+++ b/src/trace_processor/additional_modules.h
@@ -17,7 +17,7 @@
 #ifndef SRC_TRACE_PROCESSOR_ADDITIONAL_MODULES_H_
 #define SRC_TRACE_PROCESSOR_ADDITIONAL_MODULES_H_
 
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/args_tracker.h b/src/trace_processor/args_tracker.h
index 56db297..6bb6674 100644
--- a/src/trace_processor/args_tracker.h
+++ b/src/trace_processor/args_tracker.h
@@ -19,7 +19,7 @@
 
 #include "src/trace_processor/global_args_tracker.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 #include "src/trace_processor/types/variadic.h"
 
 namespace perfetto {
diff --git a/src/trace_processor/clock_tracker.cc b/src/trace_processor/clock_tracker.cc
index 483c707..934720a 100644
--- a/src/trace_processor/clock_tracker.cc
+++ b/src/trace_processor/clock_tracker.cc
@@ -24,7 +24,7 @@
 #include "perfetto/base/logging.h"
 #include "perfetto/ext/base/hash.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/trace/clock_snapshot.pbzero.h"
 
diff --git a/src/trace_processor/clock_tracker_unittest.cc b/src/trace_processor/clock_tracker_unittest.cc
index 4e1e1f6..f1f3338 100644
--- a/src/trace_processor/clock_tracker_unittest.cc
+++ b/src/trace_processor/clock_tracker_unittest.cc
@@ -20,7 +20,7 @@
 
 #include "perfetto/ext/base/optional.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 #include "test/gtest_and_gmock.h"
 
 #include "protos/perfetto/trace/clock_snapshot.pbzero.h"
diff --git a/src/trace_processor/default_modules.h b/src/trace_processor/default_modules.h
index 97da4e8..a524c43 100644
--- a/src/trace_processor/default_modules.h
+++ b/src/trace_processor/default_modules.h
@@ -17,7 +17,7 @@
 #ifndef SRC_TRACE_PROCESSOR_DEFAULT_MODULES_H_
 #define SRC_TRACE_PROCESSOR_DEFAULT_MODULES_H_
 
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/describe_slice_generator.cc b/src/trace_processor/describe_slice_generator.cc
index c508f9b..60882b5 100644
--- a/src/trace_processor/describe_slice_generator.cc
+++ b/src/trace_processor/describe_slice_generator.cc
@@ -17,7 +17,7 @@
 #include "src/trace_processor/describe_slice_generator.h"
 
 #include "src/trace_processor/analysis/describe_slice.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/event_tracker.cc b/src/trace_processor/event_tracker.cc
index 3d82682..15343dc 100644
--- a/src/trace_processor/event_tracker.cc
+++ b/src/trace_processor/event_tracker.cc
@@ -23,8 +23,8 @@
 #include "src/trace_processor/args_tracker.h"
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/storage/stats.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 #include "src/trace_processor/types/variadic.h"
 
 namespace perfetto {
diff --git a/src/trace_processor/experimental_flamegraph_generator.cc b/src/trace_processor/experimental_flamegraph_generator.cc
index 17312f0..3a13375 100644
--- a/src/trace_processor/experimental_flamegraph_generator.cc
+++ b/src/trace_processor/experimental_flamegraph_generator.cc
@@ -20,7 +20,7 @@
 
 #include "src/trace_processor/heap_profile_tracker.h"
 #include "src/trace_processor/importers/proto/heap_graph_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/export_json.cc b/src/trace_processor/export_json.cc
index 9161996..2e5516d 100644
--- a/src/trace_processor/export_json.cc
+++ b/src/trace_processor/export_json.cc
@@ -32,8 +32,8 @@
 #include "src/trace_processor/importers/json/json_utils.h"
 #include "src/trace_processor/storage/metadata.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/trace_processor_storage_impl.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #if PERFETTO_BUILDFLAG(PERFETTO_TP_JSON)
 #include <json/reader.h>
diff --git a/src/trace_processor/export_json_unittest.cc b/src/trace_processor/export_json_unittest.cc
index be4cc90..4be5687 100644
--- a/src/trace_processor/export_json_unittest.cc
+++ b/src/trace_processor/export_json_unittest.cc
@@ -28,8 +28,8 @@
 #include "src/trace_processor/args_tracker.h"
 #include "src/trace_processor/metadata_tracker.h"
 #include "src/trace_processor/process_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "test/gtest_and_gmock.h"
 
diff --git a/src/trace_processor/forwarding_trace_parser.h b/src/trace_processor/forwarding_trace_parser.h
index 14e9712..8b8a8bb 100644
--- a/src/trace_processor/forwarding_trace_parser.h
+++ b/src/trace_processor/forwarding_trace_parser.h
@@ -19,7 +19,7 @@
 
 #include "src/trace_processor/chunked_trace_reader.h"
 
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/global_args_tracker.h b/src/trace_processor/global_args_tracker.h
index 8bdb691..224a3a6 100644
--- a/src/trace_processor/global_args_tracker.h
+++ b/src/trace_processor/global_args_tracker.h
@@ -18,7 +18,7 @@
 #define SRC_TRACE_PROCESSOR_GLOBAL_ARGS_TRACKER_H_
 
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 #include "src/trace_processor/types/variadic.h"
 
 namespace perfetto {
diff --git a/src/trace_processor/heap_profile_tracker.cc b/src/trace_processor/heap_profile_tracker.cc
index 157d203..2782516 100644
--- a/src/trace_processor/heap_profile_tracker.cc
+++ b/src/trace_processor/heap_profile_tracker.cc
@@ -16,9 +16,9 @@
 
 #include "src/trace_processor/heap_profile_tracker.h"
 
-#include "src/trace_processor/process_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "perfetto/base/logging.h"
+#include "src/trace_processor/process_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/trace/profiling/profile_common.pbzero.h"
 #include "protos/perfetto/trace/profiling/profile_packet.pbzero.h"
diff --git a/src/trace_processor/heap_profile_tracker_unittest.cc b/src/trace_processor/heap_profile_tracker_unittest.cc
index 35a9c2d..55ce380 100644
--- a/src/trace_processor/heap_profile_tracker_unittest.cc
+++ b/src/trace_processor/heap_profile_tracker_unittest.cc
@@ -17,7 +17,7 @@
 #include "src/trace_processor/heap_profile_tracker.h"
 
 #include "src/trace_processor/stack_profile_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 #include "test/gtest_and_gmock.h"
 
 namespace perfetto {
diff --git a/src/trace_processor/importers/ftrace/ftrace_parser.h b/src/trace_processor/importers/ftrace/ftrace_parser.h
index 10eacfe..d3bbce2 100644
--- a/src/trace_processor/importers/ftrace/ftrace_parser.h
+++ b/src/trace_processor/importers/ftrace/ftrace_parser.h
@@ -25,7 +25,7 @@
 #include "src/trace_processor/importers/ftrace/sched_event_tracker.h"
 #include "src/trace_processor/timestamped_trace_piece.h"
 #include "src/trace_processor/trace_blob_view.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/importers/ftrace/ftrace_tokenizer.h b/src/trace_processor/importers/ftrace/ftrace_tokenizer.h
index 7febc3b..e48cb74 100644
--- a/src/trace_processor/importers/ftrace/ftrace_tokenizer.h
+++ b/src/trace_processor/importers/ftrace/ftrace_tokenizer.h
@@ -20,7 +20,7 @@
 #include "protos/perfetto/trace/ftrace/ftrace_event_bundle.pbzero.h"
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/trace_blob_view.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/importers/ftrace/rss_stat_tracker.cc b/src/trace_processor/importers/ftrace/rss_stat_tracker.cc
index e9d48de..043baee 100644
--- a/src/trace_processor/importers/ftrace/rss_stat_tracker.cc
+++ b/src/trace_processor/importers/ftrace/rss_stat_tracker.cc
@@ -18,7 +18,7 @@
 
 #include "src/trace_processor/event_tracker.h"
 #include "src/trace_processor/process_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/trace/ftrace/kmem.pbzero.h"
 
diff --git a/src/trace_processor/importers/ftrace/sched_event_tracker.cc b/src/trace_processor/importers/ftrace/sched_event_tracker.cc
index 30bf095..1bd3ba4 100644
--- a/src/trace_processor/importers/ftrace/sched_event_tracker.cc
+++ b/src/trace_processor/importers/ftrace/sched_event_tracker.cc
@@ -24,8 +24,8 @@
 #include "src/trace_processor/importers/ftrace/ftrace_descriptors.h"
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/storage/stats.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/types/task_state.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 #include "src/trace_processor/types/variadic.h"
 
 #include "protos/perfetto/trace/ftrace/ftrace_event.pbzero.h"
diff --git a/src/trace_processor/importers/ftrace/sched_event_tracker.h b/src/trace_processor/importers/ftrace/sched_event_tracker.h
index df1dee8..a4e8602 100644
--- a/src/trace_processor/importers/ftrace/sched_event_tracker.h
+++ b/src/trace_processor/importers/ftrace/sched_event_tracker.h
@@ -22,9 +22,9 @@
 
 #include "perfetto/ext/base/string_view.h"
 #include "perfetto/ext/base/utils.h"
-#include "src/trace_processor/destructible.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/destructible.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/importers/fuchsia/fuchsia_trace_tokenizer.cc b/src/trace_processor/importers/fuchsia/fuchsia_trace_tokenizer.cc
index 4a37fff..b48ac69 100644
--- a/src/trace_processor/importers/fuchsia/fuchsia_trace_tokenizer.cc
+++ b/src/trace_processor/importers/fuchsia/fuchsia_trace_tokenizer.cc
@@ -24,9 +24,9 @@
 #include "src/trace_processor/importers/fuchsia/fuchsia_record.h"
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/slice_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/trace_sorter.h"
 #include "src/trace_processor/types/task_state.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/importers/json/json_trace_parser.cc b/src/trace_processor/importers/json/json_trace_parser.cc
index 4165518..1502e9e 100644
--- a/src/trace_processor/importers/json/json_trace_parser.cc
+++ b/src/trace_processor/importers/json/json_trace_parser.cc
@@ -28,8 +28,8 @@
 #include "src/trace_processor/importers/json/json_utils.h"
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/slice_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/importers/json/json_tracker.h b/src/trace_processor/importers/json/json_tracker.h
index fda3f06..ac24083 100644
--- a/src/trace_processor/importers/json/json_tracker.h
+++ b/src/trace_processor/importers/json/json_tracker.h
@@ -17,9 +17,9 @@
 #ifndef SRC_TRACE_PROCESSOR_IMPORTERS_JSON_JSON_TRACKER_H_
 #define SRC_TRACE_PROCESSOR_IMPORTERS_JSON_JSON_TRACKER_H_
 
-#include "src/trace_processor/destructible.h"
 #include "src/trace_processor/importers/json/json_utils.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/destructible.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace Json {
 class Value;
diff --git a/src/trace_processor/importers/proto/android_probes_parser.cc b/src/trace_processor/importers/proto/android_probes_parser.cc
index 1f19644..f8eb054 100644
--- a/src/trace_processor/importers/proto/android_probes_parser.cc
+++ b/src/trace_processor/importers/proto/android_probes_parser.cc
@@ -24,7 +24,7 @@
 #include "src/trace_processor/metadata_tracker.h"
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/syscall_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/common/android_log_constants.pbzero.h"
 #include "protos/perfetto/config/trace_config.pbzero.h"
diff --git a/src/trace_processor/importers/proto/android_probes_tracker.h b/src/trace_processor/importers/proto/android_probes_tracker.h
index f50fd81..6e2e193 100644
--- a/src/trace_processor/importers/proto/android_probes_tracker.h
+++ b/src/trace_processor/importers/proto/android_probes_tracker.h
@@ -22,7 +22,7 @@
 #include "perfetto/ext/base/optional.h"
 
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/importers/proto/graphics_event_parser.cc b/src/trace_processor/importers/proto/graphics_event_parser.cc
index 87c9503..91a2ac7 100644
--- a/src/trace_processor/importers/proto/graphics_event_parser.cc
+++ b/src/trace_processor/importers/proto/graphics_event_parser.cc
@@ -25,8 +25,8 @@
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/slice_tracker.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/common/gpu_counter_descriptor.pbzero.h"
 #include "protos/perfetto/trace/android/graphics_frame_event.pbzero.h"
diff --git a/src/trace_processor/importers/proto/heap_graph_module.cc b/src/trace_processor/importers/proto/heap_graph_module.cc
index 3ca232a..572b45a 100644
--- a/src/trace_processor/importers/proto/heap_graph_module.cc
+++ b/src/trace_processor/importers/proto/heap_graph_module.cc
@@ -19,7 +19,7 @@
 #include "src/trace_processor/importers/proto/heap_graph_tracker.h"
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/trace/profiling/heap_graph.pbzero.h"
 
@@ -215,11 +215,50 @@
   }
 }
 
+void HeapGraphModule::DeobfuscateClass(
+    base::Optional<StringPool::Id> package_name_id,
+    StringPool::Id obfuscated_class_name_id,
+    const protos::pbzero::ObfuscatedClass::Decoder& cls) {
+  auto* heap_graph_tracker = HeapGraphTracker::GetOrCreate(context_);
+  const std::vector<tables::HeapGraphClassTable::Id>* cls_objects =
+      heap_graph_tracker->RowsForType(package_name_id,
+                                      obfuscated_class_name_id);
+
+  if (cls_objects) {
+    heap_graph_tracker->AddDeobfuscationMapping(
+        package_name_id, obfuscated_class_name_id,
+        context_->storage->InternString(
+            base::StringView(cls.deobfuscated_name())));
+
+    for (tables::HeapGraphClassTable::Id id : *cls_objects) {
+      uint32_t row =
+          *context_->storage->heap_graph_class_table().id().IndexOf(id);
+      const StringPool::Id obfuscated_type_name =
+          context_->storage->heap_graph_class_table().name()[row];
+      StringPool::Id deobfuscated_type_name =
+          heap_graph_tracker->MaybeDeobfuscate(package_name_id,
+                                               obfuscated_type_name);
+      PERFETTO_CHECK(!deobfuscated_type_name.is_null());
+      context_->storage->mutable_heap_graph_class_table()
+          ->mutable_deobfuscated_name()
+          ->Set(row, deobfuscated_type_name);
+    }
+  } else {
+    PERFETTO_DLOG("Class %s not found",
+                  cls.obfuscated_name().ToStdString().c_str());
+  }
+}
+
 void HeapGraphModule::ParseDeobfuscationMapping(protozero::ConstBytes blob) {
-  // TODO(fmayer): Support multiple profiles in the same trace.
   auto* heap_graph_tracker = HeapGraphTracker::GetOrCreate(context_);
   protos::pbzero::DeobfuscationMapping::Decoder deobfuscation_mapping(
       blob.data, blob.size);
+  base::Optional<StringPool::Id> package_name_id;
+  if (deobfuscation_mapping.package_name().size > 0) {
+    package_name_id = context_->storage->string_pool().GetId(
+        deobfuscation_mapping.package_name());
+  }
+
   for (auto class_it = deobfuscation_mapping.obfuscated_classes(); class_it;
        ++class_it) {
     protos::pbzero::ObfuscatedClass::Decoder cls(*class_it);
@@ -229,30 +268,12 @@
       PERFETTO_DLOG("Class string %s not found",
                     cls.obfuscated_name().ToStdString().c_str());
     } else {
-      const std::vector<tables::HeapGraphClassTable::Id>* cls_objects =
-          heap_graph_tracker->RowsForType(*obfuscated_class_name_id);
-
-      if (cls_objects) {
-        heap_graph_tracker->AddDeobfuscationMapping(
-            *obfuscated_class_name_id,
-            context_->storage->InternString(
-                base::StringView(cls.deobfuscated_name())));
-
-        for (tables::HeapGraphClassTable::Id id : *cls_objects) {
-          uint32_t row =
-              *context_->storage->heap_graph_class_table().id().IndexOf(id);
-          const StringPool::Id obfuscated_type_name =
-              context_->storage->heap_graph_class_table().name()[row];
-          StringPool::Id deobfuscated_type_name =
-              heap_graph_tracker->MaybeDeobfuscate(obfuscated_type_name);
-          PERFETTO_CHECK(!deobfuscated_type_name.is_null());
-          context_->storage->mutable_heap_graph_class_table()
-              ->mutable_deobfuscated_name()
-              ->Set(row, deobfuscated_type_name);
-        }
-      } else {
-        PERFETTO_DLOG("Class %s not found",
-                      cls.obfuscated_name().ToStdString().c_str());
+      // TODO(b/153552977): Remove this work-around for legacy traces.
+      // For traces without location information, deobfuscate all matching
+      // classes.
+      DeobfuscateClass(base::nullopt, *obfuscated_class_name_id, cls);
+      if (package_name_id) {
+        DeobfuscateClass(package_name_id, *obfuscated_class_name_id, cls);
       }
     }
     for (auto member_it = cls.obfuscated_members(); member_it; ++member_it) {
diff --git a/src/trace_processor/importers/proto/heap_graph_module.h b/src/trace_processor/importers/proto/heap_graph_module.h
index 2fdbebc..b1816b0 100644
--- a/src/trace_processor/importers/proto/heap_graph_module.h
+++ b/src/trace_processor/importers/proto/heap_graph_module.h
@@ -40,6 +40,9 @@
  private:
   void ParseHeapGraph(uint32_t seq_id, int64_t ts, protozero::ConstBytes);
   void ParseDeobfuscationMapping(protozero::ConstBytes);
+  void DeobfuscateClass(base::Optional<StringPool::Id> package_name_id,
+                        StringPool::Id obfuscated_class_id,
+                        const protos::pbzero::ObfuscatedClass::Decoder& cls);
 
   TraceProcessorContext* context_;
 };
diff --git a/src/trace_processor/importers/proto/heap_graph_tracker.cc b/src/trace_processor/importers/proto/heap_graph_tracker.cc
index 9b370e2..2e43ae8 100644
--- a/src/trace_processor/importers/proto/heap_graph_tracker.cc
+++ b/src/trace_processor/importers/proto/heap_graph_tracker.cc
@@ -16,9 +16,32 @@
 
 #include "src/trace_processor/importers/proto/heap_graph_tracker.h"
 
+#include "perfetto/ext/base/string_splitter.h"
+#include "perfetto/ext/base/string_utils.h"
+
 namespace perfetto {
 namespace trace_processor {
 
+namespace {
+base::Optional<base::StringView> PackageFromApp(base::StringView location) {
+  location = location.substr(base::StringView("/data/app/").size());
+  size_t slash = location.find('/');
+  if (slash == std::string::npos) {
+    return base::nullopt;
+  }
+  size_t second_slash = location.find('/', slash + 1);
+  if (second_slash == std::string::npos) {
+    return base::nullopt;
+  }
+  location = location.substr(slash + 1, second_slash - slash);
+  size_t minus = location.find('-');
+  if (minus == std::string::npos) {
+    return base::nullopt;
+  }
+  return location.substr(0, minus);
+}
+}  // namespace
+
 base::Optional<base::StringView> GetStaticClassTypeName(base::StringView type) {
   static const base::StringView kJavaClassTemplate("java.lang.Class<");
   if (!type.empty() && type.at(type.size() - 1) == '>' &&
@@ -71,6 +94,93 @@
 HeapGraphTracker::HeapGraphTracker(TraceProcessorContext* context)
     : context_(context) {}
 
+base::Optional<std::string> HeapGraphTracker::PackageFromLocation(
+    base::StringView location) {
+  // List of some hardcoded apps that do not follow the scheme used in
+  // PackageFromApp. Ask for yours to be added.
+  //
+  // TODO(b/153632336): Get rid of the hardcoded list of system apps.
+  base::StringView sysui(
+      "/system_ext/priv-app/SystemUIGoogle/SystemUIGoogle.apk");
+  if (location.size() >= sysui.size() &&
+      location.substr(0, sysui.size()) == sysui) {
+    return "com.android.systemui";
+  }
+
+  base::StringView phonesky("/product/priv-app/Phonesky/Phonesky.apk");
+  if (location.size() >= phonesky.size() &&
+      location.substr(0, phonesky.size()) == phonesky) {
+    return "com.android.vending";
+  }
+
+  base::StringView maps("/product/app/Maps/Maps.apk");
+  if (location.size() >= maps.size() &&
+      location.substr(0, maps.size()) == maps) {
+    return "com.google.android.apps.maps";
+  }
+
+  base::StringView launcher(
+      "/system_ext/priv-app/NexusLauncherRelease/NexusLauncherRelease.apk");
+  if (location.size() >= launcher.size() &&
+      location.substr(0, launcher.size()) == launcher) {
+    return "com.google.android.apps.nexuslauncher";
+  }
+
+  base::StringView photos("/product/app/Photos/Photos.apk");
+  if (location.size() >= photos.size() &&
+      location.substr(0, photos.size()) == photos) {
+    return "com.google.android.apps.photos";
+  }
+
+  base::StringView wellbeing(
+      "/product/priv-app/WellbeingPrebuilt/WellbeingPrebuilt.apk");
+  if (location.size() >= wellbeing.size() &&
+      location.substr(0, wellbeing.size()) == wellbeing) {
+    return "com.google.android.apps.wellbeing";
+  }
+
+  base::StringView matchmaker("MatchMaker");
+  if (location.size() >= matchmaker.size() &&
+      location.find(matchmaker) != base::StringView::npos) {
+    return "com.google.android.as";
+  }
+
+  base::StringView gm("/product/app/PrebuiltGmail/PrebuiltGmail.apk");
+  if (location.size() >= gm.size() && location.substr(0, gm.size()) == gm) {
+    return "com.google.android.gm";
+  }
+
+  base::StringView gmscore("/product/priv-app/PrebuiltGmsCore/PrebuiltGmsCore");
+  if (location.size() >= gmscore.size() &&
+      location.substr(0, gmscore.size()) == gmscore) {
+    return "com.google.android.gms";
+  }
+
+  base::StringView velvet("/product/priv-app/Velvet/Velvet.apk");
+  if (location.size() >= velvet.size() &&
+      location.substr(0, velvet.size()) == velvet) {
+    return "com.google.android.googlequicksearchbox";
+  }
+
+  base::StringView inputmethod(
+      "/product/app/LatinIMEGooglePrebuilt/LatinIMEGooglePrebuilt.apk");
+  if (location.size() >= inputmethod.size() &&
+      location.substr(0, inputmethod.size()) == inputmethod) {
+    return "com.google.android.inputmethod.latin";
+  }
+
+  base::StringView data_app("/data/app/");
+  if (location.substr(0, data_app.size()) == data_app) {
+    auto package = PackageFromApp(location);
+    if (!package) {
+      context_->storage->IncrementStats(stats::heap_graph_location_parse_error);
+      return base::nullopt;
+    }
+    return package->ToStdString();
+  }
+  return base::nullopt;
+}
+
 HeapGraphTracker::SequenceState& HeapGraphTracker::GetOrCreateSequence(
     uint32_t seq_id) {
   auto seq_it = sequence_state_.find(seq_id);
@@ -212,8 +322,39 @@
     type_id_to_db[id] = id_and_row.id;
     base::StringView normalized_type =
         NormalizeTypeName(context_->storage->GetString(interned_type.name));
-    class_to_rows_[context_->storage->InternString(normalized_type)]
-        .emplace_back(id_and_row.id);
+
+    // Annoyingly, some apps have a relative path to base.apk. We take this to
+    // mean the main package, so we treat it as if the location was unknown.
+    bool is_base_apk = false;
+    if (location_name) {
+      base::StringView base_apk("base.apk");
+      is_base_apk = context_->storage->GetString(*location_name)
+                        .substr(0, base_apk.size()) == base_apk;
+    }
+
+    if (location_name && !is_base_apk) {
+      base::Optional<std::string> package_name =
+          PackageFromLocation(context_->storage->GetString(*location_name));
+      if (package_name) {
+        class_to_rows_[std::make_pair(
+                           context_->storage->InternString(
+                               base::StringView(*package_name)),
+                           context_->storage->InternString(normalized_type))]
+            .emplace_back(id_and_row.id);
+      }
+    } else {
+      // TODO(b/153552977): Remove this workaround.
+      // For profiles collected for old versions of perfetto_hprof, we do not
+      // have any location information. We store them using the nullopt
+      // location, and assume they are all part of the main APK.
+      //
+      // This is to keep ingestion of old profiles working (especially
+      // important for the UI).
+      class_to_rows_[std::make_pair(
+                         base::nullopt,
+                         context_->storage->InternString(normalized_type))]
+          .emplace_back(id_and_row.id);
+    }
   }
 
   for (const SourceObject& obj : sequence_state.current_objects) {
@@ -403,11 +544,13 @@
   }
 }
 
-StringPool::Id HeapGraphTracker::MaybeDeobfuscate(StringPool::Id id) {
+StringPool::Id HeapGraphTracker::MaybeDeobfuscate(
+    base::Optional<StringPool::Id> package_name,
+    StringPool::Id id) {
   base::StringView type_name = context_->storage->GetString(id);
   auto normalized_type = GetNormalizedType(type_name);
-  auto it = deobfuscation_mapping_.find(
-      context_->storage->InternString(normalized_type.name));
+  auto it = deobfuscation_mapping_.find(std::make_pair(
+      package_name, context_->storage->InternString(normalized_type.name)));
   if (it == deobfuscation_mapping_.end())
     return id;
 
@@ -419,9 +562,11 @@
 }
 
 void HeapGraphTracker::AddDeobfuscationMapping(
+    base::Optional<StringPool::Id> package_name,
     StringPool::Id obfuscated_name,
     StringPool::Id deobfuscated_name) {
-  deobfuscation_mapping_.emplace(obfuscated_name, deobfuscated_name);
+  deobfuscation_mapping_.emplace(std::make_pair(package_name, obfuscated_name),
+                                 deobfuscated_name);
 }
 
 }  // namespace trace_processor
diff --git a/src/trace_processor/importers/proto/heap_graph_tracker.h b/src/trace_processor/importers/proto/heap_graph_tracker.h
index 6623235..eb25b21 100644
--- a/src/trace_processor/importers/proto/heap_graph_tracker.h
+++ b/src/trace_processor/importers/proto/heap_graph_tracker.h
@@ -26,7 +26,7 @@
 #include "protos/perfetto/trace/profiling/heap_graph.pbzero.h"
 #include "src/trace_processor/importers/proto/heap_graph_walker.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
@@ -38,6 +38,8 @@
   bool is_static_class;
   size_t number_of_arrays;
 };
+
+base::Optional<std::string> PackageFromLocation(base::StringView location);
 base::Optional<base::StringView> GetStaticClassTypeName(base::StringView type);
 size_t NumberOfArrays(base::StringView type);
 NormalizedType GetNormalizedType(base::StringView type);
@@ -100,13 +102,16 @@
                    int64_t unique_retained) override;
   void NotifyEndOfFile();
 
-  void AddDeobfuscationMapping(StringPool::Id obfuscated_name,
+  void AddDeobfuscationMapping(base::Optional<StringPool::Id> package_name,
+                               StringPool::Id obfuscated_name,
                                StringPool::Id deobfuscated_name);
-  StringPool::Id MaybeDeobfuscate(StringPool::Id);
+  StringPool::Id MaybeDeobfuscate(base::Optional<StringPool::Id> package_name,
+                                  StringPool::Id);
 
   const std::vector<tables::HeapGraphClassTable::Id>* RowsForType(
+      base::Optional<StringPool::Id> package_name,
       StringPool::Id type_name) const {
-    auto it = class_to_rows_.find(type_name);
+    auto it = class_to_rows_.find(std::make_pair(package_name, type_name));
     if (it == class_to_rows_.end())
       return nullptr;
     return &it->second;
@@ -123,6 +128,9 @@
       const int64_t current_ts,
       const UniquePid current_upid);
 
+  // public for testing.
+  base::Optional<std::string> PackageFromLocation(base::StringView location);
+
  private:
   struct InternedField {
     StringPool::Id name;
@@ -154,11 +162,14 @@
   std::map<uint32_t, SequenceState> sequence_state_;
   std::map<std::pair<UniquePid, int64_t /* ts */>, HeapGraphWalker> walkers_;
 
-  std::map<StringPool::Id, std::vector<tables::HeapGraphClassTable::Id>>
+  std::map<std::pair<base::Optional<StringPool::Id>, StringPool::Id>,
+           std::vector<tables::HeapGraphClassTable::Id>>
       class_to_rows_;
   std::map<StringPool::Id, std::vector<int64_t>> field_to_rows_;
 
-  std::map<StringPool::Id, StringPool::Id> deobfuscation_mapping_;
+  std::map<std::pair<base::Optional<StringPool::Id>, StringPool::Id>,
+           StringPool::Id>
+      deobfuscation_mapping_;
 };
 
 }  // namespace trace_processor
diff --git a/src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc b/src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc
index 7472818..796956c 100644
--- a/src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc
+++ b/src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc
@@ -25,6 +25,15 @@
 
 using ::testing::UnorderedElementsAre;
 
+TEST(HeapGraphTrackerTest, PackageFromLocationApp) {
+  TraceProcessorContext context;
+  HeapGraphTracker tracker(&context);
+  EXPECT_EQ(tracker.PackageFromLocation(
+                "/data/app/~~ASDFGH1234QWerT==/"
+                "com.twitter.android-MNBVCX7890SDTst6==/test.apk"),
+            "com.twitter.android");
+}
+
 TEST(HeapGraphTrackerTest, BuildFlamegraph) {
   //           4@A 5@B
   //             \ /
diff --git a/src/trace_processor/importers/proto/packet_sequence_state.h b/src/trace_processor/importers/proto/packet_sequence_state.h
index fe34b5c..9aa833c 100644
--- a/src/trace_processor/importers/proto/packet_sequence_state.h
+++ b/src/trace_processor/importers/proto/packet_sequence_state.h
@@ -27,7 +27,7 @@
 #include "src/trace_processor/stack_profile_tracker.h"
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/trace_blob_view.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/trace/trace_packet_defaults.pbzero.h"
 #include "protos/perfetto/trace/track_event/track_event.pbzero.h"
diff --git a/src/trace_processor/importers/proto/profile_module.cc b/src/trace_processor/importers/proto/profile_module.cc
index a6c49cd..8249452 100644
--- a/src/trace_processor/importers/proto/profile_module.cc
+++ b/src/trace_processor/importers/proto/profile_module.cc
@@ -25,8 +25,8 @@
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/tables/profiler_tables.h"
 #include "src/trace_processor/timestamped_trace_piece.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/trace_sorter.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/trace/clock_snapshot.pbzero.h"
 #include "protos/perfetto/trace/profiling/profile_packet.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 4524384..4f23f21 100644
--- a/src/trace_processor/importers/proto/proto_importer_module.cc
+++ b/src/trace_processor/importers/proto/proto_importer_module.cc
@@ -15,7 +15,7 @@
  */
 
 #include "src/trace_processor/importers/proto/proto_importer_module.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/importers/proto/proto_trace_parser.cc b/src/trace_processor/importers/proto/proto_trace_parser.cc
index 69394d6..4eda8a9 100644
--- a/src/trace_processor/importers/proto/proto_trace_parser.cc
+++ b/src/trace_processor/importers/proto/proto_trace_parser.cc
@@ -43,8 +43,8 @@
 #include "src/trace_processor/stack_profile_tracker.h"
 #include "src/trace_processor/storage/metadata.h"
 #include "src/trace_processor/timestamped_trace_piece.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 #include "src/trace_processor/types/variadic.h"
 
 #include "protos/perfetto/common/trace_stats.pbzero.h"
diff --git a/src/trace_processor/importers/proto/system_probes_parser.cc b/src/trace_processor/importers/proto/system_probes_parser.cc
index 7276e93..40251f6 100644
--- a/src/trace_processor/importers/proto/system_probes_parser.cc
+++ b/src/trace_processor/importers/proto/system_probes_parser.cc
@@ -24,7 +24,7 @@
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/storage/metadata.h"
 #include "src/trace_processor/syscall_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/trace/ps/process_stats.pbzero.h"
 #include "protos/perfetto/trace/ps/process_tree.pbzero.h"
diff --git a/src/trace_processor/importers/proto/track_event_module.cc b/src/trace_processor/importers/proto/track_event_module.cc
index bb59cd7..102614e 100644
--- a/src/trace_processor/importers/proto/track_event_module.cc
+++ b/src/trace_processor/importers/proto/track_event_module.cc
@@ -18,8 +18,8 @@
 #include "perfetto/base/build_config.h"
 #include "perfetto/ext/base/string_utils.h"
 #include "src/trace_processor/timestamped_trace_piece.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "protos/perfetto/config/data_source_config.pbzero.h"
 #include "protos/perfetto/config/trace_config.pbzero.h"
diff --git a/src/trace_processor/importers/proto/vulkan_memory_tracker.cc b/src/trace_processor/importers/proto/vulkan_memory_tracker.cc
index 03d1776..71ea4ee 100644
--- a/src/trace_processor/importers/proto/vulkan_memory_tracker.cc
+++ b/src/trace_processor/importers/proto/vulkan_memory_tracker.cc
@@ -19,7 +19,7 @@
 #include <string>
 #include "protos/perfetto/trace/interned_data/interned_data.pbzero.h"
 #include "src/trace_processor/process_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "perfetto/base/logging.h"
 
diff --git a/src/trace_processor/importers/systrace/systrace_line_parser.h b/src/trace_processor/importers/systrace/systrace_line_parser.h
index ca1dafe..00346a3 100644
--- a/src/trace_processor/importers/systrace/systrace_line_parser.h
+++ b/src/trace_processor/importers/systrace/systrace_line_parser.h
@@ -21,7 +21,7 @@
 
 #include "src/trace_processor/importers/systrace/systrace_line.h"
 #include "src/trace_processor/trace_parser.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/importers/systrace/systrace_parser.h b/src/trace_processor/importers/systrace/systrace_parser.h
index d993760..eaa6f15 100644
--- a/src/trace_processor/importers/systrace/systrace_parser.h
+++ b/src/trace_processor/importers/systrace/systrace_parser.h
@@ -19,7 +19,7 @@
 
 #include <ostream>
 
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "perfetto/base/logging.h"
 #include "perfetto/ext/base/string_utils.h"
diff --git a/src/trace_processor/importers/systrace/systrace_trace_parser.h b/src/trace_processor/importers/systrace/systrace_trace_parser.h
index 7b3c834..9e7464e 100644
--- a/src/trace_processor/importers/systrace/systrace_trace_parser.h
+++ b/src/trace_processor/importers/systrace/systrace_trace_parser.h
@@ -24,7 +24,7 @@
 #include "src/trace_processor/importers/systrace/systrace_line_parser.h"
 #include "src/trace_processor/importers/systrace/systrace_line_tokenizer.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/metadata_tracker.cc b/src/trace_processor/metadata_tracker.cc
index c8ce9e7..981e233 100644
--- a/src/trace_processor/metadata_tracker.cc
+++ b/src/trace_processor/metadata_tracker.cc
@@ -17,7 +17,7 @@
 #include "src/trace_processor/metadata_tracker.h"
 
 #include "src/trace_processor/process_tracker.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/perf_sample_tracker.cc b/src/trace_processor/perf_sample_tracker.cc
index 3a078b3..3671c70 100644
--- a/src/trace_processor/perf_sample_tracker.cc
+++ b/src/trace_processor/perf_sample_tracker.cc
@@ -23,8 +23,8 @@
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/slice_tracker.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #if !PERFETTO_BUILDFLAG(PERFETTO_OS_WIN)
 #include <cxxabi.h>
diff --git a/src/trace_processor/process_tracker.h b/src/trace_processor/process_tracker.h
index 110fa4b..86ae13b 100644
--- a/src/trace_processor/process_tracker.h
+++ b/src/trace_processor/process_tracker.h
@@ -21,7 +21,7 @@
 
 #include "perfetto/ext/base/string_view.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/slice_tracker.cc b/src/trace_processor/slice_tracker.cc
index 5d8885a..01edcee 100644
--- a/src/trace_processor/slice_tracker.cc
+++ b/src/trace_processor/slice_tracker.cc
@@ -21,8 +21,8 @@
 #include "src/trace_processor/process_tracker.h"
 #include "src/trace_processor/slice_tracker.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/slice_tracker_unittest.cc b/src/trace_processor/slice_tracker_unittest.cc
index 5eef7bb..b9a886f 100644
--- a/src/trace_processor/slice_tracker_unittest.cc
+++ b/src/trace_processor/slice_tracker_unittest.cc
@@ -19,7 +19,7 @@
 #include "src/trace_processor/args_tracker.h"
 #include "src/trace_processor/slice_tracker.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 #include "test/gtest_and_gmock.h"
 
 namespace perfetto {
diff --git a/src/trace_processor/stack_profile_tracker.cc b/src/trace_processor/stack_profile_tracker.cc
index af41a6e..2edd915 100644
--- a/src/trace_processor/stack_profile_tracker.cc
+++ b/src/trace_processor/stack_profile_tracker.cc
@@ -16,7 +16,7 @@
 
 #include "src/trace_processor/stack_profile_tracker.h"
 
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "perfetto/base/logging.h"
 #include "perfetto/ext/base/string_utils.h"
diff --git a/src/trace_processor/storage/stats.h b/src/trace_processor/storage/stats.h
index 10b91aa..d500348 100644
--- a/src/trace_processor/storage/stats.h
+++ b/src/trace_processor/storage/stats.h
@@ -115,6 +115,7 @@
   F(heap_graph_non_finalized_graph,           kSingle,  kError,    kTrace),    \
   F(heap_graph_malformed_packet,              kIndexed, kError,    kTrace),    \
   F(heap_graph_missing_packet,                kIndexed, kError,    kTrace),    \
+  F(heap_graph_location_parse_error,          kSingle,  kError,    kTrace),    \
   F(heapprofd_buffer_corrupted,               kIndexed, kError,    kTrace),    \
   F(heapprofd_hit_guardrail,                  kIndexed, kError,    kTrace),    \
   F(heapprofd_buffer_overran,                 kIndexed, kDataLoss, kTrace),    \
diff --git a/src/trace_processor/syscall_tracker.h b/src/trace_processor/syscall_tracker.h
index 455b82e..23bc329 100644
--- a/src/trace_processor/syscall_tracker.h
+++ b/src/trace_processor/syscall_tracker.h
@@ -21,11 +21,11 @@
 #include <tuple>
 
 #include "perfetto/ext/base/string_view.h"
-#include "src/trace_processor/destructible.h"
 #include "src/trace_processor/slice_tracker.h"
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/destructible.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/timestamped_trace_piece.h b/src/trace_processor/timestamped_trace_piece.h
index 437a648..74b484f 100644
--- a/src/trace_processor/timestamped_trace_piece.h
+++ b/src/trace_processor/timestamped_trace_piece.h
@@ -25,7 +25,7 @@
 #include "src/trace_processor/importers/systrace/systrace_line.h"
 #include "src/trace_processor/storage/trace_storage.h"
 #include "src/trace_processor/trace_blob_view.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 // GCC can't figure out the relationship between TimestampedTracePiece's type
 // and the union, and thus thinks that we may be moving or destroying
diff --git a/src/trace_processor/trace_processor_context.cc b/src/trace_processor/trace_processor_context.cc
index 4861033..5b00c82 100644
--- a/src/trace_processor/trace_processor_context.cc
+++ b/src/trace_processor/trace_processor_context.cc
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 #include "src/trace_processor/args_tracker.h"
 #include "src/trace_processor/chunked_trace_reader.h"
@@ -24,6 +24,7 @@
 #include "src/trace_processor/global_args_tracker.h"
 #include "src/trace_processor/heap_profile_tracker.h"
 #include "src/trace_processor/importers/ftrace/ftrace_module.h"
+#include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/importers/proto/proto_trace_parser.h"
 #include "src/trace_processor/importers/proto/track_event_module.h"
 #include "src/trace_processor/metadata_tracker.h"
@@ -33,6 +34,7 @@
 #include "src/trace_processor/stack_profile_tracker.h"
 #include "src/trace_processor/trace_sorter.h"
 #include "src/trace_processor/track_tracker.h"
+#include "src/trace_processor/types/destructible.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/trace_processor_storage_impl.cc b/src/trace_processor/trace_processor_storage_impl.cc
index 85ced50..0618e09 100644
--- a/src/trace_processor/trace_processor_storage_impl.cc
+++ b/src/trace_processor/trace_processor_storage_impl.cc
@@ -23,6 +23,7 @@
 #include "src/trace_processor/event_tracker.h"
 #include "src/trace_processor/forwarding_trace_parser.h"
 #include "src/trace_processor/heap_profile_tracker.h"
+#include "src/trace_processor/importers/proto/proto_importer_module.h"
 #include "src/trace_processor/importers/proto/proto_trace_tokenizer.h"
 #include "src/trace_processor/metadata_tracker.h"
 #include "src/trace_processor/perf_sample_tracker.h"
diff --git a/src/trace_processor/trace_processor_storage_impl.h b/src/trace_processor/trace_processor_storage_impl.h
index 41f1273..1c029c6 100644
--- a/src/trace_processor/trace_processor_storage_impl.h
+++ b/src/trace_processor/trace_processor_storage_impl.h
@@ -22,7 +22,7 @@
 #include "perfetto/trace_processor/basic_types.h"
 #include "perfetto/trace_processor/status.h"
 #include "perfetto/trace_processor/trace_processor_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/trace_sorter_unittest.cc b/src/trace_processor/trace_sorter_unittest.cc
index edbd450..a7b2fc7 100644
--- a/src/trace_processor/trace_sorter_unittest.cc
+++ b/src/trace_processor/trace_sorter_unittest.cc
@@ -21,8 +21,8 @@
 
 #include "perfetto/trace_processor/basic_types.h"
 #include "src/trace_processor/timestamped_trace_piece.h"
-#include "src/trace_processor/trace_processor_context.h"
 #include "src/trace_processor/trace_sorter.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 #include "test/gtest_and_gmock.h"
 
 namespace perfetto {
diff --git a/src/trace_processor/track_tracker.h b/src/trace_processor/track_tracker.h
index dba3c8a..3618328 100644
--- a/src/trace_processor/track_tracker.h
+++ b/src/trace_processor/track_tracker.h
@@ -18,7 +18,7 @@
 #define SRC_TRACE_PROCESSOR_TRACK_TRACKER_H_
 
 #include "src/trace_processor/storage/trace_storage.h"
-#include "src/trace_processor/trace_processor_context.h"
+#include "src/trace_processor/types/trace_processor_context.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/types/BUILD.gn b/src/trace_processor/types/BUILD.gn
index e67df83..cb719b5 100644
--- a/src/trace_processor/types/BUILD.gn
+++ b/src/trace_processor/types/BUILD.gn
@@ -14,16 +14,20 @@
 
 source_set("types") {
   sources = [
+    "destructible.cc",
+    "destructible.h",
     "gfp_flags.cc",
     "gfp_flags.h",
     "task_state.cc",
     "task_state.h",
+    "trace_processor_context.h",
     "variadic.cc",
     "variadic.h",
   ]
   deps = [
     "../../../gn:default_deps",
     "../../../include/perfetto/ext/base",
+    "../../../include/perfetto/trace_processor",
     "../containers",
   ]
 }
diff --git a/src/trace_processor/destructible.cc b/src/trace_processor/types/destructible.cc
similarity index 93%
rename from src/trace_processor/destructible.cc
rename to src/trace_processor/types/destructible.cc
index 22bcf6a..3998441 100644
--- a/src/trace_processor/destructible.cc
+++ b/src/trace_processor/types/destructible.cc
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include "src/trace_processor/destructible.h"
+#include "src/trace_processor/types/destructible.h"
 
 namespace perfetto {
 namespace trace_processor {
diff --git a/src/trace_processor/destructible.h b/src/trace_processor/types/destructible.h
similarity index 87%
rename from src/trace_processor/destructible.h
rename to src/trace_processor/types/destructible.h
index 9b9c20f..ea63202 100644
--- a/src/trace_processor/destructible.h
+++ b/src/trace_processor/types/destructible.h
@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-#ifndef SRC_TRACE_PROCESSOR_DESTRUCTIBLE_H_
-#define SRC_TRACE_PROCESSOR_DESTRUCTIBLE_H_
+#ifndef SRC_TRACE_PROCESSOR_TYPES_DESTRUCTIBLE_H_
+#define SRC_TRACE_PROCESSOR_TYPES_DESTRUCTIBLE_H_
 
 namespace perfetto {
 namespace trace_processor {
@@ -33,4 +33,4 @@
 }  // namespace trace_processor
 }  // namespace perfetto
 
-#endif  // SRC_TRACE_PROCESSOR_DESTRUCTIBLE_H_
+#endif  // SRC_TRACE_PROCESSOR_TYPES_DESTRUCTIBLE_H_
diff --git a/src/trace_processor/trace_processor_context.h b/src/trace_processor/types/trace_processor_context.h
similarity index 91%
rename from src/trace_processor/trace_processor_context.h
rename to src/trace_processor/types/trace_processor_context.h
index 1f5cd64..122b686 100644
--- a/src/trace_processor/trace_processor_context.h
+++ b/src/trace_processor/types/trace_processor_context.h
@@ -14,16 +14,14 @@
  * limitations under the License.
  */
 
-#ifndef SRC_TRACE_PROCESSOR_TRACE_PROCESSOR_CONTEXT_H_
-#define SRC_TRACE_PROCESSOR_TRACE_PROCESSOR_CONTEXT_H_
+#ifndef SRC_TRACE_PROCESSOR_TYPES_TRACE_PROCESSOR_CONTEXT_H_
+#define SRC_TRACE_PROCESSOR_TYPES_TRACE_PROCESSOR_CONTEXT_H_
 
 #include <memory>
 #include <vector>
 
 #include "perfetto/trace_processor/basic_types.h"
-#include "src/trace_processor/chunked_trace_reader.h"
-#include "src/trace_processor/destructible.h"
-#include "src/trace_processor/importers/proto/proto_importer_module.h"
+#include "src/trace_processor/types/destructible.h"
 
 namespace perfetto {
 namespace trace_processor {
@@ -40,6 +38,7 @@
 class HeapProfileTracker;
 class MetadataTracker;
 class PerfSampleTracker;
+class ProtoImporterModule;
 class ProcessTracker;
 class SliceTracker;
 class TraceParser;
@@ -59,7 +58,7 @@
 
   std::unique_ptr<ChunkedTraceReader> chunk_reader;
   std::unique_ptr<TraceSorter> sorter;
-  
+
   // Keep the global tracker before the args tracker as we access the global
   // tracker in the destructor of the args tracker. Also keep it before other
   // trackers, as they may own ArgsTrackers themselves.
@@ -111,4 +110,4 @@
 }  // namespace trace_processor
 }  // namespace perfetto
 
-#endif  // SRC_TRACE_PROCESSOR_TRACE_PROCESSOR_CONTEXT_H_
+#endif  // SRC_TRACE_PROCESSOR_TYPES_TRACE_PROCESSOR_CONTEXT_H_
diff --git a/src/tracing/ipc/consumer/consumer_ipc_client_impl.cc b/src/tracing/ipc/consumer/consumer_ipc_client_impl.cc
index fc4ac95..7d60e5e 100644
--- a/src/tracing/ipc/consumer/consumer_ipc_client_impl.cc
+++ b/src/tracing/ipc/consumer/consumer_ipc_client_impl.cc
@@ -338,18 +338,54 @@
     return;
   }
 
+  auto it = pending_query_svc_reqs_.insert(pending_query_svc_reqs_.end(),
+                                           {std::move(callback), {}});
   protos::gen::QueryServiceStateRequest req;
   ipc::Deferred<protos::gen::QueryServiceStateResponse> async_response;
+  auto weak_this = weak_ptr_factory_.GetWeakPtr();
   async_response.Bind(
-      [callback](
-          ipc::AsyncResult<protos::gen::QueryServiceStateResponse> response) {
-        if (!response)
-          callback(false, TracingServiceState());
-        callback(true, response->service_state());
+      [weak_this,
+       it](ipc::AsyncResult<protos::gen::QueryServiceStateResponse> response) {
+        if (weak_this)
+          weak_this->OnQueryServiceStateResponse(std::move(response), it);
       });
   consumer_port_.QueryServiceState(req, std::move(async_response));
 }
 
+void ConsumerIPCClientImpl::OnQueryServiceStateResponse(
+    ipc::AsyncResult<protos::gen::QueryServiceStateResponse> response,
+    PendingQueryServiceRequests::iterator req_it) {
+  PERFETTO_DCHECK(req_it->callback);
+
+  if (!response) {
+    auto callback = std::move(req_it->callback);
+    pending_query_svc_reqs_.erase(req_it);
+    callback(false, TracingServiceState());
+    return;
+  }
+
+  // The QueryServiceState response can be split in several chunks if the
+  // service has several data sources. The client is supposed to merge all the
+  // replies. The easiest way to achieve this is to re-serialize the partial
+  // response and then re-decode the merged result in one shot.
+  std::vector<uint8_t>& merged_resp = req_it->merged_resp;
+  std::vector<uint8_t> part = response->service_state().SerializeAsArray();
+  merged_resp.insert(merged_resp.end(), part.begin(), part.end());
+
+  if (response.has_more())
+    return;
+
+  // All replies have been received. Decode the merged result and reply to the
+  // callback.
+  protos::gen::TracingServiceState svc_state;
+  bool ok = svc_state.ParseFromArray(merged_resp.data(), merged_resp.size());
+  if (!ok)
+    PERFETTO_ELOG("Failed to decode merged QueryServiceStateResponse");
+  auto callback = std::move(req_it->callback);
+  pending_query_svc_reqs_.erase(req_it);
+  callback(ok, std::move(svc_state));
+}
+
 void ConsumerIPCClientImpl::QueryCapabilities(
     QueryCapabilitiesCallback callback) {
   if (!connected_) {
diff --git a/src/tracing/ipc/consumer/consumer_ipc_client_impl.h b/src/tracing/ipc/consumer/consumer_ipc_client_impl.h
index eb139dc..ff305f6 100644
--- a/src/tracing/ipc/consumer/consumer_ipc_client_impl.h
+++ b/src/tracing/ipc/consumer/consumer_ipc_client_impl.h
@@ -19,6 +19,7 @@
 
 #include <stdint.h>
 
+#include <list>
 #include <vector>
 
 #include "perfetto/ext/base/scoped_file.h"
@@ -80,10 +81,23 @@
   void OnDisconnect() override;
 
  private:
+  struct PendingQueryServiceRequest {
+    QueryServiceStateCallback callback;
+
+    // All the replies will be appended here until |has_more| == false.
+    std::vector<uint8_t> merged_resp;
+  };
+
+  // List because we need stable iterators.
+  using PendingQueryServiceRequests = std::list<PendingQueryServiceRequest>;
+
   void OnReadBuffersResponse(
       ipc::AsyncResult<protos::gen::ReadBuffersResponse>);
   void OnEnableTracingResponse(
       ipc::AsyncResult<protos::gen::EnableTracingResponse>);
+  void OnQueryServiceStateResponse(
+      ipc::AsyncResult<protos::gen::QueryServiceStateResponse>,
+      PendingQueryServiceRequests::iterator);
 
   // TODO(primiano): think to dtor order, do we rely on any specific sequence?
   Consumer* const consumer_;
@@ -97,6 +111,8 @@
 
   bool connected_ = false;
 
+  PendingQueryServiceRequests pending_query_svc_reqs_;
+
   // When a packet is too big to fit into a ReadBuffersResponse IPC, the service
   // will chunk it into several IPCs, each containing few slices of the packet
   // (a packet's slice is always guaranteed to be << kIPCBufferSize). When
diff --git a/src/tracing/ipc/service/consumer_ipc_service.cc b/src/tracing/ipc/service/consumer_ipc_service.cc
index 12cbdd9..073bc74 100644
--- a/src/tracing/ipc/service/consumer_ipc_service.cc
+++ b/src/tracing/ipc/service/consumer_ipc_service.cc
@@ -206,14 +206,67 @@
     PendingQuerySvcResponses::iterator pending_response_it) {
   DeferredQueryServiceStateResponse response(std::move(*pending_response_it));
   pending_query_service_responses_.erase(pending_response_it);
-  if (success) {
+  if (!success) {
+    response.Reject();
+    return;
+  }
+
+  // The TracingServiceState object might be too big to fit into a single IPC
+  // message because it contains the DataSourceDescriptor of each data source.
+  // Here we split it in chunks to fit in the IPC limit, observing the
+  // following rule: each chunk must be invididually a valid TracingServiceState
+  // message; all the chunks concatenated together must form the original
+  // message. This is to deal with the legacy API that was just sending one
+  // whole message (failing in presence of too many data sources, b/153142114).
+  // The message is split as follows: we take the whole TracingServiceState,
+  // take out the data sources section (which is a top-level repeated field)
+  // and re-add them one-by-one. If, in the process of appending, the IPC msg
+  // size is reached, a new chunk is created. This assumes that the rest of
+  // TracingServiceState fits in one IPC message and each DataSourceDescriptor
+  // fits in the worst case in a dedicated message (which is true, because
+  // otherwise the RegisterDataSource() which passes the descriptor in the first
+  // place would fail).
+
+  std::vector<uint8_t> chunked_reply;
+
+  // Transmits the current chunk and starts a new one.
+  bool sent_eof = false;
+  auto send_chunked_reply = [&chunked_reply, &response,
+                             &sent_eof](bool has_more) {
+    PERFETTO_CHECK(!sent_eof);
+    sent_eof = !has_more;
     auto resp =
         ipc::AsyncResult<protos::gen::QueryServiceStateResponse>::Create();
-    *resp->mutable_service_state() = svc_state;
+    resp.set_has_more(has_more);
+    PERFETTO_CHECK(resp->mutable_service_state()->ParseFromArray(
+        chunked_reply.data(), chunked_reply.size()));
+    chunked_reply.clear();
     response.Resolve(std::move(resp));
-  } else {
-    response.Reject();
+  };
+
+  // Create a copy of the whole response and cut away the data_sources section.
+  protos::gen::TracingServiceState svc_state_copy = svc_state;
+  auto data_sources = std::move(*svc_state_copy.mutable_data_sources());
+  chunked_reply = svc_state_copy.SerializeAsArray();
+
+  // Now re-add them fitting within the IPC message limits (- some margin for
+  // the outer IPC frame).
+  constexpr size_t kMaxMsgSize = ipc::kIPCBufferSize - 128;
+  for (const auto& data_source : data_sources) {
+    protos::gen::TracingServiceState tmp;
+    tmp.mutable_data_sources()->emplace_back(std::move(data_source));
+    std::vector<uint8_t> chunk = tmp.SerializeAsArray();
+    if (chunked_reply.size() + chunk.size() < kMaxMsgSize) {
+      chunked_reply.insert(chunked_reply.end(), chunk.begin(), chunk.end());
+    } else {
+      send_chunked_reply(/*has_more=*/true);
+      chunked_reply = std::move(chunk);
+    }
   }
+
+  PERFETTO_DCHECK(!chunked_reply.empty());
+  send_chunked_reply(/*has_more=*/false);
+  PERFETTO_CHECK(sent_eof);
 }
 
 // Called by the service in response to a service_endpoint->Flush() request.
diff --git a/test/BUILD.gn b/test/BUILD.gn
index 771e940..c7b7c15 100644
--- a/test/BUILD.gn
+++ b/test/BUILD.gn
@@ -22,6 +22,7 @@
     ":test_helper",
     "../gn:default_deps",
     "../gn:gtest_and_gmock",
+    "../include/perfetto/ext/ipc",
     "../include/perfetto/ext/traced",
     "../include/perfetto/protozero",
     "../protos/perfetto/config:cpp",
diff --git a/test/end_to_end_integrationtest.cc b/test/end_to_end_integrationtest.cc
index 6564996..026c95f 100644
--- a/test/end_to_end_integrationtest.cc
+++ b/test/end_to_end_integrationtest.cc
@@ -27,12 +27,16 @@
 #include "perfetto/ext/base/file_utils.h"
 #include "perfetto/ext/base/pipe.h"
 #include "perfetto/ext/base/scoped_file.h"
+#include "perfetto/ext/base/string_utils.h"
 #include "perfetto/ext/base/subprocess.h"
 #include "perfetto/ext/base/temp_file.h"
+#include "perfetto/ext/ipc/basic_types.h"
 #include "perfetto/ext/traced/traced.h"
+#include "perfetto/ext/tracing/core/commit_data_request.h"
 #include "perfetto/ext/tracing/core/trace_packet.h"
 #include "perfetto/ext/tracing/ipc/default_socket.h"
 #include "perfetto/protozero/scattered_heap_buffer.h"
+#include "perfetto/tracing/core/tracing_service_state.h"
 #include "src/base/test/test_task_runner.h"
 #include "src/base/test/utils.h"
 #include "src/traced/probes/ftrace/ftrace_controller.h"
@@ -58,6 +62,7 @@
 namespace {
 
 using ::testing::ContainsRegex;
+using ::testing::ElementsAreArray;
 using ::testing::HasSubstr;
 
 constexpr size_t kBuiltinPackets = 8;
@@ -640,6 +645,71 @@
   ASSERT_TRUE(packets[1].has_for_testing());
 }
 
+// Regression test for b/153142114.
+TEST_F(PerfettoTest, QueryServiceStateLargeResponse) {
+  base::TestTaskRunner task_runner;
+
+  TestHelper helper(&task_runner);
+  helper.StartServiceIfRequired();
+  FakeProducer* producer = helper.ConnectFakeProducer();
+
+  helper.ConnectConsumer();
+  helper.WaitForConsumerConnect();
+
+  // Start a short tracing session. This is not really required for
+  // QueryServiceState. It's used only to put the FakeProducer in a state where
+  // it can call CommitData(). In turn CommitData() is used to linearize the
+  // producer with the service.
+  TraceConfig trace_config;
+  trace_config.add_buffers()->set_size_kb(1024);
+  trace_config.set_duration_ms(1);
+  auto* ds_config = trace_config.add_data_sources()->mutable_config();
+  ds_config->set_name("android.perfetto.FakeProducer");
+  helper.StartTracing(trace_config);
+  helper.WaitForProducerEnabled();
+  helper.WaitForTracingDisabled();
+
+  // Register 5 data sources with very large descriptors. Each descriptor will
+  // max out the IPC message size, so that the service has no other choice
+  // than chunking them.
+  std::map<std::string, std::string> ds_expected;
+  for (int i = 0; i < 5; i++) {
+    DataSourceDescriptor dsd;
+    std::string name = "big_ds_" + std::to_string(i);
+    dsd.set_name(name);
+    std::string descriptor(ipc::kIPCBufferSize - 64, (' ' + i) % 64);
+    dsd.set_track_event_descriptor_raw(descriptor);
+    ds_expected[name] = std::move(descriptor);
+    producer->RegisterDataSource(dsd);
+  }
+
+  // CommitData() here is used only to linearize the producer with the service.
+  // We need to make sure that all the RegisterDataSource() calls above have
+  // been seen by the service before continuing. The CommitData() callback will
+  // force a round-trip to the service and hence linearization.
+  auto all_reg_cb = task_runner.CreateCheckpoint("all_registered");
+  producer->CommitData(CommitDataRequest(), [&task_runner, &all_reg_cb] {
+    task_runner.PostTask(all_reg_cb);
+  });
+  task_runner.RunUntilCheckpoint("all_registered");
+
+  // Now invoke QueryServiceState() and wait for the reply. The service will
+  // send 6 (1 + 5) IPCs which will be merged together in
+  // producer_ipc_client_impl.cc.
+  auto svc_state = helper.QueryServiceStateAndWait();
+
+  ASSERT_GE(svc_state.producers().size(), 1u);
+
+  std::map<std::string, std::string> ds_found;
+  for (const auto& ds : svc_state.data_sources()) {
+    if (!base::StartsWith(ds.ds_descriptor().name(), "big_ds_"))
+      continue;
+    ds_found[ds.ds_descriptor().name()] =
+        ds.ds_descriptor().track_event_descriptor_raw();
+  }
+  EXPECT_THAT(ds_found, ElementsAreArray(ds_expected));
+}
+
 // Disable cmdline tests on sanitizets because they use fork() and that messes
 // up leak / races detections, which has been fixed only recently (see
 // https://github.com/google/sanitizers/issues/836 ).
diff --git a/test/fake_producer.cc b/test/fake_producer.cc
index 784b711..3645d4a 100644
--- a/test/fake_producer.cc
+++ b/test/fake_producer.cc
@@ -22,6 +22,7 @@
 #include "perfetto/base/time.h"
 #include "perfetto/ext/base/utils.h"
 #include "perfetto/ext/traced/traced.h"
+#include "perfetto/ext/tracing/core/commit_data_request.h"
 #include "perfetto/ext/tracing/core/shared_memory_arbiter.h"
 #include "perfetto/ext/tracing/core/trace_packet.h"
 #include "perfetto/ext/tracing/core/trace_writer.h"
@@ -120,12 +121,23 @@
     callback();
   });
 }
+
 // Note: this can be called on a different thread.
 void FakeProducer::ProduceEventBatch(std::function<void()> callback) {
   task_runner_->PostTask(
       [this, callback] { EmitEventBatchOnTaskRunner(callback); });
 }
 
+void FakeProducer::RegisterDataSource(const DataSourceDescriptor& desc) {
+  task_runner_->PostTask([this, desc] { endpoint_->RegisterDataSource(desc); });
+}
+
+void FakeProducer::CommitData(const CommitDataRequest& req,
+                              std::function<void()> callback) {
+  task_runner_->PostTask(
+      [this, req, callback] { endpoint_->CommitData(req, callback); });
+}
+
 void FakeProducer::OnTracingSetup() {}
 
 void FakeProducer::Flush(FlushRequestID flush_request_id,
diff --git a/test/fake_producer.h b/test/fake_producer.h
index f71d90e..82fe685 100644
--- a/test/fake_producer.h
+++ b/test/fake_producer.h
@@ -59,6 +59,9 @@
   // posts a callback when the service acknowledges the commit.
   void ProduceEventBatch(std::function<void()> callback = [] {});
 
+  void RegisterDataSource(const DataSourceDescriptor&);
+  void CommitData(const CommitDataRequest&, std::function<void()> callback);
+
   bool IsShmemProvidedByProducer() const {
     return endpoint_->IsShmemProvidedByProducer();
   }
diff --git a/test/test_helper.cc b/test/test_helper.cc
index 34114b8..e21af41 100644
--- a/test/test_helper.cc
+++ b/test/test_helper.cc
@@ -19,6 +19,7 @@
 #include "perfetto/ext/traced/traced.h"
 #include "perfetto/ext/tracing/core/trace_packet.h"
 #include "perfetto/ext/tracing/ipc/default_socket.h"
+#include "perfetto/tracing/core/tracing_service_state.h"
 
 #include "protos/perfetto/trace/trace_packet.pbzero.h"
 
@@ -176,6 +177,18 @@
                      timeout_ms);
 }
 
+TracingServiceState TestHelper::QueryServiceStateAndWait() {
+  TracingServiceState res;
+  auto checkpoint = CreateCheckpoint("query_svc_state");
+  auto callback = [&checkpoint, &res](bool, const TracingServiceState& tss) {
+    res = tss;
+    checkpoint();
+  };
+  endpoint_->QueryServiceState(callback);
+  RunUntilCheckpoint("query_svc_state");
+  return res;
+}
+
 std::function<void()> TestHelper::WrapTask(
     const std::function<void()>& function) {
   return [this, function] { task_runner_->PostTask(function); };
diff --git a/test/test_helper.h b/test/test_helper.h
index 8617466..07705a1 100644
--- a/test/test_helper.h
+++ b/test/test_helper.h
@@ -188,6 +188,7 @@
   void WaitForProducerEnabled();
   void WaitForTracingDisabled(uint32_t timeout_ms = 5000);
   void WaitForReadData(uint32_t read_count = 0, uint32_t timeout_ms = 5000);
+  TracingServiceState QueryServiceStateAndWait();
 
   std::string AddID(const std::string& checkpoint) {
     return checkpoint + "." + std::to_string(instance_num_);
diff --git a/test/trace_processor/heap_graph.textproto b/test/trace_processor/heap_graph.textproto
index 7587323..b0a304d 100644
--- a/test/trace_processor/heap_graph.textproto
+++ b/test/trace_processor/heap_graph.textproto
@@ -91,25 +91,34 @@
   timestamp: 10
   heap_graph {
     pid: 2
-    type_names {
+    location_names {
       iid: 1
-      str: "FactoryProducerDelegateImplActor"
+      str: "/data/app/ASDFG/invalid.test.android-SDASD/test.apk"
     }
-    type_names {
-      iid: 2
-      str: "Foo"
+    types {
+      id: 1
+      class_name: "FactoryProducerDelegateImplActor"
+      location_id: 1
     }
-    type_names {
-      iid: 3
-      str: "a"
+    types {
+      id: 2
+      class_name: "Foo"
+      location_id: 1
     }
-    type_names {
-      iid: 4
-      str: "a[]"
+    types {
+      id: 3
+      class_name: "a"
+      location_id: 1
     }
-    type_names {
-      iid: 5
-      str: "java.lang.Class<a[]>"
+    types {
+      id: 4
+      class_name: "a[]"
+      location_id: 1
+    }
+    types {
+      id: 5
+      class_name: "java.lang.Class<a[]>"
+      location_id: 1
     }
     field_names {
       iid: 1
@@ -129,6 +138,7 @@
 }
 packet {
   deobfuscation_mapping {
+    package_name: "invalid.test.android"
     obfuscated_classes {
       obfuscated_name: "a"
       deobfuscated_name: "DeobfuscatedA"
diff --git a/test/trace_processor/heap_graph_baseapk.textproto b/test/trace_processor/heap_graph_baseapk.textproto
new file mode 100644
index 0000000..e0ea8e9
--- /dev/null
+++ b/test/trace_processor/heap_graph_baseapk.textproto
@@ -0,0 +1,155 @@
+packet {
+  process_tree {
+    processes {
+      pid: 1
+      ppid: 0
+      cmdline: "init"
+      uid: 0
+    }
+    processes {
+      pid: 2
+      ppid: 1
+      cmdline: "system_server"
+      uid: 1000
+    }
+  }
+}
+packet {
+  timestamp: 1
+  process_stats {
+    processes {
+      pid: 2
+      rss_anon_kb: 1000
+      vm_swap_kb: 0
+      oom_score_adj: 0
+    }
+  }
+}
+packet {
+  timestamp: 10
+  process_stats {
+    processes {
+      pid: 2
+      rss_anon_kb: 1000
+      vm_swap_kb: 3000
+      oom_score_adj: 0
+    }
+  }
+}
+packet {
+  trusted_packet_sequence_id: 999
+  timestamp: 10
+  heap_graph {
+    pid: 2
+    roots {
+      root_type: ROOT_JAVA_FRAME
+      object_ids: 0x01
+      object_ids: 0x05
+    }
+    objects {
+      id: 0x01
+      type_id: 1
+      self_size: 64
+      reference_field_id: 1
+      reference_object_id: 0x02
+      reference_field_id: 3
+      reference_object_id: 0x02
+    }
+    objects {
+      id: 0x02
+      type_id: 2
+      self_size: 32
+    }
+    objects {
+      id: 0x03
+      type_id: 2
+      self_size: 128
+    }
+    objects {
+      id: 0x04
+      type_id: 3
+      self_size: 256
+      reference_field_id: 2
+      reference_object_id: 0x01
+    }
+    objects {
+      id: 0x05
+      type_id: 4
+      self_size: 256
+    }
+    objects {
+      id: 0x06
+      type_id: 5
+      self_size: 256
+    }
+    continued: true
+    index: 0
+  }
+}
+packet {
+  trusted_packet_sequence_id: 999
+  timestamp: 10
+  heap_graph {
+    pid: 2
+    location_names {
+      iid: 1
+      str: "base.apk!classes.dex"
+    }
+    types {
+      id: 1
+      class_name: "FactoryProducerDelegateImplActor"
+      location_id: 1
+    }
+    types {
+      id: 2
+      class_name: "Foo"
+      location_id: 1
+    }
+    types {
+      id: 3
+      class_name: "a"
+      location_id: 1
+    }
+    types {
+      id: 4
+      class_name: "a[]"
+      location_id: 1
+    }
+    types {
+      id: 5
+      class_name: "java.lang.Class<a[]>"
+      location_id: 1
+    }
+    field_names {
+      iid: 1
+      str: "FactoryProducerDelegateImplActor.foo"
+    }
+    field_names {
+      iid: 2
+      str: "int a.a"
+    }
+    field_names {
+      iid: 3
+      str: "a.b"
+    }
+    continued: false
+    index: 1
+  }
+}
+packet {
+  deobfuscation_mapping {
+    package_name: "invalid.test.android"
+    obfuscated_classes {
+      obfuscated_name: "a"
+      deobfuscated_name: "DeobfuscatedA"
+      obfuscated_members {
+        obfuscated_name: "a"
+        deobfuscated_name: "deobfuscatedA"
+      }
+      obfuscated_members {
+        obfuscated_name: "b"
+        deobfuscated_name: "Other.deobfuscatedA"
+      }
+    }
+  }
+}
diff --git a/test/trace_processor/heap_graph_interleaved.textproto b/test/trace_processor/heap_graph_interleaved.textproto
index eb6a482..fa736a4 100644
--- a/test/trace_processor/heap_graph_interleaved.textproto
+++ b/test/trace_processor/heap_graph_interleaved.textproto
@@ -65,9 +65,14 @@
       type_id: 1
       self_size: 64
     }
-    type_names {
+    location_names {
       iid: 1
-      str: "FactoryProducerDelegateImplActor"
+      str: "TEST:TEST"
+    }
+    types {
+      id: 1
+      class_name: "FactoryProducerDelegateImplActor"
+      location_id: 1
     }
     continued: false
     index: 0
@@ -77,17 +82,24 @@
   trusted_packet_sequence_id: 999
   heap_graph {
     pid: 2
-    type_names {
+    location_names {
       iid: 1
-      str: "FactoryProducerDelegateImplActor"
+      str: "/data/app/ASDFG/invalid.test.android-SDASD/test.apk"
     }
-    type_names {
-      iid: 2
-      str: "Foo"
+    types {
+      id: 1
+      class_name: "FactoryProducerDelegateImplActor"
+      location_id: 1
     }
-    type_names {
-      iid: 3
-      str: "a"
+    types {
+      id: 2
+      class_name: "Foo"
+      location_id: 1
+    }
+    types {
+      id: 3
+      class_name: "a"
+      location_id: 1
     }
     field_names {
       iid: 1
@@ -103,6 +115,7 @@
 }
 packet {
   deobfuscation_mapping {
+    package_name: "invalid.test.android"
     obfuscated_classes {
       obfuscated_name: "a"
       deobfuscated_name: "DeobfuscatedA"
diff --git a/test/trace_processor/heap_graph_legacy.textproto b/test/trace_processor/heap_graph_legacy.textproto
new file mode 100644
index 0000000..13fc018
--- /dev/null
+++ b/test/trace_processor/heap_graph_legacy.textproto
@@ -0,0 +1,146 @@
+packet {
+  process_tree {
+    processes {
+      pid: 1
+      ppid: 0
+      cmdline: "init"
+      uid: 0
+    }
+    processes {
+      pid: 2
+      ppid: 1
+      cmdline: "system_server"
+      uid: 1000
+    }
+  }
+}
+packet {
+  timestamp: 1
+  process_stats {
+    processes {
+      pid: 2
+      rss_anon_kb: 1000
+      vm_swap_kb: 0
+      oom_score_adj: 0
+    }
+  }
+}
+packet {
+  timestamp: 10
+  process_stats {
+    processes {
+      pid: 2
+      rss_anon_kb: 1000
+      vm_swap_kb: 3000
+      oom_score_adj: 0
+    }
+  }
+}
+packet {
+  trusted_packet_sequence_id: 999
+  timestamp: 10
+  heap_graph {
+    pid: 2
+    roots {
+      root_type: ROOT_JAVA_FRAME
+      object_ids: 0x01
+      object_ids: 0x05
+    }
+    objects {
+      id: 0x01
+      type_id: 1
+      self_size: 64
+      reference_field_id: 1
+      reference_object_id: 0x02
+      reference_field_id: 3
+      reference_object_id: 0x02
+    }
+    objects {
+      id: 0x02
+      type_id: 2
+      self_size: 32
+    }
+    objects {
+      id: 0x03
+      type_id: 2
+      self_size: 128
+    }
+    objects {
+      id: 0x04
+      type_id: 3
+      self_size: 256
+      reference_field_id: 2
+      reference_object_id: 0x01
+    }
+    objects {
+      id: 0x05
+      type_id: 4
+      self_size: 256
+    }
+    objects {
+      id: 0x06
+      type_id: 5
+      self_size: 256
+    }
+    continued: true
+    index: 0
+  }
+}
+packet {
+  trusted_packet_sequence_id: 999
+  timestamp: 10
+  heap_graph {
+    pid: 2
+    type_names {
+      iid: 1
+      str: "FactoryProducerDelegateImplActor"
+    }
+    type_names {
+      iid: 2
+      str: "Foo"
+    }
+    type_names {
+      iid: 3
+      str: "a"
+    }
+    type_names {
+      iid: 4
+      str: "a[]"
+    }
+    type_names {
+      iid: 5
+      str: "java.lang.Class<a[]>"
+    }
+    field_names {
+      iid: 1
+      str: "FactoryProducerDelegateImplActor.foo"
+    }
+    field_names {
+      iid: 2
+      str: "int a.a"
+    }
+    field_names {
+      iid: 3
+      str: "a.b"
+    }
+    continued: false
+    index: 1
+  }
+}
+packet {
+  deobfuscation_mapping {
+    package_name: "invalid.example.android"
+    obfuscated_classes {
+      obfuscated_name: "a"
+      deobfuscated_name: "DeobfuscatedA"
+      obfuscated_members {
+        obfuscated_name: "a"
+        deobfuscated_name: "deobfuscatedA"
+      }
+      obfuscated_members {
+        obfuscated_name: "b"
+        deobfuscated_name: "Other.deobfuscatedA"
+      }
+    }
+  }
+}
diff --git a/test/trace_processor/heap_graph_two_locations.out b/test/trace_processor/heap_graph_two_locations.out
new file mode 100644
index 0000000..fdd18bf
--- /dev/null
+++ b/test/trace_processor/heap_graph_two_locations.out
@@ -0,0 +1,7 @@
+"id","type","upid","graph_sample_ts","object_id","self_size","retained_size","unique_retained_size","reference_set_id","reachable","type_name","deobfuscated_type_name","root_type"
+0,"heap_graph_object",2,10,1,64,-1,-1,0,1,"FactoryProducerDelegateImplActor","[NULL]","ROOT_JAVA_FRAME"
+1,"heap_graph_object",2,10,2,32,-1,-1,2,1,"a","Foo","[NULL]"
+2,"heap_graph_object",2,10,3,128,-1,-1,2,0,"a","Foo","[NULL]"
+3,"heap_graph_object",2,10,4,256,-1,-1,2,0,"a","DeobfuscatedA","[NULL]"
+4,"heap_graph_object",2,10,5,256,-1,-1,3,1,"a[]","DeobfuscatedA[]","ROOT_JAVA_FRAME"
+5,"heap_graph_object",2,10,6,256,-1,-1,3,0,"java.lang.Class<a[]>","java.lang.Class<DeobfuscatedA[]>","[NULL]"
diff --git a/test/trace_processor/heap_graph_two_locations.textproto b/test/trace_processor/heap_graph_two_locations.textproto
new file mode 100644
index 0000000..e70e525
--- /dev/null
+++ b/test/trace_processor/heap_graph_two_locations.textproto
@@ -0,0 +1,168 @@
+packet {
+  process_tree {
+    processes {
+      pid: 1
+      ppid: 0
+      cmdline: "init"
+      uid: 0
+    }
+    processes {
+      pid: 2
+      ppid: 1
+      cmdline: "system_server"
+      uid: 1000
+    }
+  }
+}
+packet {
+  timestamp: 1
+  process_stats {
+    processes {
+      pid: 2
+      rss_anon_kb: 1000
+      vm_swap_kb: 0
+      oom_score_adj: 0
+    }
+  }
+}
+packet {
+  timestamp: 10
+  process_stats {
+    processes {
+      pid: 2
+      rss_anon_kb: 1000
+      vm_swap_kb: 3000
+      oom_score_adj: 0
+    }
+  }
+}
+packet {
+  trusted_packet_sequence_id: 999
+  timestamp: 10
+  heap_graph {
+    pid: 2
+    roots {
+      root_type: ROOT_JAVA_FRAME
+      object_ids: 0x01
+      object_ids: 0x05
+    }
+    objects {
+      id: 0x01
+      type_id: 1
+      self_size: 64
+      reference_field_id: 1
+      reference_object_id: 0x02
+      reference_field_id: 3
+      reference_object_id: 0x02
+    }
+    objects {
+      id: 0x02
+      type_id: 2
+      self_size: 32
+    }
+    objects {
+      id: 0x03
+      type_id: 2
+      self_size: 128
+    }
+    objects {
+      id: 0x04
+      type_id: 3
+      self_size: 256
+      reference_field_id: 2
+      reference_object_id: 0x01
+    }
+    objects {
+      id: 0x05
+      type_id: 4
+      self_size: 256
+    }
+    objects {
+      id: 0x06
+      type_id: 5
+      self_size: 256
+    }
+    continued: true
+    index: 0
+  }
+}
+packet {
+  trusted_packet_sequence_id: 999
+  timestamp: 10
+  heap_graph {
+    pid: 2
+    location_names {
+      iid: 1
+      str: "/data/app/ASDFG/invalid.test.android-SDASD/test.apk"
+    }
+    location_names {
+      iid: 2
+      str: "/data/app/ASDFG/invalid.test2.android-SDASD/test.apk"
+    }
+    types {
+      id: 1
+      class_name: "FactoryProducerDelegateImplActor"
+      location_id: 1
+    }
+    types {
+      id: 2
+      class_name: "a"
+      location_id: 2
+    }
+    types {
+      id: 3
+      class_name: "a"
+      location_id: 1
+    }
+    types {
+      id: 4
+      class_name: "a[]"
+      location_id: 1
+    }
+    types {
+      id: 5
+      class_name: "java.lang.Class<a[]>"
+      location_id: 1
+    }
+    field_names {
+      iid: 1
+      str: "FactoryProducerDelegateImplActor.foo"
+    }
+    field_names {
+      iid: 2
+      str: "int a.a"
+    }
+    field_names {
+      iid: 3
+      str: "a.b"
+    }
+    continued: false
+    index: 1
+  }
+}
+packet {
+  deobfuscation_mapping {
+    package_name: "invalid.test.android"
+    obfuscated_classes {
+      obfuscated_name: "a"
+      deobfuscated_name: "DeobfuscatedA"
+      obfuscated_members {
+        obfuscated_name: "a"
+        deobfuscated_name: "deobfuscatedA"
+      }
+      obfuscated_members {
+        obfuscated_name: "b"
+        deobfuscated_name: "Other.deobfuscatedA"
+      }
+    }
+  }
+}
+packet {
+  deobfuscation_mapping {
+    package_name: "invalid.test2.android"
+    obfuscated_classes {
+      obfuscated_name: "a"
+      deobfuscated_name: "Foo"
+    }
+  }
+}
diff --git a/test/trace_processor/index b/test/trace_processor/index
index e874636..efbb0aa 100644
--- a/test/trace_processor/index
+++ b/test/trace_processor/index
@@ -152,9 +152,21 @@
 heap_profile_jit.textproto heap_profile_frames.sql heap_profile_jit.out
 
 profiler_smaps.textproto profiler_smaps.sql profiler_smaps.out
+
+heap_graph_baseapk.textproto heap_graph_flamegraph.sql heap_graph_flamegraph.out
+heap_graph_baseapk.textproto heap_graph_object.sql heap_graph_object.out
+heap_graph_baseapk.textproto heap_graph_reference.sql heap_graph_reference.out
+
 heap_graph.textproto heap_graph_flamegraph.sql heap_graph_flamegraph.out
 heap_graph.textproto heap_graph_object.sql heap_graph_object.out
 heap_graph.textproto heap_graph_reference.sql heap_graph_reference.out
+heap_graph_two_locations.textproto heap_graph_object.sql heap_graph_two_locations.out
+# TODO(b/153552977): Stop supporting legacy heap graphs. These never made
+#                    it into a public release, so we should eventually stop
+#                    supporting workarounds for them.
+heap_graph_legacy.textproto heap_graph_flamegraph.sql heap_graph_flamegraph.out
+heap_graph_legacy.textproto heap_graph_object.sql heap_graph_object.out
+heap_graph_legacy.textproto heap_graph_reference.sql heap_graph_reference.out
 heap_graph_interleaved.textproto heap_graph_object.sql heap_graph_interleaved_object.out
 heap_graph_interleaved.textproto heap_graph_reference.sql heap_graph_interleaved_reference.out
 ../data/system-server-heap-graph.pftrace heap_graph_flamegraph.sql heap_graph_flamegraph_system-server-heap-graph.out
diff --git a/ui/src/assets/record.scss b/ui/src/assets/record.scss
index b2170bb..4f98127 100644
--- a/ui/src/assets/record.scss
+++ b/ui/src/assets/record.scss
@@ -273,14 +273,14 @@
   background: #fff;
   transition: opacity 0.25s ease;
   opacity: 0;
-  visibility: hidden;
+  display: none;
 
   &:not(.active) {
     max-height: 0;
   }
 
   &.active {
-    visibility: visible;
+    display: grid;
     opacity: 1;
   }
 
@@ -730,7 +730,7 @@
   }
 
   .code-snippet {
-    display: block;
+    display: grid;
     position: relative;
     padding: 0;
     margin: var(--record-section-padding);
diff --git a/ui/src/controller/selection_controller.ts b/ui/src/controller/selection_controller.ts
index 9970680..13cad15 100644
--- a/ui/src/controller/selection_controller.ts
+++ b/ui/src/controller/selection_controller.ts
@@ -161,10 +161,13 @@
         const cpu = result.columns[5].longValues![0] as number;
         const selected: SliceDetails =
             {ts: timeFromStart, dur, priority, endState, cpu, id, utid};
-        this.schedulingDetails(ts, utid).then(wakeResult => {
-          Object.assign(selected, wakeResult);
-          globals.publish('SliceDetails', selected);
-        });
+        this.schedulingDetails(ts, utid)
+            .then(wakeResult => {
+              Object.assign(selected, wakeResult);
+            })
+            .finally(() => {
+              globals.publish('SliceDetails', selected);
+            });
       }
     });
   }
diff --git a/ui/src/tracks/thread_state/controller.ts b/ui/src/tracks/thread_state/controller.ts
index b8c23ea..521da7e 100644
--- a/ui/src/tracks/thread_state/controller.ts
+++ b/ui/src/tracks/thread_state/controller.ts
@@ -135,7 +135,7 @@
         summary.ends[outIndex] = fromNs(start + dur);
         summary.state[outIndex] =
             internString(numStates === 1 ? state : 'Various states');
-        summary.cpu[outIndex] = +cols[2].doubleValues![row];
+        summary.cpu[outIndex] = +cols[2].longValues![row];
         outIndex++;
       }
     }