Add service protocol method to facilitate getting snapshots (flutter#32628)

Author: iskakaushik

runtime/service_protocol.cc

@@ -37,6 +37,9 @@ const std::string_view ServiceProtocol::kGetSkSLsExtensionName =
 const std::string_view
     ServiceProtocol::kEstimateRasterCacheMemoryExtensionName =
         "_flutter.estimateRasterCacheMemory";
+const std::string_view
+    ServiceProtocol::kRenderFrameWithRasterStatsExtensionName =
+        "_flutter.renderFrameWithRasterStats";
 
 static constexpr std::string_view kViewIdPrefx = "_flutterView/";
 static constexpr std::string_view kListViewsExtensionName =
@@ -56,6 +59,7 @@ ServiceProtocol::ServiceProtocol()
           kGetDisplayRefreshRateExtensionName,
           kGetSkSLsExtensionName,
           kEstimateRasterCacheMemoryExtensionName,
+          kRenderFrameWithRasterStatsExtensionName,
       }),
       handlers_mutex_(fml::SharedMutex::Create()) {}
 

runtime/service_protocol.h

@@ -29,6 +29,7 @@ class ServiceProtocol {
   static const std::string_view kGetDisplayRefreshRateExtensionName;
   static const std::string_view kGetSkSLsExtensionName;
   static const std::string_view kEstimateRasterCacheMemoryExtensionName;
+  static const std::string_view kRenderFrameWithRasterStatsExtensionName;
 
   class Handler {
    public:

shell/common/fixtures/shell_test.dart

@@ -299,3 +299,25 @@ void frameCallback(_Image, int) {
   // 'ItDoesNotCrashThatSkiaUnrefQueueDrainAfterIOManagerReset'.
   // The test is a regression test and doesn't care about images, so it is empty.
 }
+
+Picture CreateRedBox(Size size) {
+  Paint paint = Paint()
+    ..color = Color.fromARGB(255, 255, 0, 0)
+    ..style = PaintingStyle.fill;
+  PictureRecorder baseRecorder = PictureRecorder();
+  Canvas canvas = Canvas(baseRecorder);
+  canvas.drawRect(Rect.fromLTRB(0.0, 0.0, size.width, size.height), paint);
+  return baseRecorder.endRecording();
+}
+
+@pragma('vm:entry-point')
+void scene_with_red_box() {
+  PlatformDispatcher.instance.onBeginFrame = (Duration duration) {
+    SceneBuilder builder = SceneBuilder();
+    builder.pushOffset(0.0, 0.0);
+    builder.addPicture(Offset(0.0, 0.0), CreateRedBox(Size(2.0, 2.0)));
+    builder.pop();
+    PlatformDispatcher.instance.views.first.render(builder.build());
+  };
+  PlatformDispatcher.instance.scheduleFrame();
+}

shell/common/shell.cc

@@ -449,6 +449,11 @@ Shell::Shell(DartVMRef vm,
           task_runners_.GetRasterTaskRunner(),
           std::bind(&Shell::OnServiceProtocolEstimateRasterCacheMemory, this,
                     std::placeholders::_1, std::placeholders::_2)};
+  service_protocol_handlers_
+      [ServiceProtocol::kRenderFrameWithRasterStatsExtensionName] = {
+          task_runners_.GetRasterTaskRunner(),
+          std::bind(&Shell::OnServiceProtocolRenderFrameWithRasterStats, this,
+                    std::placeholders::_1, std::placeholders::_2)};
 }
 
 Shell::~Shell() {
@@ -1812,6 +1817,75 @@ bool Shell::OnServiceProtocolSetAssetBundlePath(
   return false;
 }
 
+static rapidjson::Value SerializeLayerSnapshot(
+    const LayerSnapshotData& snapshot,
+    rapidjson::Document* response) {
+  auto& allocator = response->GetAllocator();
+  rapidjson::Value result;
+  result.SetObject();
+  result.AddMember("layer_unique_id", snapshot.GetLayerUniqueId(), allocator);
+  result.AddMember("duration_micros", snapshot.GetDuration().ToMicroseconds(),
+                   allocator);
+  sk_sp<SkData> snapshot_bytes = snapshot.GetSnapshot();
+  if (snapshot_bytes) {
+    rapidjson::Value image;
+    image.SetArray();
+    const uint8_t* data =
+        reinterpret_cast<const uint8_t*>(snapshot_bytes->data());
+    for (size_t i = 0; i < snapshot_bytes->size(); i++) {
+      image.PushBack(data[i], allocator);
+    }
+    result.AddMember("snapshot", image, allocator);
+  }
+  return result;
+}
+
+bool Shell::OnServiceProtocolRenderFrameWithRasterStats(
+    const ServiceProtocol::Handler::ServiceProtocolMap& params,
+    rapidjson::Document* response) {
+  FML_DCHECK(task_runners_.GetRasterTaskRunner()->RunsTasksOnCurrentThread());
+
+  if (auto last_layer_tree = rasterizer_->GetLastLayerTree()) {
+    auto& allocator = response->GetAllocator();
+    response->SetObject();
+    response->AddMember("type", "RenderFrameWithRasterStats", allocator);
+
+    // When rendering the last layer tree, we do not need to build a frame,
+    // invariants in FrameTimingRecorder enforce that raster timings can not be
+    // set before build-end.
+    auto frame_timings_recorder = std::make_unique<FrameTimingsRecorder>();
+    const auto now = fml::TimePoint::Now();
+    frame_timings_recorder->RecordVsync(now, now);
+    frame_timings_recorder->RecordBuildStart(now);
+    frame_timings_recorder->RecordBuildEnd(now);
+
+    last_layer_tree->enable_leaf_layer_tracing(true);
+    rasterizer_->DrawLastLayerTree(std::move(frame_timings_recorder));
+    last_layer_tree->enable_leaf_layer_tracing(false);
+
+    rapidjson::Value snapshots;
+    snapshots.SetArray();
+
+    LayerSnapshotStore& store =
+        rasterizer_->compositor_context()->snapshot_store();
+    for (const LayerSnapshotData& data : store) {
+      snapshots.PushBack(SerializeLayerSnapshot(data, response), allocator);
+    }
+
+    response->AddMember("snapshots", snapshots, allocator);
+    return true;
+  } else {
+    const char* error =
+        "Failed to render the last frame with raster stats."
+        " Rasterizer does not hold a valid last layer tree."
+        " This could happen if this method was invoked before a frame was "
+        "rendered";
+    FML_DLOG(ERROR) << error;
+    ServiceProtocolFailureError(response, error);
+    return false;
+  }
+}
+
 Rasterizer::Screenshot Shell::Screenshot(
     Rasterizer::ScreenshotType screenshot_type,
     bool base64_encode) {

shell/common/shell.h

@@ -706,6 +706,15 @@ class Shell final : public PlatformView::Delegate,
       const ServiceProtocol::Handler::ServiceProtocolMap& params,
       rapidjson::Document* response);
 
+  // Service protocol handler
+  //
+  // Renders a frame and responds with various statistics pertaining to the
+  // raster call. These include time taken to raster every leaf layer and also
+  // leaf layer snapshots.
+  bool OnServiceProtocolRenderFrameWithRasterStats(
+      const ServiceProtocol::Handler::ServiceProtocolMap& params,
+      rapidjson::Document* response);
+
   // |ResourceCacheLimitItem|
   size_t GetResourceCacheLimit() override { return resource_cache_limit_; };
 

shell/common/shell_test.cc

@@ -256,24 +256,27 @@ void ShellTest::OnServiceProtocol(
     const ServiceProtocol::Handler::ServiceProtocolMap& params,
     rapidjson::Document* response) {
   std::promise<bool> finished;
-  fml::TaskRunner::RunNowOrPostTask(
-      task_runner, [shell, some_protocol, params, response, &finished]() {
-        switch (some_protocol) {
-          case ServiceProtocolEnum::kGetSkSLs:
-            shell->OnServiceProtocolGetSkSLs(params, response);
-            break;
-          case ServiceProtocolEnum::kEstimateRasterCacheMemory:
-            shell->OnServiceProtocolEstimateRasterCacheMemory(params, response);
-            break;
-          case ServiceProtocolEnum::kSetAssetBundlePath:
-            shell->OnServiceProtocolSetAssetBundlePath(params, response);
-            break;
-          case ServiceProtocolEnum::kRunInView:
-            shell->OnServiceProtocolRunInView(params, response);
-            break;
-        }
-        finished.set_value(true);
-      });
+  fml::TaskRunner::RunNowOrPostTask(task_runner, [shell, some_protocol, params,
+                                                  response, &finished]() {
+    switch (some_protocol) {
+      case ServiceProtocolEnum::kGetSkSLs:
+        shell->OnServiceProtocolGetSkSLs(params, response);
+        break;
+      case ServiceProtocolEnum::kEstimateRasterCacheMemory:
+        shell->OnServiceProtocolEstimateRasterCacheMemory(params, response);
+        break;
+      case ServiceProtocolEnum::kSetAssetBundlePath:
+        shell->OnServiceProtocolSetAssetBundlePath(params, response);
+        break;
+      case ServiceProtocolEnum::kRunInView:
+        shell->OnServiceProtocolRunInView(params, response);
+        break;
+      case ServiceProtocolEnum::kRenderFrameWithRasterStats:
+        shell->OnServiceProtocolRenderFrameWithRasterStats(params, response);
+        break;
+    }
+    finished.set_value(true);
+  });
   finished.get_future().wait();
 }
 

shell/common/shell_test.h

@@ -106,6 +106,7 @@ class ShellTest : public FixtureTest {
     kEstimateRasterCacheMemory,
     kSetAssetBundlePath,
     kRunInView,
+    kRenderFrameWithRasterStats,
   };
 
   // Helper method to test private method Shell::OnServiceProtocolGetSkSLs.

shell/common/shell_unittests.cc

@@ -2529,6 +2529,49 @@ TEST_F(ShellTest, OnServiceProtocolEstimateRasterCacheMemoryWorks) {
   DestroyShell(std::move(shell));
 }
 
+// ktz
+TEST_F(ShellTest, OnServiceProtocolRenderFrameWithRasterStatsWorks) {
+  auto settings = CreateSettingsForFixture();
+  std::unique_ptr<Shell> shell = CreateShell(settings);
+
+  // Create the surface needed by rasterizer
+  PlatformViewNotifyCreated(shell.get());
+
+  auto configuration = RunConfiguration::InferFromSettings(settings);
+  configuration.SetEntrypoint("scene_with_red_box");
+
+  RunEngine(shell.get(), std::move(configuration));
+  PumpOneFrame(shell.get());
+
+  ServiceProtocol::Handler::ServiceProtocolMap empty_params;
+  rapidjson::Document document;
+  OnServiceProtocol(
+      shell.get(), ServiceProtocolEnum::kRenderFrameWithRasterStats,
+      shell->GetTaskRunners().GetRasterTaskRunner(), empty_params, &document);
+  rapidjson::StringBuffer buffer;
+  rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
+  document.Accept(writer);
+
+  // It would be better to parse out the json and check for the validity of
+  // fields. Below checks approximate what needs to be checked, this can not be
+  // an exact check since duration will not exactly match.
+  std::string expected_json =
+      "\"snapshot\":[137,80,78,71,13,10,26,10,0,"
+      "0,0,13,73,72,68,82,0,0,0,1,0,0,0,1,8,6,0,0,0,31,21,196,137,0,0,0,1,115,"
+      "82,71,66,0,174,206,28,233,0,0,0,4,115,66,73,84,8,8,8,8,124,8,100,136,0,"
+      "0,0,13,73,68,65,84,8,153,99,248,207,192,240,31,0,5,0,1,255,171,206,54,"
+      "137,0,0,0,0,73,69,78,68,174,66,96,130]";
+  std::string actual_json = buffer.GetString();
+
+  EXPECT_THAT(actual_json, ::testing::HasSubstr(expected_json));
+  EXPECT_THAT(actual_json,
+              ::testing::HasSubstr("{\"type\":\"RenderFrameWithRasterStats\""));
+  EXPECT_THAT(actual_json, ::testing::HasSubstr("\"duration_micros\""));
+
+  PlatformViewNotifyDestroyed(shell.get());
+  DestroyShell(std::move(shell));
+}
+
 // TODO(https://github.com/flutter/flutter/issues/100273): Disabled due to
 // flakiness.
 // TODO(https://github.com/flutter/flutter/issues/100299): Fix it when
@@ -2918,8 +2961,8 @@ TEST_F(ShellTest, AssetManagerMultiSubdir) {
 
   std::vector<std::string> filenames = {
       "bad0",
-      "notgood",  // this is to make sure the pattern (.*)good(.*) only matches
-                  // things in the subdirectory
+      "notgood",  // this is to make sure the pattern (.*)good(.*) only
+                  // matches things in the subdirectory
   };
 
   std::vector<std::string> subdir_filenames = {
@@ -2994,8 +3037,9 @@ TEST_F(ShellTest, Spawn) {
   // Fulfill native function for the second Shell's entrypoint.
   fml::CountDownLatch second_latch(2);
   AddNativeCallback(
-      // The Dart native function names aren't very consistent but this is just
-      // the native function name of the second vm entrypoint in the fixture.
+      // The Dart native function names aren't very consistent but this is
+      // just the native function name of the second vm entrypoint in the
+      // fixture.
       "NotifyNative",
       CREATE_NATIVE_ENTRY([&](auto args) { second_latch.CountDown(); }));