Add a simple setup latency profiling

PiperOrigin-RevId: 941651211
Change-Id: Id63fabbebea6b8a1e78387ce371ea3a1b3f79d5a

Author: happyCoder92

sandboxed_api/sandbox2/BUILD

@@ -659,21 +659,20 @@ cc_library(
         ":comms",
         ":fork_client",
         ":forkserver_cc_proto",
+        ":latency_stop_watch",
         ":mounts",
         ":namespace",
         ":policy",
         ":sanitizer",
+        ":setup_latency_breakdown",
         ":syscall",
         ":util",
-        ":version",
         "//sandboxed_api/sandbox2/util:bpf_helper",
         "//sandboxed_api/util:fileops",
         "//sandboxed_api/util:raw_logging",
         "@abseil-cpp//absl/base:core_headers",
         "@abseil-cpp//absl/container:flat_hash_map",
-        "@abseil-cpp//absl/container:flat_hash_set",
         "@abseil-cpp//absl/status",
-        "@abseil-cpp//absl/status:statusor",
         "@abseil-cpp//absl/strings",
         "@libcap",
     ],
@@ -689,7 +688,9 @@ cc_library(
         ":fork_client",
         ":forkedprocess",
         ":forkserver_cc_proto",
+        ":latency_stop_watch",
         ":namespace",
+        ":setup_latency_breakdown",
         ":util",
         "//sandboxed_api/util:fileops",
         "//sandboxed_api/util:raw_logging",
@@ -711,10 +712,12 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":comms",
+        ":flags",
         ":forkserver_cc_proto",
-        ":version",
+        ":setup_latency_breakdown",
         "//sandboxed_api/util:fileops",
         "@abseil-cpp//absl/base:core_headers",
+        "@abseil-cpp//absl/flags:flag",
         "@abseil-cpp//absl/log",
         "@abseil-cpp//absl/log:check",
         "@abseil-cpp//absl/status",
@@ -1309,3 +1312,36 @@ cc_library(
         "@abseil-cpp//absl/strings",
     ],
 )
+
+cc_library(
+    name = "setup_latency_breakdown",
+    srcs = ["setup_latency_breakdown.cc"],
+    hdrs = ["setup_latency_breakdown.h"],
+    copts = sapi_platform_copts(),
+    deps = [
+        ":comms",
+        "//sandboxed_api/util:raw_logging",
+        "@abseil-cpp//absl/time",
+    ],
+)
+
+cc_library(
+    name = "latency_stop_watch",
+    srcs = ["latency_stop_watch.cc"],
+    hdrs = ["latency_stop_watch.h"],
+    copts = sapi_platform_copts(),
+    deps = [
+        "@abseil-cpp//absl/time",
+    ],
+)
+
+cc_test(
+    name = "latency_stop_watch_test",
+    srcs = ["latency_stop_watch_test.cc"],
+    copts = sapi_platform_copts(),
+    deps = [
+        ":latency_stop_watch",
+        "@abseil-cpp//absl/time",
+        "@googletest//:gtest_main",
+    ],
+)

sandboxed_api/sandbox2/CMakeLists.txt

