Trace Redaction - Remap pids to synth pids

Whenever a pid maps to a process/thread outside of the target package,
replace it with the pid reserved for the cpu hosting the thread.

This will cause threads in the process tree to have no sched waking or
sched switch events. To remove those entries, an additional primitive
will be needed.

Because synth threads don't appear in the process tree, the synth
threads won't appear in the same process, so each will be treated as
a main thread.

Bug: 336807771
Change-Id: Idd87f7773ebc94569e086b6c33db07994099bf4e

7 files changed, 1155 insertions, 0 deletions

diff --git a/Android.bp b/Android.bp
index ae8225dae..c43694efc 100644
--- a/Android.bp
+++ b/Android.bp
@@ -13477,6 +13477,7 @@ filegroup {
         "src/trace_redaction/redact_process_free.cc",
         "src/trace_redaction/redact_sched_switch.cc",
         "src/trace_redaction/redact_task_newtask.cc",
+        "src/trace_redaction/remap_scheduling_events.cc",
         "src/trace_redaction/scrub_ftrace_events.cc",
         "src/trace_redaction/scrub_process_stats.cc",
         "src/trace_redaction/scrub_process_trees.cc",
@@ -13505,6 +13506,7 @@ filegroup {
         "src/trace_redaction/redact_process_free_unittest.cc",
         "src/trace_redaction/redact_sched_switch_unittest.cc",
         "src/trace_redaction/redact_task_newtask_unittest.cc",
+        "src/trace_redaction/remap_scheduling_events_unittest.cc",
         "src/trace_redaction/suspend_resume_unittest.cc",
     ],
 }
