[XProf: trace viewer] Introduce a scheduler to manage redraws in Trace Viewer v2.

frontend/app/components/trace_viewer_v2/BUILD

@@ -23,6 +23,31 @@ bzl_library(
     ],
 )
 
+wasm_cc_library(
+    name = "scheduler",
+    srcs = ["scheduler.cc"],
+    hdrs = ["scheduler.h"],
+    deps = [
+        "@com_google_absl//absl/base:no_destructor",
+    ],
+)
+
+# This test won't be run by TAP because emscripten tests require special build flags.
+# To run this test locally, use:
+# blaze test //third_party/xprof/frontend/app/components/trace_viewer_v2:scheduler_test
+wasm_cc_test(
+    name = "scheduler_test",
+    srcs = ["scheduler_test.cc"],
+    linkopts = [
+        "-sASYNCIFY=1",
+        "-sNO_EXIT_RUNTIME=1",
+    ],
+    deps = [
+        ":scheduler",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "animation",
     hdrs = ["animation.h"],
@@ -70,6 +95,7 @@ wasm_cc_library(
         ":canvas_state",
         ":event_manager",
         ":input_handler",
+        ":scheduler",
         ":webgpu_render_platform",
         "@org_xprof//frontend/app/components/trace_viewer_v2/fonts",
         "@org_xprof//frontend/app/components/trace_viewer_v2/timeline",
@@ -86,6 +112,7 @@ wasm_cc_library(
     srcs = ["input_handler.cc"],
     hdrs = ["input_handler.h"],
     deps = [
+        ":scheduler",
         "//third_party/dear_imgui",
         "//util/gtl:flat_map",
         "@com_google_absl//absl/strings",
@@ -115,6 +142,9 @@ wasm_cc_library(
     ],
 )
 
+# This test won't be run by TAP because emscripten tests require special build flags.
+# To run this test locally, use:
+# blaze test //third_party/xprof/frontend/app/components/trace_viewer_v2:event_manager_test
 wasm_cc_test(
     name = "event_manager_test",
     srcs = ["event_manager_test.cc"],
@@ -182,6 +212,7 @@ cc_binary(
         ":canvas_state",
         ":imgui_webgpu_backend",
         ":input_handler",
+        ":scheduler",
         ":webgpu_render_platform",
         "//third_party/emscripten:embind",
         "@org_xprof//frontend/app/components/trace_viewer_v2/fonts",

frontend/app/components/trace_viewer_v2/animation.h

@@ -38,6 +38,9 @@ class Animation {
     finished_->clear();
   }
 
+  // Returns true if there are any active animations.
+  static bool HasActiveAnimations() { return !animations_->empty(); }
+
  protected:
   virtual void on_finished() = 0;
 

frontend/app/components/trace_viewer_v2/application.cc

@@ -7,6 +7,7 @@
 #include <stdlib.h>
 
 #include <memory>
+#include <tuple>
 
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
@@ -16,6 +17,7 @@
 #include "xprof/frontend/app/components/trace_viewer_v2/event_manager.h"
 #include "xprof/frontend/app/components/trace_viewer_v2/fonts/fonts.h"
 #include "xprof/frontend/app/components/trace_viewer_v2/input_handler.h"  // NO_LINT
+#include "xprof/frontend/app/components/trace_viewer_v2/scheduler.h"
 #include "xprof/frontend/app/components/trace_viewer_v2/timeline/timeline.h"
 #include "xprof/frontend/app/components/trace_viewer_v2/webgpu_render_platform.h"
 
@@ -27,6 +29,15 @@ const char* const kWindowTarget = EMSCRIPTEN_EVENT_TARGET_WINDOW;
 const char* const kCanvasTarget = "#canvas";
 constexpr float kScrollbarSize = 10.0f;
 
+void ApplyDefaultTraceViewerStyles() {
+  ImGuiStyle& style = ImGui::GetStyle();
+  style.ScrollbarSize = kScrollbarSize;
+  style.WindowRounding = 0.0f;
+  style.WindowPadding = ImVec2(0.0f, 0.0f);
+  style.CellPadding = ImVec2(0.0f, 0.0f);
+  style.ItemSpacing = ImVec2(0.0f, 0.0f);
+}
+
 void ApplyLightTheme() {
   ImGui::StyleColorsLight();
   ImGuiStyle& style = ImGui::GetStyle();
@@ -39,6 +50,12 @@ void ApplyLightTheme() {
   style.Colors[ImGuiCol_TableBorderLight] = ImVec4(0.4f, 0.4f, 0.4f, 1.0f);
 }
 
+EM_BOOL OnResize(int eventType, const EmscriptenUiEvent* uiEvent,
+                 void* userData) {
+  Application::Instance().RequestRedraw();
+  return EM_FALSE;
+}
+
 }  // namespace
 
 // This function initializes the application, setting up the ImGui context,
@@ -57,7 +74,7 @@ void Application::Initialize() {
   fonts::LoadFonts(initial_canvas_state.device_pixel_ratio());
   // TODO: b/450584482 - Add a dark theme for the timeline.
   ApplyLightTheme();
-  ImGui::GetStyle().ScrollbarSize = kScrollbarSize;
+  ApplyDefaultTraceViewerStyles();
 
   // Initialize the platform
   platform_ = std::make_unique<WGPURenderPlatform>();
@@ -95,8 +112,14 @@ void Application::Initialize() {
   // Register wheel event handlers to the canvas element.
   emscripten_set_wheel_callback(kCanvasTarget, /*user_data=*/this,
                                 /*use_capture=*/true, HandleWheel);
+  emscripten_set_resize_callback(kWindowTarget, /*user_data=*/nullptr,
+                                 /*use_capture=*/true, OnResize);
+
+  Scheduler::Instance().SetMainLoopCallback([this]() { MainLoop(); });
 }
 
+void Application::RequestRedraw() { Scheduler::Instance().RequestRedraw(); }
+
 void Application::MainLoop() {
   // TODO: b/454172203 - Replace polling `CanvasState::Update()` with a
   // push-based model. Use the `ResizeObserver` API in TypeScript to listen for
@@ -117,14 +140,10 @@ void Application::MainLoop() {
   platform_->NewFrame();
   timeline_->Draw();
   platform_->RenderFrame();
-}
 
-void Application::Main() {
-  emscripten_set_main_loop_arg(
-      [](void* app) {
-        static_cast<traceviewer::Application*>(app)->MainLoop();
-      },
-      this, 0, true);
+  if (Animation::HasActiveAnimations()) {
+    RequestRedraw();
+  }
 }
 
 float Application::GetDeltaTime() {

frontend/app/components/trace_viewer_v2/application.h

@@ -35,7 +35,8 @@ class Application {
   ~Application() { ImGui::DestroyContext(); }
 
   void Initialize();
-  void Main();
+  void Main() { RequestRedraw(); }
+  void RequestRedraw();
 
   Timeline& timeline() { return *timeline_; };
   const std::vector<std::string> process_list() {

frontend/app/components/trace_viewer_v2/color/colors.h

@@ -24,14 +24,6 @@ inline constexpr ImU32 kPurple70 = 0xFFFF97C5;
 inline constexpr ImU32 kYellow90 = 0xFF7CE0FF;
 // go/keep-sorted end
 
-// Baseline palette:
-// go/keep-sorted start
-// Baseline Primary: #0B57D0
-inline constexpr ImU32 kPrimary40 = 0xFFD0570B;
-// Baseline Secondary: #C2E7FF
-inline constexpr ImU32 kSecondary90 = 0xFFFFE7C2;
-// go/keep-sorted end
-
 }  // namespace traceviewer
 
 #endif  // THIRD_PARTY_XPROF_FRONTEND_APP_COMPONENTS_TRACE_VIEWER_V2_COLOR_COLORS_H_

frontend/app/components/trace_viewer_v2/input_handler.cc

@@ -5,6 +5,7 @@
 
 #include "absl/strings/string_view.h"
 #include "third_party/dear_imgui/imgui.h"
+#include "xprof/frontend/app/components/trace_viewer_v2/scheduler.h"
 #include "util/gtl/flat_map.h"
 
 namespace traceviewer {
@@ -45,7 +46,24 @@ int IsActiveElementInput() {
   });
 }
 
+// Input Event Handlers
+//
+// Common behavior for all the following input handlers:
+//
+// 1. Redraw Requests:
+//    All handlers unconditionally request a redraw for the next animation frame
+//    via Scheduler::Instance().RequestRedraw(). Since drawing is asynchronous
+//    (scheduled for the next frame), the input updates processed here will be
+//    correctly applied before that frame is rendered.
+//
+// 2. Return Values (Event Propagation):
+//    The return value (EM_BOOL) indicates whether the event was handled:
+//    - EM_TRUE (true): The event was handled and should NOT be propagated.
+//    - EM_FALSE (false): The event was not handled and SHOULD be propagated to
+//      other listeners (e.g., browser default behavior).
+
 EM_BOOL HandleKeyDown(int, const EmscriptenKeyboardEvent* event, void*) {
+  Scheduler::Instance().RequestRedraw();
   UpdateModifierKeys(event);
 
   // If a native input element has focus, do not let ImGui capture the keyboard.
@@ -61,6 +79,7 @@ EM_BOOL HandleKeyDown(int, const EmscriptenKeyboardEvent* event, void*) {
 }
 
 EM_BOOL HandleKeyUp(int, const EmscriptenKeyboardEvent* event, void*) {
+  Scheduler::Instance().RequestRedraw();
   UpdateModifierKeys(event);
 
   // If a native input element has focus, do not let ImGui capture the keyboard.
@@ -76,24 +95,28 @@ EM_BOOL HandleKeyUp(int, const EmscriptenKeyboardEvent* event, void*) {
 }
 
 EM_BOOL HandleMouseMove(int, const EmscriptenMouseEvent* event, void*) {
+  Scheduler::Instance().RequestRedraw();
   ImGuiIO& io = ImGui::GetIO();
   io.AddMousePosEvent(event->targetX, event->targetY);
   return io.WantCaptureMouse;
 }
 
 EM_BOOL HandleMouseDown(int, const EmscriptenMouseEvent* event, void*) {
+  Scheduler::Instance().RequestRedraw();
   ImGuiIO& io = ImGui::GetIO();
   io.AddMouseButtonEvent(event->button, true);
   return io.WantCaptureMouse;
 }
 
 EM_BOOL HandleMouseUp(int, const EmscriptenMouseEvent* event, void*) {
+  Scheduler::Instance().RequestRedraw();
   ImGuiIO& io = ImGui::GetIO();
   io.AddMouseButtonEvent(event->button, false);
   return io.WantCaptureMouse;
 }
 
 EM_BOOL HandleWheel(int, const EmscriptenWheelEvent* event, void*) {
+  Scheduler::Instance().RequestRedraw();
   ImGuiIO& io = ImGui::GetIO();
 
   io.AddKeyEvent(ImGuiMod_Ctrl, event->mouse.ctrlKey);

frontend/app/components/trace_viewer_v2/scheduler.cc

@@ -0,0 +1,43 @@
+#include "xprof/frontend/app/components/trace_viewer_v2/scheduler.h"
+
+#include <emscripten/html5.h>
+
+#include <utility>
+
+#include "absl/base/no_destructor.h"
+
+namespace traceviewer {
+
+Scheduler& Scheduler::Instance() {
+  static absl::NoDestructor<Scheduler> instance;
+  return *instance;
+}
+
+void Scheduler::SetMainLoopCallback(std::function<void()> callback) {
+  callback_ = std::move(callback);
+}
+
+void Scheduler::Reset() {
+  frame_scheduled_ = false;
+  callback_ = {};
+}
+
+void Scheduler::RequestRedraw() {
+  if (!frame_scheduled_) {
+    emscripten_request_animation_frame(LoopOnce, this);
+    frame_scheduled_ = true;
+  }
+}
+
+EM_BOOL Scheduler::LoopOnce(double time, void* user_data) {
+  auto* scheduler = static_cast<Scheduler*>(user_data);
+  scheduler->frame_scheduled_ = false;
+  if (scheduler->callback_) {
+    scheduler->callback_();
+  }
+  // Return EM_FALSE because this is a demand-driven system. Redraws should
+  // only happen when explicitly requested via RequestRedraw().
+  return EM_FALSE;
+}
+
+}  // namespace traceviewer

frontend/app/components/trace_viewer_v2/scheduler.h

@@ -0,0 +1,40 @@
+#ifndef THIRD_PARTY_XPROF_FRONTEND_APP_COMPONENTS_TRACE_VIEWER_V2_SCHEDULER_H_
+#define THIRD_PARTY_XPROF_FRONTEND_APP_COMPONENTS_TRACE_VIEWER_V2_SCHEDULER_H_
+
+#include <emscripten/html5.h>
+
+#include <functional>
+
+#include "absl/base/no_destructor.h"
+
+namespace traceviewer {
+
+class Scheduler {
+ public:
+  static Scheduler& Instance();
+
+  Scheduler(const Scheduler&) = delete;
+  Scheduler& operator=(const Scheduler&) = delete;
+  Scheduler(Scheduler&&) = delete;
+  Scheduler& operator=(Scheduler&&) = delete;
+
+  void RequestRedraw();
+
+  void SetMainLoopCallback(std::function<void()> callback);
+
+  // Resets the scheduler state.
+  void Reset();
+
+ private:
+  friend class absl::NoDestructor<Scheduler>;
+  Scheduler() = default;
+
+  static EM_BOOL LoopOnce(double time, void* user_data);
+
+  bool frame_scheduled_ = false;
+  std::function<void()> callback_;
+};
+
+}  // namespace traceviewer
+
+#endif  // THIRD_PARTY_XPROF_FRONTEND_APP_COMPONENTS_TRACE_VIEWER_V2_SCHEDULER_H_

frontend/app/components/trace_viewer_v2/scheduler_test.cc

@@ -0,0 +1,50 @@
+#include "xprof/frontend/app/components/trace_viewer_v2/scheduler.h"
+
+#include <emscripten.h>
+
+#include "<gtest/gtest.h>"
+
+namespace traceviewer {
+
+namespace {
+
+class SchedulerTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    Scheduler::Instance().Reset();
+  }
+};
+
+TEST_F(SchedulerTest, CallbackIsCalledOnlyOnceForMultipleRequests) {
+  int call_count = 0;
+  Scheduler::Instance().SetMainLoopCallback([&]() { call_count++; });
+
+  Scheduler::Instance().RequestRedraw();
+  Scheduler::Instance().RequestRedraw();
+  Scheduler::Instance().RequestRedraw();
+
+  // Sleep for 100ms to allow RAF to fire.
+  // This requires Asyncify (-sASYNCIFY=1).
+  emscripten_sleep(100);
+
+  EXPECT_EQ(call_count, 1);
+}
+
+TEST_F(SchedulerTest, CallbackIsCalledAgainIfRequestedAfterFrame) {
+  int call_count = 0;
+  Scheduler::Instance().SetMainLoopCallback([&]() { call_count++; });
+
+  Scheduler::Instance().RequestRedraw();
+
+  emscripten_sleep(100);
+
+  EXPECT_EQ(call_count, 1);
+
+  Scheduler::Instance().RequestRedraw();
+  emscripten_sleep(100);
+
+  EXPECT_EQ(call_count, 2);
+}
+
+}  // namespace
+}  // namespace traceviewer