@@ -637,10 +637,12 @@ target_link_libraries(sandbox2_forkedprocess
           sandbox2::bpf_helper
           sandbox2::client
           sandbox2::fork_client
+          sandbox2::latency_stop_watch
           sandbox2::mounts
           sandbox2::namespace
           sandbox2::policy
           sandbox2::sanitizer
+          sandbox2::setup_latency_breakdown
           sandbox2::syscall
           sandbox2::util
           sapi::base
@@ -666,7 +668,9 @@ target_link_libraries(sandbox2_forkserver
           sandbox2::fork_client
           sandbox2::forkedprocess
           sandbox2::forkserver_proto
+          sandbox2::latency_stop_watch
           sandbox2::namespace
+          sandbox2::setup_latency_breakdown
           sandbox2::util
           sapi::base
           sandbox2::version
@@ -683,7 +687,9 @@ add_library(sandbox2_fork_client ${SAPI_LIB_TYPE}
 add_library(sandbox2::fork_client ALIAS sandbox2_fork_client)
 target_link_libraries(sandbox2_fork_client
   PRIVATE absl::status
+          sandbox2::flags
           sandbox2::forkserver_proto
+          sandbox2::setup_latency_breakdown
           sandbox2::version
   PUBLIC absl::core_headers
          absl::statusor
@@ -914,6 +920,29 @@ target_link_libraries(sandbox2_version
   PUBLIC absl::strings
   PRIVATE sapi::base)
 
+# sandboxed_api/sandbox2:setup_latency_breakdown
+add_library(sandbox2_setup_latency_breakdown STATIC
+  setup_latency_breakdown.cc
+  setup_latency_breakdown.h
+)
+add_library(sandbox2::setup_latency_breakdown ALIAS sandbox2_setup_latency_breakdown)
+target_link_libraries(sandbox2_setup_latency_breakdown
+  PRIVATE sapi::base
+  PUBLIC sandbox2::comms
+         sapi::raw_logging
+         absl::time
+)
+
+# sandboxed_api/sandbox2:latency_stop_watch
+add_library(sandbox2_latency_stop_watch STATIC
+  latency_stop_watch.cc
+  latency_stop_watch.h
+)
+add_library(sandbox2::latency_stop_watch ALIAS sandbox2_latency_stop_watch)
+target_link_libraries(sandbox2_latency_stop_watch
+  PRIVATE sapi::base
+  PUBLIC absl::time
+)
 
 if(BUILD_TESTING AND SAPI_BUILD_TESTING)
   add_subdirectory(testcases)
@@ -1402,6 +1431,22 @@ if(BUILD_TESTING AND SAPI_BUILD_TESTING)
     ENVIRONMENT "TEST_SRCDIR=${PROJECT_BINARY_DIR}"
   )
 
+  # sandboxed_api/sandbox2:latency_stop_watch_test
+  add_executable(sandbox2_latency_stop_watch_test
+    latency_stop_watch_test.cc
+  )
+  set_target_properties(sandbox2_latency_stop_watch_test PROPERTIES
+    OUTPUT_NAME latency_stop_watch_test
+  )
+  target_link_libraries(sandbox2_latency_stop_watch_test
+    PRIVATE sandbox2::latency_stop_watch
+            sapi::test_main
+  )
+  gtest_discover_tests_xcompile(sandbox2_latency_stop_watch_test PROPERTIES
+    ENVIRONMENT "TEST_TMPDIR=/tmp"
+    ENVIRONMENT "TEST_SRCDIR=${PROJECT_BINARY_DIR}"
+  )
+
 endif()
 
 configure_file(

sandboxed_api/sandbox2/comms.h

@@ -76,6 +76,7 @@ class Comms {
   static constexpr uint32_t kTagCreds = 0X80000202;
   static constexpr uint32_t kTagCommsUpgrade = 0x80000203;
   static constexpr uint32_t kTagCommsNoUpgrade = 0x80000204;
+  static constexpr uint32_t kTagLatencyBreakdown = 0x80000205;
 
   // Any payload size above this limit will LOG(WARNING).
   static constexpr size_t kWarnMsgSize = (256ULL << 20);

sandboxed_api/sandbox2/flags.cc

@@ -67,6 +67,11 @@ ABSL_FLAG(int, sandbox2_deadline_manager_signal, SIGRTMAX - 1,
           "Signal to use for deadline notifications - must be not otherwise "
           "used by the process (default: SIGRTMAX - 1)");
 
+// sandbox2:fork_client
+ABSL_FLAG(bool, sandbox2_log_setup_latency_breakdown, false,
+          "If set, sandbox2 will log the setup latency breakdown for each "
+          "started sandboxee");
+
 namespace sandbox2 {
 
 bool AbslParseFlag(absl::string_view text, GlobalForkserverStartModeSet* out,

sandboxed_api/sandbox2/flags.h

@@ -89,4 +89,7 @@ ABSL_DECLARE_FLAG(bool, sandbox_libunwind_crash_handler);  // retired flag
 // sandbox2/util:deadline_manager
 ABSL_DECLARE_FLAG(int, sandbox2_deadline_manager_signal);
 
+// sandbox2:fork_client
+ABSL_DECLARE_FLAG(bool, sandbox2_log_setup_latency_breakdown);
+
 #endif  // SANDBOXED_API_SANDBOX2_FLAGS_H_

sandboxed_api/sandbox2/fork_client.cc

@@ -16,21 +16,41 @@
 
 #include <sys/types.h>
 
+#include <string>
 #include <utility>
 
+#include "absl/flags/flag.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/synchronization/mutex.h"
 #include "sandboxed_api/sandbox2/comms.h"
+#include "sandboxed_api/sandbox2/flags.h"
 #include "sandboxed_api/sandbox2/forkserver.pb.h"
+#include "sandboxed_api/sandbox2/setup_latency_breakdown.h"
 #include "sandboxed_api/util/fileops.h"
 
 namespace sandbox2 {
-
+namespace {
 using ::sapi::file_util::fileops::FDCloser;
 
+void ReportLatency(const SetupLatencyBreakdown& latency_breakdown) {
+  if (!absl::GetFlag(FLAGS_sandbox2_log_setup_latency_breakdown)) {
+    return;
+  }
+  LOG(INFO) << "Setup latency breakdown:";
+  for (int i = 0; i < SetupLatencyBreakdown::kNumSetupSteps; ++i) {
+    LOG(INFO) << "  "
+              << SetupLatencyBreakdown::SetupStepToString(
+                     static_cast<SetupLatencyBreakdown::SetupStep>(i))
+              << ": "
+              << latency_breakdown.GetLatency(
+                     static_cast<SetupLatencyBreakdown::SetupStep>(i));
+  }
+}
+}  // namespace
+
 ForkClient::ForkClient(pid_t pid, Comms* comms, bool is_global)
     : pid_(pid), comms_(comms), is_global_(is_global) {
 }
@@ -117,6 +137,15 @@ absl::StatusOr<SandboxeeProcess> ForkClient::PendingRequest::Finalize(
     return absl::InternalError(
         "Receiving sandboxee PID from the ForkServer failed");
   }
+
+  // Receive setup latency tracer data and report it.
+  if (SetupLatencyBreakdown latency_breakdown;
+      latency_breakdown.Receive(setup_comms_)) {
+    ReportLatency(latency_breakdown);
+  } else {
+    LOG(WARNING) << "Failed to receive setup latency tracer data";
+  }
+
   return process;
 }
 

sandboxed_api/sandbox2/forkedprocess.cc

@@ -44,10 +44,12 @@
 #include "sandboxed_api/sandbox2/comms.h"
 #include "sandboxed_api/sandbox2/fork_client.h"
 #include "sandboxed_api/sandbox2/forkserver.pb.h"
+#include "sandboxed_api/sandbox2/latency_stop_watch.h"
 #include "sandboxed_api/sandbox2/mounts.h"
 #include "sandboxed_api/sandbox2/namespace.h"
 #include "sandboxed_api/sandbox2/policy.h"
 #include "sandboxed_api/sandbox2/sanitizer.h"
+#include "sandboxed_api/sandbox2/setup_latency_breakdown.h"
 #include "sandboxed_api/sandbox2/syscall.h"
 #include "sandboxed_api/sandbox2/util.h"
 #include "sandboxed_api/sandbox2/util/bpf_helper.h"
@@ -379,6 +381,8 @@ void ForkedProcess::SetupNamespaces(FDCloser initial_userns_fd,
       // We'll continue just in the child.
       _exit(0);
     }
+    latency_breakdown_.SetLatency(SetupLatencyBreakdown::kInitFork,
+                                  latency_stop_watch_.LapTime());
   }
   SanitizeEnvironment();
   JoinNamespaces(std::move(initial_mntns_fd), std::move(shared_netns_fd));
@@ -388,13 +392,23 @@ void ForkedProcess::SetupNamespaces(FDCloser initial_userns_fd,
   if (request_.netns_mode() == NETNS_MODE_SHARED_PER_FORKSERVER) {
     unshare_flags &= ~CLONE_NEWNET;
   }
+
+  latency_breakdown_.SetLatency(SetupLatencyBreakdown::kTillNamespacesUnshare,
+                                latency_stop_watch_.LapTime());
   SAPI_RAW_PCHECK(unshare(unshare_flags) == 0, "unsharing namespaces");
+  latency_breakdown_.SetLatency(SetupLatencyBreakdown::kNamespacesUnshare,
+                                latency_stop_watch_.LapTime());
   Namespace::InitializeNamespaces(
       uid, gid, request_.clone_flags(), Mounts(request_.mount_tree()),
       request_.hostname(), request_.allow_mount_propagation(),
       request_.allow_write_executable());
+  latency_breakdown_.SetLatency(
+      SetupLatencyBreakdown::kNamespacesInitialization,
+      latency_stop_watch_.LapTime());
   if (has_newpid) {
     LaunchInit();
+    latency_breakdown_.SetLatency(SetupLatencyBreakdown::kInitLaunch,
+                                  latency_stop_watch_.LapTime());
   }
 }
 
@@ -434,6 +448,9 @@ Comms ForkedProcess::Setup(FDCloser initial_userns_fd,
   DropAllCapabilities();
   // Send sandboxee pid.
   SAPI_RAW_CHECK(setup_comms_.SendCreds(), "Failed to send sandboxee_pid");
+  latency_breakdown_.SetLatency(SetupLatencyBreakdown::kTillAlmostDone,
+                                latency_stop_watch_.LapTime());
+  latency_breakdown_.Send(setup_comms_);
   setup_comms_.Terminate();
   if (will_exec) {
     LaunchSandboxee();

sandboxed_api/sandbox2/forkedprocess.h

@@ -21,14 +21,19 @@
 
 #include "sandboxed_api/sandbox2/comms.h"
 #include "sandboxed_api/sandbox2/forkserver.pb.h"
+#include "sandboxed_api/sandbox2/latency_stop_watch.h"
+#include "sandboxed_api/sandbox2/setup_latency_breakdown.h"
 #include "sandboxed_api/util/fileops.h"
 
 namespace sandbox2 {
 
 class ForkedProcess {
  public:
-  ForkedProcess(ForkRequest request, Comms setup_comms)
-      : request_(std::move(request)), setup_comms_(std::move(setup_comms)) {}
+  ForkedProcess(ForkRequest request, Comms setup_comms,
+                SetupLatencyBreakdown latency_breakdown)
+      : request_(std::move(request)),
+        setup_comms_(std::move(setup_comms)),
+        latency_breakdown_(std::move(latency_breakdown)) {}
   // Returns the comms channel for the newly setup sandboxee.
   Comms Setup(sapi::file_util::fileops::FDCloser initial_userns_fd,
               sapi::file_util::fileops::FDCloser initial_mntns_fd,
@@ -52,6 +57,8 @@ class ForkedProcess {
 
   ForkRequest request_;
   Comms setup_comms_;
+  SetupLatencyBreakdown latency_breakdown_;
+  LatencyStopWatch latency_stop_watch_;
   std::vector<std::string> args_;
   std::vector<std::string> envp_;
   sapi::file_util::fileops::FDCloser comms_fd_;

sandboxed_api/sandbox2/forkserver.cc

@@ -47,7 +47,9 @@
 #include "sandboxed_api/sandbox2/fork_client.h"
 #include "sandboxed_api/sandbox2/forkedprocess.h"
 #include "sandboxed_api/sandbox2/forkserver.pb.h"
+#include "sandboxed_api/sandbox2/latency_stop_watch.h"
 #include "sandboxed_api/sandbox2/namespace.h"
+#include "sandboxed_api/sandbox2/setup_latency_breakdown.h"
 #include "sandboxed_api/sandbox2/util.h"
 #include "sandboxed_api/util/fileops.h"
 #include "sandboxed_api/util/raw_logging.h"
@@ -117,6 +119,8 @@ pid_t ForkServer::ServeRequest() {
     }
     SAPI_RAW_LOG(FATAL, "Failed to receive ForkServer request");
   }
+  SetupLatencyBreakdown latency_breakdown;
+  LatencyStopWatch latency_stop_watch;
   SAPI_RAW_CHECK(fork_request.mode() != FORKSERVER_FORK_UNSPECIFIED,
                  "Forkserver mode is unspecified");
   if (fork_request.clone_flags() & CLONE_NEWNS) {
@@ -125,6 +129,8 @@ pid_t ForkServer::ServeRequest() {
   if (fork_request.netns_mode() == NETNS_MODE_SHARED_PER_FORKSERVER) {
     CreateForkserverSharedNetworkNamespace();
   }
+  latency_breakdown.SetLatency(SetupLatencyBreakdown::kSharedNamespacesCreation,
+                               latency_stop_watch.LapTime());
 
   // Create a new comms channel to coordinate the child setup.
   Comms setup_comms = [this] {
@@ -133,6 +139,8 @@ pid_t ForkServer::ServeRequest() {
                     "Failed to send setup socket");
     return Comms(setup_socketpair.sock[0].Release());
   }();
+  latency_breakdown.SetLatency(SetupLatencyBreakdown::kSetupCommsCreation,
+                               latency_stop_watch.LapTime());
 
   // We fork a child early on to do the rest of the setup.
   const bool has_namespaces = fork_request.clone_flags() & CLONE_NEWUSER;
@@ -148,9 +156,12 @@ pid_t ForkServer::ServeRequest() {
   }
   SAPI_RAW_PCHECK(pid != -1, "fork failed");
   if (pid == 0) {
+    latency_breakdown.SetLatency(SetupLatencyBreakdown::kSetupProcessFork,
+                                 latency_stop_watch.LapTime());
     // Make sure we don't use the forkserver's comms in the forked process.
     comms_->Terminate();
-    ForkedProcess forked(fork_request, std::move(setup_comms));
+    ForkedProcess forked(fork_request, std::move(setup_comms),
+                         latency_breakdown);
     *comms_ =
         forked.Setup(std::move(initial_userns_fd_),
                      std::move(initial_mntns_fd_), std::move(shared_netns_fd_));