Optimize PrepareSnapshotPartitionsForUpdate runtime

During PrepareSnapshotPartitionsForUpdate, we attempt to connect to
snapuserd with a 5s timeout, only to tell snapuserd to shutdown
immediately. If snapuserd isn't running, we will wait-out the whole 5
seconds. Change the logic to return early if socket_connect() calls
return ENOENT, indicating that snapuserd socket isn't used by any
process. This reduces allocateSpaceForPayload() time from 6s to 1s.

Test: th
Bug: 315215541
Change-Id: Ib24d7c63733a896c082ac92aaa88ad52d050a2a5

fs_mgr/libsnapshot/snapshot.cpp

@@ -3330,7 +3330,7 @@ Return SnapshotManager::CreateUpdateSnapshots(const DeltaArchiveManifest& manife
         // Terminate stale daemon if any
         std::unique_ptr<SnapuserdClient> snapuserd_client = std::move(snapuserd_client_);
         if (!snapuserd_client) {
-            snapuserd_client = SnapuserdClient::Connect(kSnapuserdSocket, 5s);
+            snapuserd_client = SnapuserdClient::TryConnect(kSnapuserdSocket, 5s);
         }
         if (snapuserd_client) {
             snapuserd_client->DetachSnapuserd();
@@ -3661,7 +3661,7 @@ std::unique_ptr<ICowWriter> SnapshotManager::OpenCompressedSnapshotWriter(
     cow_options.compression = status.compression_algorithm();
     cow_options.max_blocks = {status.device_size() / cow_options.block_size};
     cow_options.batch_write = status.batched_writes();
-    cow_options.num_compress_threads = status.enable_threading() ? 2 : 0;
+    cow_options.num_compress_threads = status.enable_threading() ? 2 : 1;
     // TODO(b/313962438) Improve op_count estimate. For now, use number of
     // blocks as an upper bound.
     cow_options.op_count_max = status.device_size() / cow_options.block_size;

fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_client.h

@@ -17,11 +17,7 @@
 #include <unistd.h>
 
 #include <chrono>
-#include <cstring>
-#include <iostream>
 #include <string>
-#include <thread>
-#include <vector>
 
 #include <android-base/unique_fd.h>
 
@@ -53,9 +49,14 @@ class SnapuserdClient {
     explicit SnapuserdClient(android::base::unique_fd&& sockfd);
     SnapuserdClient(){};
 
+    // Attempt to connect to snapsuerd, wait for the daemon to start if
+    // connection failed.
     static std::unique_ptr<SnapuserdClient> Connect(const std::string& socket_name,
                                                     std::chrono::milliseconds timeout_ms);
-
+    // Attempt to connect to snapsuerd, but does not wait for the daemon to
+    // start.
+    static std::unique_ptr<SnapuserdClient> TryConnect(const std::string& socket_name,
+                                                       std::chrono::milliseconds timeout_ms);
     bool StopSnapuserd();
 
     // Initializing a snapuserd handler is a three-step process:

fs_mgr/libsnapshot/snapuserd/snapuserd_client.cpp

@@ -27,7 +27,7 @@
 #include <unistd.h>
 
 #include <chrono>
-#include <sstream>
+#include <thread>
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
@@ -64,6 +64,40 @@ static inline bool IsRetryErrno() {
     return errno == ECONNREFUSED || errno == EINTR || errno == ENOENT;
 }
 
+std::unique_ptr<SnapuserdClient> SnapuserdClient::TryConnect(const std::string& socket_name,
+                                                             std::chrono::milliseconds timeout_ms) {
+    unique_fd fd;
+    const auto start = std::chrono::steady_clock::now();
+    while (true) {
+        fd.reset(TEMP_FAILURE_RETRY(socket_local_client(
+                socket_name.c_str(), ANDROID_SOCKET_NAMESPACE_RESERVED, SOCK_STREAM)));
+        if (fd >= 0) {
+            auto client = std::make_unique<SnapuserdClient>(std::move(fd));
+            if (!client->ValidateConnection()) {
+                return nullptr;
+            }
+            return client;
+        }
+        if (errno == ENOENT) {
+            LOG(INFO) << "Daemon socket " << socket_name
+                      << " does not exist, return without waiting.";
+            return nullptr;
+        }
+        if (errno == ECONNREFUSED) {
+            const auto now = std::chrono::steady_clock::now();
+            const auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(now - start);
+            if (elapsed >= timeout_ms) {
+                LOG(ERROR) << "Timed out connecting to snapuserd socket: " << socket_name;
+                return nullptr;
+            }
+            std::this_thread::sleep_for(10ms);
+        } else {
+            PLOG(ERROR) << "connect failed: " << socket_name;
+            return nullptr;
+        }
+    }
+}
+
 std::unique_ptr<SnapuserdClient> SnapuserdClient::Connect(const std::string& socket_name,
                                                           std::chrono::milliseconds timeout_ms) {
     unique_fd fd;

init/first_stage_init.cpp

@@ -30,6 +30,7 @@
 #include <chrono>
 #include <filesystem>
 #include <string>
+#include <thread>
 #include <vector>
 
 #include <android-base/chrono_utils.h>

init/snapuserd_transition.cpp

@@ -25,6 +25,7 @@
 #include <filesystem>
 #include <string>
 #include <string_view>
+#include <thread>
 
 #include <android-base/file.h>
 #include <android-base/logging.h>