init: Avoid killing all processes during reboot if update is in progress

Bug: 262321167
Test: Reboot device post OTA reboot
Change-Id: I5d13844e4e157d24a414dce668163bb1915bf65e
Signed-off-by: Akilesh Kailash <akailash@google.com>
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index cdff06e..9eb89b6 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -395,6 +395,10 @@
     // first-stage to decide whether to launch snapuserd.
     bool IsSnapuserdRequired();
 
+    // This is primarily used to device reboot. If OTA update is in progress,
+    // init will avoid killing processes
+    bool IsUserspaceSnapshotUpdateInProgress();
+
     enum class SnapshotDriver {
         DM_SNAPSHOT,
         DM_USER,
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 6fed09c..10d2f18 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -4349,5 +4349,16 @@
     return status.source_build_fingerprint();
 }
 
+bool SnapshotManager::IsUserspaceSnapshotUpdateInProgress() {
+    auto slot = GetCurrentSlot();
+    if (slot == Slot::Target) {
+        if (IsSnapuserdRequired()) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 }  // namespace snapshot
 }  // namespace android
diff --git a/init/reboot.cpp b/init/reboot.cpp
index a3fc534..27a7876 100644
--- a/init/reboot.cpp
+++ b/init/reboot.cpp
@@ -51,6 +51,7 @@
 #include <bootloader_message/bootloader_message.h>
 #include <cutils/android_reboot.h>
 #include <fs_mgr.h>
+#include <libsnapshot/snapshot.h>
 #include <logwrap/logwrap.h>
 #include <private/android_filesystem_config.h>
 #include <selinux/selinux.h>
@@ -422,11 +423,31 @@
     if (run_fsck && !FindPartitionsToUmount(&block_devices, &emulated_devices, false)) {
         return UMOUNT_STAT_ERROR;
     }
-
+    auto sm = snapshot::SnapshotManager::New();
+    bool ota_update_in_progress = false;
+    if (sm->IsUserspaceSnapshotUpdateInProgress()) {
+        LOG(INFO) << "OTA update in progress";
+        ota_update_in_progress = true;
+    }
     UmountStat stat = UmountPartitions(timeout - t.duration());
     if (stat != UMOUNT_STAT_SUCCESS) {
         LOG(INFO) << "umount timeout, last resort, kill all and try";
         if (DUMP_ON_UMOUNT_FAILURE) DumpUmountDebuggingInfo();
+        // Since umount timedout, we will try to kill all processes
+        // and do one more attempt to umount the partitions.
+        //
+        // However, if OTA update is in progress, we don't want
+        // to kill the snapuserd daemon as the daemon will
+        // be serving I/O requests. Killing the daemon will
+        // end up with I/O failures. If the update is in progress,
+        // we will just return the umount failure status immediately.
+        // This is ok, given the fact that killing the processes
+        // and doing an umount is just a last effort. We are
+        // still not doing fsck when all processes are killed.
+        //
+        if (ota_update_in_progress) {
+            return stat;
+        }
         KillAllProcesses();
         // even if it succeeds, still it is timeout and do not run fsck with all processes killed
         UmountStat st = UmountPartitions(0ms);