diff --git a/src/trace_redaction/BUILD.gn b/src/trace_redaction/BUILD.gn
index 8ff31a347..29c4a75e0 100644
--- a/src/trace_redaction/BUILD.gn
+++ b/src/trace_redaction/BUILD.gn
@@ -67,6 +67,8 @@ source_set("trace_redaction") {
     "redact_sched_switch.h",
     "redact_task_newtask.cc",
     "redact_task_newtask.h",
+    "remap_scheduling_events.cc",
+    "remap_scheduling_events.h",
     "scrub_ftrace_events.cc",
     "scrub_ftrace_events.h",
     "scrub_process_stats.cc",
@@ -105,6 +107,7 @@ source_set("integrationtests") {
     "filter_task_rename_integrationtest.cc",
     "prune_package_list_integrationtest.cc",
     "redact_sched_switch_integrationtest.cc",
+    "remap_scheduling_events_integrationtest.cc",
     "scrub_ftrace_events_integrationtest.cc",
     "scrub_process_stats_integrationtest.cc",
     "scrub_process_trees_integrationtest.cc",
@@ -143,6 +146,7 @@ perfetto_unittest_source_set("unittests") {
     "redact_process_free_unittest.cc",
     "redact_sched_switch_unittest.cc",
     "redact_task_newtask_unittest.cc",
+    "remap_scheduling_events_unittest.cc",
     "suspend_resume_unittest.cc",
   ]
   deps = [
diff --git a/src/trace_redaction/main.cc b/src/trace_redaction/main.cc
index 04cfa1eb0..051a424ac 100644
--- a/src/trace_redaction/main.cc
+++ b/src/trace_redaction/main.cc
@@ -32,6 +32,7 @@
 #include "src/trace_redaction/redact_process_free.h"
 #include "src/trace_redaction/redact_sched_switch.h"
 #include "src/trace_redaction/redact_task_newtask.h"
+#include "src/trace_redaction/remap_scheduling_events.h"
 #include "src/trace_redaction/scrub_ftrace_events.h"
 #include "src/trace_redaction/scrub_process_stats.h"
 #include "src/trace_redaction/scrub_process_trees.h"
@@ -88,6 +89,23 @@ static base::Status Main(std::string_view input,
   redact_ftrace_events
       ->emplace_back<RedactProcessFree::kFieldId, RedactProcessFree>();
 
+  // This set of transformations will change pids. This will break the
+  // connections between pids and the timeline (the synth threads are not in the
+  // timeline). If a transformation uses the timeline, it must be before this
+  // transformation.
+  auto* redact_sched_events = redactor.emplace_transform<RedactFtraceEvent>();
+  redact_sched_events->emplace_back<ThreadMergeRemapFtraceEventPid::kFieldId,
+                                    ThreadMergeRemapFtraceEventPid>();
+  redact_sched_events->emplace_back<ThreadMergeRemapSchedSwitchPid::kFieldId,
+                                    ThreadMergeRemapSchedSwitchPid>();
+  redact_sched_events->emplace_back<ThreadMergeRemapSchedWakingPid::kFieldId,
+                                    ThreadMergeRemapSchedWakingPid>();
+  redact_sched_events->emplace_back<
+      ThreadMergeDropField::kTaskNewtaskFieldNumber, ThreadMergeDropField>();
+  redact_sched_events
+      ->emplace_back<ThreadMergeDropField::kSchedProcessFreeFieldNumber,
+                     ThreadMergeDropField>();
+
   Context context;
   context.package_name = package_name;
 
diff --git a/src/trace_redaction/remap_scheduling_events.cc b/src/trace_redaction/remap_scheduling_events.cc
new file mode 100644
index 000000000..4817f1972
--- /dev/null
+++ b/src/trace_redaction/remap_scheduling_events.cc
@@ -0,0 +1,261 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "src/trace_redaction/remap_scheduling_events.h"
+
+#include "src/trace_redaction/proto_util.h"
+
+#include "protos/perfetto/trace/ftrace/ftrace_event.pbzero.h"
+#include "protos/perfetto/trace/ftrace/ftrace_event_bundle.pbzero.h"
+#include "protos/perfetto/trace/ftrace/sched.pbzero.h"
+
+namespace perfetto::trace_redaction {
+
+namespace {
+int32_t RemapPid(const Context& context,
+                 uint64_t timestamp,
+                 uint32_t cpu,
+                 int32_t pid) {
+  PERFETTO_DCHECK(context.package_uid.value());
+  PERFETTO_DCHECK(cpu < context.synthetic_threads->tids.size());
+
+  auto slice = context.timeline->Search(timestamp, pid);
+
+  auto expected_uid = NormalizeUid(slice.uid);
+  auto actual_uid = NormalizeUid(context.package_uid.value());
+
+  return !pid || expected_uid == actual_uid
+             ? pid
+             : context.synthetic_threads->tids[cpu];
+}
+}  // namespace
+
+base::Status ThreadMergeRemapFtraceEventPid::Redact(
+    const Context& context,
+    const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+    protozero::ProtoDecoder& event,
+    protos::pbzero::FtraceEvent* event_message) const {
+  if (!context.package_uid.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: missing package uid");
+  }
+
+  if (!context.synthetic_threads.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: missing synthetic threads");
+  }
+
+  // This should never happen. A bundle should have a cpu.
+  if (!bundle.has_cpu()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: Invalid ftrace event, missing cpu.");
+  }
+
+  if (bundle.cpu() >= context.synthetic_threads->tids.size()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: synthetic thread count");
+  }
+
+  auto timestamp =
+      event.FindField(protos::pbzero::FtraceEvent::kTimestampFieldNumber);
+
+  // This should never happen. An event should have a timestamp.
+  if (!timestamp.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapFtraceEventPid: Invalid ftrace event, missing "
+        "timestamp.");
+  }
+
+  // This handler should only be called for the pid field.
+  auto pid = event.FindField(protos::pbzero::FtraceEvent::kPidFieldNumber);
+  PERFETTO_DCHECK(pid.valid());
+
+  // The event's pid is technically a uint, but we need it as a int.
+  auto new_pid =
+      RemapPid(context, timestamp.as_uint64(), bundle.cpu(), pid.as_int32());
+  event_message->set_pid(static_cast<uint32_t>(new_pid));
+
+  return base::OkStatus();
+}
+
+base::Status ThreadMergeRemapSchedSwitchPid::Redact(
+    const Context& context,
+    const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+    protozero::ProtoDecoder& event,
+    protos::pbzero::FtraceEvent* event_message) const {
+  if (!context.package_uid.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: missing package uid");
+  }
+
+  if (!context.synthetic_threads.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: missing synthetic threads");
+  }
+
+  // This should never happen. A bundle should have a cpu.
+  if (!bundle.has_cpu()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: Invalid ftrace event, missing cpu.");
+  }
+
+  if (bundle.cpu() >= context.synthetic_threads->tids.size()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: synthetic thread count");
+  }
+
+  auto timestamp =
+      event.FindField(protos::pbzero::FtraceEvent::kTimestampFieldNumber);
+
+  // This should never happen. An event should have a timestamp.
+  if (!timestamp.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: Invalid ftrace event, missing "
+        "timestamp.");
+  }
+
+  // This handler should only be called for the sched switch field.
+  auto sched_switch =
+      event.FindField(protos::pbzero::FtraceEvent::kSchedSwitchFieldNumber);
+  PERFETTO_DCHECK(sched_switch.valid());
+
+  protozero::ProtoDecoder sched_switch_decoder(sched_switch.as_bytes());
+
+  auto old_prev_pid_field = sched_switch_decoder.FindField(
+      protos::pbzero::SchedSwitchFtraceEvent::kPrevPidFieldNumber);
+  auto old_next_pid_field = sched_switch_decoder.FindField(
+      protos::pbzero::SchedSwitchFtraceEvent::kNextPidFieldNumber);
+
+  if (!old_prev_pid_field.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: Invalid sched switch event, missing "
+        "prev pid");
+  }
+
+  if (!old_next_pid_field.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedSwitchPid: Invalid sched switch event, missing "
+        "next pid");
+  }
+
+  auto new_prev_pid_field =
+      RemapPid(context, timestamp.as_uint64(), bundle.cpu(),
+               old_prev_pid_field.as_int32());
+  auto new_next_pid_field =
+      RemapPid(context, timestamp.as_uint64(), bundle.cpu(),
+               old_next_pid_field.as_int32());
+
+  auto* sched_switch_message = event_message->set_sched_switch();
+
+  for (auto f = sched_switch_decoder.ReadField(); f.valid();
+       f = sched_switch_decoder.ReadField()) {
+    switch (f.id()) {
+      case protos::pbzero::SchedSwitchFtraceEvent::kPrevPidFieldNumber:
+        sched_switch_message->set_prev_pid(new_prev_pid_field);
+        break;
+
+      case protos::pbzero::SchedSwitchFtraceEvent::kNextPidFieldNumber:
+        sched_switch_message->set_next_pid(new_next_pid_field);
+        break;
+
+      default:
+        proto_util::AppendField(f, sched_switch_message);
+        break;
+    }
+  }
+
+  return base::OkStatus();
+}
+
+base::Status ThreadMergeRemapSchedWakingPid::Redact(
+    const Context& context,
+    const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+    protozero::ProtoDecoder& event,
+    protos::pbzero::FtraceEvent* event_message) const {
+  if (!context.package_uid.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: missing package uid");
+  }
+
+  if (!context.synthetic_threads.has_value()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: missing synthetic threads");
+  }
+
+  // This should never happen. A bundle should have a cpu.
+  if (!bundle.has_cpu()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: Invalid ftrace event, missing cpu.");
+  }
+
+  if (bundle.cpu() >= context.synthetic_threads->tids.size()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: synthetic thread count");
+  }
+
+  auto timestamp =
+      event.FindField(protos::pbzero::FtraceEvent::kTimestampFieldNumber);
+
+  // This should never happen. An event should have a timestamp.
+  if (!timestamp.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: Invalid ftrace event, missing "
+        "timestamp.");
+  }
+
+  // This handler should only be called for the sched waking field.
+  auto sched_waking =
+      event.FindField(protos::pbzero::FtraceEvent::kSchedWakingFieldNumber);
+  PERFETTO_DCHECK(sched_waking.valid());
+
+  protozero::ProtoDecoder sched_waking_decoder(sched_waking.as_bytes());
+
+  auto old_pid = sched_waking_decoder.FindField(
+      protos::pbzero::SchedWakingFtraceEvent::kPidFieldNumber);
+
+  if (!old_pid.valid()) {
+    return base::ErrStatus(
+        "ThreadMergeRemapSchedWakingPid: Invalid sched waking event, missing "
+        "pid");
+  }
+
+  auto new_pid_field = RemapPid(context, timestamp.as_uint64(), bundle.cpu(),
+                                old_pid.as_int32());
+
+  auto* sched_waking_message = event_message->set_sched_waking();
+
+  for (auto f = sched_waking_decoder.ReadField(); f.valid();
+       f = sched_waking_decoder.ReadField()) {
+    if (f.id() == protos::pbzero::SchedWakingFtraceEvent::kPidFieldNumber) {
+      sched_waking_message->set_pid(new_pid_field);
+    } else {
+      proto_util::AppendField(f, sched_waking_message);
+    }
+  }
+
+  return base::OkStatus();
+}
+
+// By doing nothing, the field gets dropped.
+base::Status ThreadMergeDropField::Redact(
+    const Context&,
+    const protos::pbzero::FtraceEventBundle::Decoder&,
+    protozero::ProtoDecoder&,
+    protos::pbzero::FtraceEvent*) const {
+  return base::OkStatus();
+}
+
+}  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/remap_scheduling_events.h b/src/trace_redaction/remap_scheduling_events.h
new file mode 100644
index 000000000..016053489
--- /dev/null
+++ b/src/trace_redaction/remap_scheduling_events.h
@@ -0,0 +1,137 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SRC_TRACE_REDACTION_REMAP_SCHEDULING_EVENTS_H_
+#define SRC_TRACE_REDACTION_REMAP_SCHEDULING_EVENTS_H_
+
+#include "perfetto/protozero/proto_decoder.h"
+#include "src/trace_redaction/redact_ftrace_event.h"
+#include "src/trace_redaction/trace_redaction_framework.h"
+
+#include "protos/perfetto/trace/ftrace/ftrace_event.pbzero.h"
+
+namespace perfetto::trace_redaction {
+
+// Reads the Ftrace event's pid and replaces it with a synthetic thread id (if
+// necessary).
+class ThreadMergeRemapFtraceEventPid : public FtraceEventRedaction {
+ public:
+  static constexpr auto kFieldId = protos::pbzero::FtraceEvent::kPidFieldNumber;
+
+  base::Status Redact(
+      const Context& context,
+      const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+      protozero::ProtoDecoder& event,
+      protos::pbzero::FtraceEvent* event_message) const override;
+};
+
+// Reads the sched switch pid and replaces it with a synthetic thread id (if
+// necessary).
+//
+//  event {
+//    timestamp: 6702093743539938
+//    pid: 0
+//    sched_switch {
+//      prev_comm: "swapper/7"
+//      prev_pid: 0
+//      prev_prio: 120
+//      prev_state: 0
+//      next_comm: "FMOD stream thr"
+//      next_pid: 7174
+//      next_prio: 104
+//    }
+//  }
+class ThreadMergeRemapSchedSwitchPid : public FtraceEventRedaction {
+ public:
+  static constexpr auto kFieldId =
+      protos::pbzero::FtraceEvent::kSchedSwitchFieldNumber;
+
+  base::Status Redact(
+      const Context& context,
+      const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+      protozero::ProtoDecoder& event,
+      protos::pbzero::FtraceEvent* event_message) const override;
+};
+
+// Reads the sched waking pid and replaces it with a synthetic thread id (if
+// necessary).
+//
+//  event {
+//    timestamp: 6702093743527386
+//    pid: 0
+//    sched_waking {
+//      comm: "FMOD stream thr"
+//      pid: 7174
+//      prio: 104
+//      success: 1
+//      target_cpu: 7
+//    }
+//  }
+class ThreadMergeRemapSchedWakingPid : public FtraceEventRedaction {
+ public:
+  static constexpr auto kFieldId =
+      protos::pbzero::FtraceEvent::kSchedWakingFieldNumber;
+
+  base::Status Redact(
+      const Context& context,
+      const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+      protozero::ProtoDecoder& event,
+      protos::pbzero::FtraceEvent* event_message) const override;
+};
+
+// Drop "new task" events because it's safe to assume that the threads always
+// exist.
+//
+//  event {
+//    timestamp: 6702094133317685
+//    pid: 6167
+//    task_newtask {
+//      pid: 7972                 <-- Pid being started
+//      comm: "adbd"
+//      clone_flags: 4001536
+//      oom_score_adj: -1000
+//    }
+//  }
+//
+// Drop "process free" events because it's safe to assume that the threads
+// always exist.
+//
+//  event {
+//    timestamp: 6702094703942898
+//    pid: 10
+//    sched_process_free {
+//      comm: "shell svc 7973"
+//      pid: 7974                 <-- Pid being freed
+//      prio: 120
+//    }
+//  }
+class ThreadMergeDropField : public FtraceEventRedaction {
+ public:
+  static constexpr auto kTaskNewtaskFieldNumber =
+      protos::pbzero::FtraceEvent::kTaskNewtaskFieldNumber;
+  static constexpr auto kSchedProcessFreeFieldNumber =
+      protos::pbzero::FtraceEvent::kSchedProcessFreeFieldNumber;
+
+  base::Status Redact(
+      const Context& context,
+      const protos::pbzero::FtraceEventBundle::Decoder& bundle,
+      protozero::ProtoDecoder& event,
+      protos::pbzero::FtraceEvent* event_message) const override;
+};
+
+}  // namespace perfetto::trace_redaction
+
+#endif  // SRC_TRACE_REDACTION_REMAP_SCHEDULING_EVENTS_H_
diff --git a/src/trace_redaction/remap_scheduling_events_integrationtest.cc b/src/trace_redaction/remap_scheduling_events_integrationtest.cc
new file mode 100644
index 000000000..2f68c95bf
--- /dev/null
+++ b/src/trace_redaction/remap_scheduling_events_integrationtest.cc
@@ -0,0 +1,292 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "src/base/test/status_matchers.h"
+#include "src/trace_redaction/collect_system_info.h"
+#include "src/trace_redaction/collect_timeline_events.h"
+#include "src/trace_redaction/find_package_uid.h"
+#include "src/trace_redaction/redact_ftrace_event.h"
+#include "src/trace_redaction/remap_scheduling_events.h"
+#include "src/trace_redaction/trace_redaction_integration_fixture.h"
+#include "test/gtest_and_gmock.h"
+
+#include "protos/perfetto/trace/ftrace/sched.pbzero.h"
+#include "protos/perfetto/trace/ftrace/task.pbzero.h"
+
+namespace perfetto::trace_redaction {
+
+// Runs ThreadMergeRemapFtraceEventPid, ThreadMergeRemapSchedSwitchPid,
+// ThreadMergeRemapSchedWakingPid, and ThreadMergeDropField to replace pids with
+// synthetic pids (for all threads outside of the target package);
+class RemapSchedulingEventsIntegrationTest
+    : public testing::Test,
+      protected TraceRedactionIntegrationFixure {
+ public:
+  static constexpr auto kPackageName =
+      "com.Unity.com.unity.multiplayer.samples.coop";
+  static constexpr uint64_t kPackageId = 10252;
+  static constexpr int32_t kPid = 7105;
+
+  // Threads belonging to pid 7105. Collected using trace processors.
+  static constexpr auto kTids = {
+      0,  // pid 0 will always be included because CPU idle uses it.
+      7105, 7111, 7112, 7113, 7114, 7115, 7116, 7117, 7118, 7119, 7120,
+      7124, 7125, 7127, 7129, 7130, 7131, 7132, 7133, 7134, 7135, 7136,
+      7137, 7139, 7141, 7142, 7143, 7144, 7145, 7146, 7147, 7148, 7149,
+      7150, 7151, 7152, 7153, 7154, 7155, 7156, 7157, 7158, 7159, 7160,
+      7161, 7162, 7163, 7164, 7165, 7166, 7167, 7171, 7172, 7174, 7178,
+      7180, 7184, 7200, 7945, 7946, 7947, 7948, 7950, 7969,
+  };
+
+ protected:
+  void SetUp() override {
+    trace_redactor()->emplace_collect<FindPackageUid>();
+
+    // In order to remap threads, we need to have synth threads.
+    trace_redactor()->emplace_collect<CollectSystemInfo>();
+    trace_redactor()->emplace_build<BuildSyntheticThreads>();
+
+    // Timeline information is needed to know if a pid belongs to a package.
+    trace_redactor()->emplace_collect<CollectTimelineEvents>();
+
+    auto* redactions = trace_redactor()->emplace_transform<RedactFtraceEvent>();
+    redactions->emplace_back<ThreadMergeRemapFtraceEventPid::kFieldId,
+                             ThreadMergeRemapFtraceEventPid>();
+    redactions->emplace_back<ThreadMergeRemapSchedSwitchPid::kFieldId,
+                             ThreadMergeRemapSchedSwitchPid>();
+    redactions->emplace_back<ThreadMergeRemapSchedWakingPid::kFieldId,
+                             ThreadMergeRemapSchedWakingPid>();
+    redactions->emplace_back<ThreadMergeDropField::kSchedProcessFreeFieldNumber,
+                             ThreadMergeDropField>();
+    redactions->emplace_back<ThreadMergeDropField::kTaskNewtaskFieldNumber,
+                             ThreadMergeDropField>();
+
+    context()->package_name = kPackageName;
+  }
+
+  struct Index {
+    // List of FtraceEvent
+    std::vector<protozero::ConstBytes> events;
+
+    // List of SchedSwitchFtraceEvent
+    std::vector<protozero::ConstBytes> events_sched_switch;
+
+    // List of SchedWakingFtraceEvent
+    std::vector<protozero::ConstBytes> events_sched_waking;
+
+    // List of SchedProcessFreeFtraceEvent
+    std::vector<protozero::ConstBytes> events_sched_process_free;
+
+    // List of TaskNewtaskFtraceEvent
+    std::vector<protozero::ConstBytes> events_task_newtask;
+  };
+
+  void UpdateFtraceIndex(protozero::ConstBytes bytes, Index* index) {
+    protos::pbzero::FtraceEventBundle::Decoder bundle(bytes);
+
+    for (auto event = bundle.event(); event; ++event) {
+      index->events.push_back(event->as_bytes());
+
+      // protos::pbzero::FtraceEvent
+      protozero::ProtoDecoder ftrace_event(event->as_bytes());
+
+      auto sched_switch = ftrace_event.FindField(
+          protos::pbzero::FtraceEvent::kSchedSwitchFieldNumber);
+      if (sched_switch.valid()) {
+        index->events_sched_switch.push_back(sched_switch.as_bytes());
+      }
+
+      auto sched_waking = ftrace_event.FindField(
+          protos::pbzero::FtraceEvent::kSchedWakingFieldNumber);
+      if (sched_waking.valid()) {
+        index->events_sched_waking.push_back(sched_waking.as_bytes());
+      }
+
+      auto sched_process_free = ftrace_event.FindField(
+          protos::pbzero::FtraceEvent::kSchedProcessFreeFieldNumber);
+      if (sched_process_free.valid()) {
+        index->events_sched_process_free.push_back(
+            sched_process_free.as_bytes());
+      }
+
+      auto task_newtask = ftrace_event.FindField(
+          protos::pbzero::FtraceEvent::kTaskNewtaskFieldNumber);
+      if (task_newtask.valid()) {
+        index->events_task_newtask.push_back(task_newtask.as_bytes());
+      }
+    }
+  }
+
+  // Bytes should be TracePacket
+  Index CreateFtraceIndex(const std::string& bytes) {
+    Index index;
+
+    protozero::ProtoDecoder packet_decoder(bytes);
+
+    for (auto packet = packet_decoder.ReadField(); packet.valid();
+         packet = packet_decoder.ReadField()) {
+      auto events = packet_decoder.FindField(
+          protos::pbzero::TracePacket::kFtraceEventsFieldNumber);
+
+      if (events.valid()) {
+        UpdateFtraceIndex(events.as_bytes(), &index);
+      }
+    }
+
+    return index;
+  }
+
+  base::StatusOr<std::string> LoadAndRedactTrace() {
+    auto source = LoadOriginal();
+
+    if (!source.ok()) {
+      return source.status();
+    }
+
+    auto redact = Redact();
+
+    if (!redact.ok()) {
+      return redact;
+    }
+
+    // Double-check the package id with the one from trace processor. If this
+    // was wrong and this check was missing, finding the problem would be much
+    // harder.
+    if (!context()->package_uid.has_value()) {
+      return base::ErrStatus("Missing package uid.");
+    }
+
+    if (context()->package_uid.value() != kPackageId) {
+      return base::ErrStatus("Unexpected package uid found.");
+    }
+
+    auto redacted = LoadRedacted();
+
+    if (redacted.ok()) {
+      return redacted;
+    }
+
+    // System info is used to initialize the synth threads. If these are wrong,
+    // then the synth threads will be wrong.
+    if (!context()->system_info.has_value()) {
+      return base::ErrStatus("Missing system info.");
+    }
+
+    if (context()->system_info->last_cpu() != 7u) {
+      return base::ErrStatus("Unexpected cpu count.");
+    }
+
+    // The synth threads should have been initialized. They will be used here to
+    // verify which threads exist in the redacted trace.
+    if (!context()->synthetic_threads.has_value()) {
+      return base::ErrStatus("Missing synthetic threads.");
+    }
+
+    if (context()->synthetic_threads->tids.size() != 8u) {
+      return base::ErrStatus("Unexpected synthentic thread count.");
+    }
+
+    return redacted;
+  }
+
+  // Should be called after redaction since it requires data from the context.
+  std::unordered_set<int32_t> CopyAllowedTids(const Context& context) const {
+    std::unordered_set<int32_t> tids(kTids.begin(), kTids.end());
+
+    tids.insert(context.synthetic_threads->tgid);
+    tids.insert(context.synthetic_threads->tids.begin(),
+                context.synthetic_threads->tids.end());
+
+    return tids;
+  }
+
+ private:
+  std::unordered_set<int32_t> allowed_tids_;
+};
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FilterFtraceEventPid) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events) {
+    protos::pbzero::FtraceEvent::Decoder decoder(event);
+    auto pid = static_cast<int32_t>(decoder.pid());
+    ASSERT_TRUE(allowlist.count(pid));
+  }
+}
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FiltersSchedSwitch) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events_sched_switch) {
+    protos::pbzero::SchedSwitchFtraceEvent::Decoder decoder(event);
+    ASSERT_TRUE(allowlist.count(decoder.prev_pid()));
+    ASSERT_TRUE(allowlist.count(decoder.next_pid()));
+  }
+}
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FiltersSchedWaking) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events_sched_waking) {
+    protos::pbzero::SchedWakingFtraceEvent::Decoder decoder(event);
+    ASSERT_TRUE(allowlist.count(decoder.pid()));
+  }
+}
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FiltersProcessFree) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events_sched_process_free) {
+    protos::pbzero::SchedProcessFreeFtraceEvent::Decoder decoder(event);
+    ASSERT_TRUE(allowlist.count(decoder.pid()));
+  }
+}
+
+TEST_F(RemapSchedulingEventsIntegrationTest, FiltersNewTask) {
+  auto redacted = LoadAndRedactTrace();
+  ASSERT_OK(redacted);
+
+  auto allowlist = CopyAllowedTids(*context());
+
+  auto index = CreateFtraceIndex(*redacted);
+
+  for (const auto& event : index.events_task_newtask) {
+    protos::pbzero::TaskNewtaskFtraceEvent::Decoder decoder(event);
+    ASSERT_TRUE(allowlist.count(decoder.pid()));
+  }
+}
+
+}  // namespace perfetto::trace_redaction
diff --git a/src/trace_redaction/remap_scheduling_events_unittest.cc b/src/trace_redaction/remap_scheduling_events_unittest.cc
new file mode 100644
index 000000000..5cdc75373
--- /dev/null
+++ b/src/trace_redaction/remap_scheduling_events_unittest.cc
@@ -0,0 +1,441 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "src/trace_redaction/remap_scheduling_events.h"
+
+#include "perfetto/protozero/scattered_heap_buffer.h"
+#include "src/base/test/status_matchers.h"
+#include "src/trace_redaction/trace_redaction_framework.h"
+#include "test/gtest_and_gmock.h"
+
+#include "protos/perfetto/trace/ftrace/ftrace_event.gen.h"
+#include "protos/perfetto/trace/ftrace/ftrace_event_bundle.gen.h"
+#include "protos/perfetto/trace/ftrace/sched.gen.h"
+
+namespace perfetto::trace_redaction {
+
+template <class T>
+class ThreadMergeTest {
+ protected:
+  struct Process {
+    uint64_t uid;
+    int32_t ppid;
+    int32_t pid;
+  };
+
+  base::Status Redact(protos::pbzero::FtraceEvent* event_message) {
+    T redact;
+
+    auto bundle_str = bundle_.SerializeAsString();
+    protos::pbzero::FtraceEventBundle::Decoder bundle_decoder(bundle_str);
+
+    auto event_str = bundle_.event().back().SerializeAsString();
+    protos::pbzero::FtraceEvent::Decoder event_decoder(event_str);
+
+    return redact.Redact(context_, bundle_decoder, event_decoder,
+                         event_message);
+  }
+
+  Context context_;
+  protos::gen::FtraceEventBundle bundle_;
+};
+
+// All ftrace events have a timestamp and a pid. This test focuses on the
+// event's pid value. When that pid doesn't belong to the target package, it
+// should be replaced with a synthetic thread id.
+//
+//  event {
+//    timestamp: 6702093743539938
+//    pid: 0
+//    sched_switch { ... }
+//  }
+class ThreadMergeRemapFtraceEventPidTest
+    : public testing::Test,
+      protected ThreadMergeTest<ThreadMergeRemapFtraceEventPid> {
+ protected:
+  static constexpr uint32_t kCpu = 3;
+
+  static constexpr auto kTimestamp = 123456789;
+
+  // This process will be connected to the target package.
+  static constexpr Process kProcess = {12, 5, 7};
+
+  // This process will not be connected to the target package.
+  static constexpr Process kOtherProcess = {120, 50, 70};
+
+  void SetUp() override {
+    bundle_.add_event();
+
+    context_.package_uid = kProcess.uid;
+
+    context_.timeline = std::make_unique<ProcessThreadTimeline>();
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kProcess.pid, kProcess.ppid, kProcess.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kOtherProcess.pid, kOtherProcess.ppid, kOtherProcess.uid));
+    context_.timeline->Sort();
+
+    // Because kCpu is 3, it means that there are four CPUs (id 0, id 1, ...).
+    context_.synthetic_threads.emplace();
+    context_.synthetic_threads->tids.assign({100, 101, 102, 103});
+  }
+};
+
+// This should never happen, a bundle should always have a cpu. If it doesn't
+// have a CPU, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapFtraceEventPidTest, MissingCpuReturnsError) {
+  // Do not call set_cpu(uint32_t value). There should be no cpu for this case.
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+// This should never happen, an event should always have a timestamp. If it
+// doesn't have a timestamp, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapFtraceEventPidTest, MissingTimestampReturnsError) {
+  bundle_.set_cpu(kCpu);
+  // Do not call set_timestamp(uint64_t value). There should be no timestamp for
+  // this case.
+  bundle_.mutable_event()->back().set_pid(kProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+TEST_F(ThreadMergeRemapFtraceEventPidTest, NoopWhenPidIsInPackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.pid()), kProcess.pid);
+}
+
+TEST_F(ThreadMergeRemapFtraceEventPidTest, ChangesPidWhenPidIsOutsidePackage) {
+  bundle_.set_cpu(kCpu);  // The CPU is used to select the pid.
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kOtherProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.pid()),
+            context_.synthetic_threads->tids[kCpu]);
+}
+
+// When creating a sched_switch event, the event pid and the previous pid should
+// be the same pid.
+//
+//  event {
+//    timestamp: 6702093743539938
+//    pid: 0
+//    sched_switch {
+//      prev_comm: "swapper/7"
+//      prev_pid: 0
+//      prev_prio: 120
+//      prev_state: 0
+//      next_comm: "FMOD stream thr"
+//      next_pid: 7174
+//      next_prio: 104
+//    }
+//  }
+class ThreadMergeRemapSchedSwitchPidTest
+    : public testing::Test,
+      protected ThreadMergeTest<ThreadMergeRemapSchedSwitchPid> {
+ protected:
+  static constexpr uint32_t kCpu = 3;
+
+  static constexpr auto kTimestamp = 123456789;
+
+  // This process will be connected to the target package.
+  static constexpr Process kPrevProcess = {12, 5, 7};
+  static constexpr Process kNextProcess = {12, 5, 8};
+
+  // This process will not be connected to the target package.
+  static constexpr Process kOtherProcess = {120, 50, 70};
+
+  void SetUp() override {
+    bundle_.add_event();
+
+    context_.package_uid = kPrevProcess.uid;
+
+    context_.timeline = std::make_unique<ProcessThreadTimeline>();
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kPrevProcess.pid, kPrevProcess.ppid, kPrevProcess.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kNextProcess.pid, kNextProcess.ppid, kNextProcess.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kOtherProcess.pid, kOtherProcess.ppid, kOtherProcess.uid));
+
+    context_.timeline->Sort();
+
+    // Because kCpu is 3, it means that there are four CPUs (id 0, id 1, ...).
+    context_.synthetic_threads.emplace();
+    context_.synthetic_threads->tids.assign({100, 101, 102, 103});
+  }
+};
+
+// This should never happen, a bundle should always have a cpu. If it doesn't
+// have a CPU, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapSchedSwitchPidTest, MissingCpuReturnsError) {
+  // Do not call set_cpu(uint32_t value). There should be no cpu for this case.
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+// This should never happen, an event should always have a timestamp. If it
+// doesn't have a timestamp, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapSchedSwitchPidTest, MissingTimestampReturnsError) {
+  bundle_.set_cpu(kCpu);
+  // Do not call set_timestamp(uint64_t value). There should be no timestamp for
+  // this case.
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+TEST_F(ThreadMergeRemapSchedSwitchPidTest, NoopWhenPidIsInPackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  auto* sched_switch = bundle_.mutable_event()->back().mutable_sched_switch();
+  sched_switch->set_prev_pid(kPrevProcess.pid);
+  sched_switch->set_next_pid(kNextProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_switch());
+
+  ASSERT_TRUE(event.sched_switch().has_prev_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().prev_pid()),
+            kPrevProcess.pid);
+
+  ASSERT_TRUE(event.sched_switch().has_next_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().next_pid()),
+            kNextProcess.pid);
+}
+
+TEST_F(ThreadMergeRemapSchedSwitchPidTest,
+       ChangesPrevPidWhenPidIsOutsidePackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  auto* sched_switch = bundle_.mutable_event()->back().mutable_sched_switch();
+  sched_switch->set_prev_pid(kOtherProcess.pid);
+  sched_switch->set_next_pid(kNextProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_switch());
+
+  ASSERT_TRUE(event.sched_switch().has_prev_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().prev_pid()),
+            context_.synthetic_threads->tids[kCpu]);
+
+  ASSERT_TRUE(event.sched_switch().has_next_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().next_pid()),
+            kNextProcess.pid);
+}
+
+TEST_F(ThreadMergeRemapSchedSwitchPidTest,
+       ChangesNextPidWhenPidIsOutsidePackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kPrevProcess.pid);
+
+  auto* sched_switch = bundle_.mutable_event()->back().mutable_sched_switch();
+  sched_switch->set_prev_pid(kPrevProcess.pid);
+  sched_switch->set_next_pid(kOtherProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_switch());
+
+  ASSERT_TRUE(event.sched_switch().has_prev_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().prev_pid()),
+            kPrevProcess.pid);
+
+  ASSERT_TRUE(event.sched_switch().has_next_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_switch().next_pid()),
+            context_.synthetic_threads->tids[kCpu]);
+}
+
+//  event {
+//    timestamp: 6702093743527386
+//    pid: 0
+//    sched_waking {
+//      comm: "FMOD stream thr"
+//      pid: 7174
+//      prio: 104
+//      success: 1
+//      target_cpu: 7
+//    }
+//  }
+class ThreadMergeRemapSchedWakingPidTest
+    : public testing::Test,
+      protected ThreadMergeTest<ThreadMergeRemapSchedWakingPid> {
+ protected:
+  static constexpr uint32_t kCpu = 3;
+
+  static constexpr auto kTimestamp = 123456789;
+
+  // This process will be connected to the target package.
+  static constexpr Process kWakerProcess = {12, 5, 7};
+  static constexpr Process kWakeTarget = {12, 5, 8};
+
+  // This process will not be connected to the target package.
+  static constexpr Process kOtherProcess = {120, 50, 70};
+
+  void SetUp() override {
+    bundle_.add_event();
+
+    context_.package_uid = kWakerProcess.uid;
+
+    context_.timeline = std::make_unique<ProcessThreadTimeline>();
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kWakerProcess.pid, kWakerProcess.ppid, kWakerProcess.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kWakeTarget.pid, kWakeTarget.ppid, kWakeTarget.uid));
+    context_.timeline->Append(ProcessThreadTimeline::Event::Open(
+        0, kOtherProcess.pid, kOtherProcess.ppid, kOtherProcess.uid));
+
+    context_.timeline->Sort();
+
+    // Because kCpu is 3, it means that there are four CPUs (id 0, id 1, ...).
+    context_.synthetic_threads.emplace();
+    context_.synthetic_threads->tids.assign({100, 101, 102, 103});
+  }
+};
+
+// This should never happen, a bundle should always have a cpu. If it doesn't
+// have a CPU, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapSchedWakingPidTest, MissingCpuReturnsError) {
+  // Do not call set_cpu(uint32_t value). There should be no cpu for this case.
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kWakerProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+// This should never happen, an event should always have a timestamp. If it
+// doesn't have a timestamp, the event field should be dropped (safest option).
+//
+// TODO(vaage): This will create an invalid trace. It can also leak information
+// if other primitives don't strip the remaining information. To be safe, these
+// cases should be replaced with errors.
+TEST_F(ThreadMergeRemapSchedWakingPidTest, MissingTimestampReturnsError) {
+  bundle_.set_cpu(kCpu);
+  // Do not call set_timestamp(uint64_t value). There should be no timestamp for
+  // this case.
+  bundle_.mutable_event()->back().set_pid(kWakerProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_FALSE(Redact(event_message.get()).ok());
+}
+
+TEST_F(ThreadMergeRemapSchedWakingPidTest, NoopWhenPidIsInPackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kWakerProcess.pid);
+
+  auto* sched_waking = bundle_.mutable_event()->back().mutable_sched_waking();
+  sched_waking->set_pid(kWakeTarget.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_waking());
+
+  ASSERT_TRUE(event.sched_waking().has_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_waking().pid()), kWakeTarget.pid);
+}
+
+TEST_F(ThreadMergeRemapSchedWakingPidTest, ChangesPidWhenPidIsOutsidePackage) {
+  bundle_.set_cpu(kCpu);
+  bundle_.mutable_event()->back().set_timestamp(kTimestamp);
+  bundle_.mutable_event()->back().set_pid(kWakerProcess.pid);
+
+  auto* sched_switch = bundle_.mutable_event()->back().mutable_sched_waking();
+  sched_switch->set_pid(kOtherProcess.pid);
+
+  protozero::HeapBuffered<protos::pbzero::FtraceEvent> event_message;
+  ASSERT_OK(Redact(event_message.get()));
+
+  protos::gen::FtraceEvent event;
+  event.ParseFromString(event_message.SerializeAsString());
+
+  ASSERT_TRUE(event.has_sched_waking());
+
+  ASSERT_TRUE(event.sched_waking().has_pid());
+  ASSERT_EQ(static_cast<int32_t>(event.sched_waking().pid()),
+            context_.synthetic_threads->tids[kCpu]);
+}
+
+}  // namespace perfetto::trace_redaction