Merge "Revert "Fixed a typo in the version name""
diff --git a/CleanSpec.mk b/CleanSpec.mk
index c84bd24..0a534a2 100644
--- a/CleanSpec.mk
+++ b/CleanSpec.mk
@@ -90,3 +90,4 @@
 $(call add-clean-step, rm -rf $(PRODUCT_OUT)/debug_ramdisk/product_services)
 $(call add-clean-step, find $(PRODUCT_OUT) -type l -name "charger" -print0 | xargs -0 rm -f)
 $(call add-clean-step, rm -f $(PRODUCT_OUT)/system/bin/adbd)
+$(call add-clean-step, rm -rf $(PRODUCT_OUT)/system/etc/init/snapshotctl.rc)
diff --git a/base/logging.cpp b/base/logging.cpp
index f42b996..9360a56 100644
--- a/base/logging.cpp
+++ b/base/logging.cpp
@@ -447,7 +447,7 @@
     // See the comment in SetLogger().
     static std::atomic<AbortFunction*> abort_function(nullptr);
     auto* old_abort_function = abort_function.exchange(new AbortFunction(aborter));
-    __android_log_set_aborter([](const char* abort_message) {
+    liblog_functions->__android_log_set_aborter([](const char* abort_message) {
       auto& function = *abort_function.load(std::memory_order_acquire);
       function(abort_message);
     });
@@ -578,7 +578,7 @@
   if (liblog_functions) {
     __android_logger_data logger_data = {
         sizeof(__android_logger_data), LOG_ID_DEFAULT, priority, tag, file, line};
-    __android_log_write_logger_data(&logger_data, message);
+    liblog_functions->__android_log_write_logger_data(&logger_data, message);
   } else {
     if (tag == nullptr) {
       std::lock_guard<std::recursive_mutex> lock(TagLock());
diff --git a/debuggerd/crash_dump.cpp b/debuggerd/crash_dump.cpp
index 3e99880..6a38145 100644
--- a/debuggerd/crash_dump.cpp
+++ b/debuggerd/crash_dump.cpp
@@ -254,9 +254,7 @@
 }
 
 static void ReadCrashInfo(unique_fd& fd, siginfo_t* siginfo,
-                          std::unique_ptr<unwindstack::Regs>* regs, uintptr_t* abort_msg_address,
-                          uintptr_t* fdsan_table_address, uintptr_t* gwp_asan_state,
-                          uintptr_t* gwp_asan_metadata) {
+                          std::unique_ptr<unwindstack::Regs>* regs, ProcessInfo* process_info) {
   std::aligned_storage<sizeof(CrashInfo) + 1, alignof(CrashInfo)>::type buf;
   CrashInfo* crash_info = reinterpret_cast<CrashInfo*>(&buf);
   ssize_t rc = TEMP_FAILURE_RETRY(read(fd.get(), &buf, sizeof(buf)));
@@ -288,19 +286,16 @@
     }
   }
 
-  *fdsan_table_address = 0;
-  *gwp_asan_state = 0;
-  *gwp_asan_metadata = 0;
   switch (crash_info->header.version) {
     case 3:
-      *gwp_asan_state = crash_info->data.v3.gwp_asan_state;
-      *gwp_asan_metadata = crash_info->data.v3.gwp_asan_metadata;
+      process_info->gwp_asan_state = crash_info->data.v3.gwp_asan_state;
+      process_info->gwp_asan_metadata = crash_info->data.v3.gwp_asan_metadata;
       FALLTHROUGH_INTENDED;
     case 2:
-      *fdsan_table_address = crash_info->data.v2.fdsan_table_address;
+      process_info->fdsan_table_address = crash_info->data.v2.fdsan_table_address;
       FALLTHROUGH_INTENDED;
     case 1:
-      *abort_msg_address = crash_info->data.v1.abort_msg_address;
+      process_info->abort_msg_address = crash_info->data.v1.abort_msg_address;
       *siginfo = crash_info->data.v1.siginfo;
       regs->reset(unwindstack::Regs::CreateFromUcontext(unwindstack::Regs::CurrentArch(),
                                                         &crash_info->data.v1.ucontext));
@@ -425,10 +420,7 @@
   ATRACE_NAME("after reparent");
   pid_t pseudothread_tid;
   DebuggerdDumpType dump_type;
-  uintptr_t abort_msg_address = 0;
-  uintptr_t fdsan_table_address = 0;
-  uintptr_t gwp_asan_state = 0;
-  uintptr_t gwp_asan_metadata = 0;
+  ProcessInfo process_info;
 
   Initialize(argv);
   ParseArgs(argc, argv, &pseudothread_tid, &dump_type);
@@ -489,8 +481,7 @@
 
       if (thread == g_target_thread) {
         // Read the thread's registers along with the rest of the crash info out of the pipe.
-        ReadCrashInfo(input_pipe, &siginfo, &info.registers, &abort_msg_address,
-                      &fdsan_table_address, &gwp_asan_state, &gwp_asan_metadata);
+        ReadCrashInfo(input_pipe, &siginfo, &info.registers, &process_info);
         info.siginfo = &siginfo;
         info.signo = info.siginfo->si_signo;
       } else {
@@ -599,14 +590,14 @@
   } else {
     {
       ATRACE_NAME("fdsan table dump");
-      populate_fdsan_table(&open_files, unwinder.GetProcessMemory(), fdsan_table_address);
+      populate_fdsan_table(&open_files, unwinder.GetProcessMemory(),
+                           process_info.fdsan_table_address);
     }
 
     {
       ATRACE_NAME("engrave_tombstone");
-      engrave_tombstone(std::move(g_output_fd), &unwinder, thread_info, g_target_thread,
-                        abort_msg_address, &open_files, &amfd_data, gwp_asan_state,
-                        gwp_asan_metadata);
+      engrave_tombstone(std::move(g_output_fd), &unwinder, thread_info, g_target_thread, process_info,
+                        &open_files, &amfd_data);
     }
   }
 
diff --git a/debuggerd/include/debuggerd/handler.h b/debuggerd/include/debuggerd/handler.h
index 4f24360..665d24a 100644
--- a/debuggerd/include/debuggerd/handler.h
+++ b/debuggerd/include/debuggerd/handler.h
@@ -19,7 +19,9 @@
 #include <bionic/reserved_signals.h>
 #include <signal.h>
 #include <stdint.h>
+#include <string.h>
 #include <sys/cdefs.h>
+#include <sys/system_properties.h>
 #include <sys/types.h>
 
 __BEGIN_DECLS
@@ -50,16 +52,21 @@
 #define DEBUGGER_SIGNAL BIONIC_SIGNAL_DEBUGGER
 
 static void __attribute__((__unused__)) debuggerd_register_handlers(struct sigaction* action) {
-  sigaction(SIGABRT, action, nullptr);
-  sigaction(SIGBUS, action, nullptr);
-  sigaction(SIGFPE, action, nullptr);
-  sigaction(SIGILL, action, nullptr);
-  sigaction(SIGSEGV, action, nullptr);
-#if defined(SIGSTKFLT)
-  sigaction(SIGSTKFLT, action, nullptr);
-#endif
-  sigaction(SIGSYS, action, nullptr);
-  sigaction(SIGTRAP, action, nullptr);
+  char value[PROP_VALUE_MAX] = "";
+  bool enabled =
+      !(__system_property_get("ro.debuggable", value) > 0 && !strcmp(value, "1") &&
+        __system_property_get("debug.debuggerd.disable", value) > 0 && !strcmp(value, "1"));
+  if (enabled) {
+    sigaction(SIGABRT, action, nullptr);
+    sigaction(SIGBUS, action, nullptr);
+    sigaction(SIGFPE, action, nullptr);
+    sigaction(SIGILL, action, nullptr);
+    sigaction(SIGSEGV, action, nullptr);
+    sigaction(SIGSTKFLT, action, nullptr);
+    sigaction(SIGSYS, action, nullptr);
+    sigaction(SIGTRAP, action, nullptr);
+  }
+
   sigaction(BIONIC_SIGNAL_DEBUGGER, action, nullptr);
 }
 
diff --git a/debuggerd/libdebuggerd/gwp_asan.cpp b/debuggerd/libdebuggerd/gwp_asan.cpp
index 53df783..fe3a173 100644
--- a/debuggerd/libdebuggerd/gwp_asan.cpp
+++ b/debuggerd/libdebuggerd/gwp_asan.cpp
@@ -63,12 +63,11 @@
 }
 
 GwpAsanCrashData::GwpAsanCrashData(unwindstack::Memory* process_memory,
-                                   uintptr_t gwp_asan_state_ptr, uintptr_t gwp_asan_metadata_ptr,
-                                   const ThreadInfo& thread_info) {
-  if (!process_memory || !gwp_asan_metadata_ptr || !gwp_asan_state_ptr) return;
+                                   const ProcessInfo& process_info, const ThreadInfo& thread_info) {
+  if (!process_memory || !process_info.gwp_asan_metadata || !process_info.gwp_asan_state) return;
   // Extract the GWP-ASan regions from the dead process.
-  if (!retrieve_gwp_asan_state(process_memory, gwp_asan_state_ptr, &state_)) return;
-  metadata_.reset(retrieve_gwp_asan_metadata(process_memory, state_, gwp_asan_metadata_ptr));
+  if (!retrieve_gwp_asan_state(process_memory, process_info.gwp_asan_state, &state_)) return;
+  metadata_.reset(retrieve_gwp_asan_metadata(process_memory, state_, process_info.gwp_asan_metadata));
   if (!metadata_.get()) return;
 
   // Get the external crash address from the thread info.
diff --git a/debuggerd/libdebuggerd/include/libdebuggerd/gwp_asan.h b/debuggerd/libdebuggerd/include/libdebuggerd/gwp_asan.h
index aef4c62..6c88733 100644
--- a/debuggerd/libdebuggerd/include/libdebuggerd/gwp_asan.h
+++ b/debuggerd/libdebuggerd/include/libdebuggerd/gwp_asan.h
@@ -38,8 +38,8 @@
   // still be responsible, as it terminates when it detects an internal error
   // (double free, invalid free). In these cases, we will retrieve the fault
   // address from the GWP-ASan allocator's state.
-  GwpAsanCrashData(unwindstack::Memory* process_memory, uintptr_t gwp_asan_state_ptr,
-                   uintptr_t gwp_asan_metadata_ptr, const ThreadInfo& thread_info);
+  GwpAsanCrashData(unwindstack::Memory* process_memory, const ProcessInfo& process_info,
+                   const ThreadInfo& thread_info);
 
   // Is GWP-ASan responsible for this crash.
   bool CrashIsMine() const;
diff --git a/debuggerd/libdebuggerd/include/libdebuggerd/tombstone.h b/debuggerd/libdebuggerd/include/libdebuggerd/tombstone.h
index 291d994..3ff7d62 100644
--- a/debuggerd/libdebuggerd/include/libdebuggerd/tombstone.h
+++ b/debuggerd/libdebuggerd/include/libdebuggerd/tombstone.h
@@ -44,18 +44,13 @@
 int open_tombstone(std::string* path);
 
 /* Creates a tombstone file and writes the crash dump to it. */
-void engrave_tombstone(int tombstone_fd, unwindstack::Unwinder* unwinder,
-                       const OpenFilesList* open_files, pid_t pid, pid_t tid,
-                       const std::string& process_name, const std::map<pid_t, std::string>& threads,
-                       uint64_t abort_msg_address, std::string* amfd_data);
+void engrave_tombstone(android::base::unique_fd output_fd, unwindstack::Unwinder* unwinder,
+                       const std::map<pid_t, ThreadInfo>& thread_info, pid_t target_thread,
+                       const ProcessInfo& process_info, OpenFilesList* open_files,
+                       std::string* amfd_data);
 
 void engrave_tombstone_ucontext(int tombstone_fd, uint64_t abort_msg_address, siginfo_t* siginfo,
                                 ucontext_t* ucontext);
 
-void engrave_tombstone(android::base::unique_fd output_fd, unwindstack::Unwinder* unwinder,
-                       const std::map<pid_t, ThreadInfo>& thread_info, pid_t target_thread,
-                       uint64_t abort_msg_address, OpenFilesList* open_files,
-                       std::string* amfd_data, uintptr_t gwp_asan_state,
-                       uintptr_t gwp_asan_metadata);
 
 #endif  // _DEBUGGERD_TOMBSTONE_H
diff --git a/debuggerd/libdebuggerd/include/libdebuggerd/types.h b/debuggerd/libdebuggerd/include/libdebuggerd/types.h
index eb4b1b8..4f681c2 100644
--- a/debuggerd/libdebuggerd/include/libdebuggerd/types.h
+++ b/debuggerd/libdebuggerd/include/libdebuggerd/types.h
@@ -35,3 +35,10 @@
   int signo = 0;
   siginfo_t* siginfo = nullptr;
 };
+
+struct ProcessInfo {
+  uintptr_t abort_msg_address = 0;
+  uintptr_t fdsan_table_address = 0;
+  uintptr_t gwp_asan_state = 0;
+  uintptr_t gwp_asan_metadata = 0;
+};
diff --git a/debuggerd/libdebuggerd/test/tombstone_test.cpp b/debuggerd/libdebuggerd/test/tombstone_test.cpp
index eed95bc..aec8c60 100644
--- a/debuggerd/libdebuggerd/test/tombstone_test.cpp
+++ b/debuggerd/libdebuggerd/test/tombstone_test.cpp
@@ -371,7 +371,7 @@
   GwpAsanCrashDataTest(
       gwp_asan::Error error,
       const gwp_asan::AllocationMetadata *responsible_allocation) :
-      GwpAsanCrashData(nullptr, 0u, 0u, ThreadInfo{}) {
+      GwpAsanCrashData(nullptr, ProcessInfo{}, ThreadInfo{}) {
     is_gwp_asan_responsible_ = true;
     error_ = error;
     responsible_allocation_ = responsible_allocation;
diff --git a/debuggerd/libdebuggerd/tombstone.cpp b/debuggerd/libdebuggerd/tombstone.cpp
index fd52e81..b3f059c 100644
--- a/debuggerd/libdebuggerd/tombstone.cpp
+++ b/debuggerd/libdebuggerd/tombstone.cpp
@@ -376,8 +376,7 @@
 }
 
 static bool dump_thread(log_t* log, unwindstack::Unwinder* unwinder, const ThreadInfo& thread_info,
-                        uint64_t abort_msg_address, bool primary_thread,
-                        const GwpAsanCrashData& gwp_asan_crash_data) {
+                        const ProcessInfo& process_info, bool primary_thread) {
   log->current_tid = thread_info.tid;
   if (!primary_thread) {
     _LOG(log, logtype::THREAD, "--- --- --- --- --- --- --- --- --- --- --- --- --- --- --- ---\n");
@@ -388,15 +387,21 @@
     dump_signal_info(log, thread_info, unwinder->GetProcessMemory().get());
   }
 
-  if (primary_thread && gwp_asan_crash_data.CrashIsMine()) {
-    gwp_asan_crash_data.DumpCause(log);
+  std::unique_ptr<GwpAsanCrashData> gwp_asan_crash_data;
+  if (primary_thread) {
+    gwp_asan_crash_data = std::make_unique<GwpAsanCrashData>(unwinder->GetProcessMemory().get(),
+                                                             process_info, thread_info);
+  }
+
+  if (primary_thread && gwp_asan_crash_data->CrashIsMine()) {
+    gwp_asan_crash_data->DumpCause(log);
   } else if (thread_info.siginfo) {
     dump_probable_cause(log, thread_info.siginfo, unwinder->GetMaps(),
                         thread_info.registers.get());
   }
 
   if (primary_thread) {
-    dump_abort_message(log, unwinder->GetProcessMemory().get(), abort_msg_address);
+    dump_abort_message(log, unwinder->GetProcessMemory().get(), process_info.abort_msg_address);
   }
 
   dump_registers(log, thread_info.registers.get());
@@ -413,12 +418,12 @@
   }
 
   if (primary_thread) {
-    if (gwp_asan_crash_data.HasDeallocationTrace()) {
-      gwp_asan_crash_data.DumpDeallocationTrace(log, unwinder);
+    if (gwp_asan_crash_data->HasDeallocationTrace()) {
+      gwp_asan_crash_data->DumpDeallocationTrace(log, unwinder);
     }
 
-    if (gwp_asan_crash_data.HasAllocationTrace()) {
-      gwp_asan_crash_data.DumpAllocationTrace(log, unwinder);
+    if (gwp_asan_crash_data->HasAllocationTrace()) {
+      gwp_asan_crash_data->DumpAllocationTrace(log, unwinder);
     }
 
     unwindstack::Maps* maps = unwinder->GetMaps();
@@ -601,15 +606,16 @@
     LOG(FATAL) << "Failed to init unwinder object.";
   }
 
-  engrave_tombstone(unique_fd(dup(tombstone_fd)), &unwinder, threads, tid, abort_msg_address,
-                    nullptr, nullptr, 0u, 0u);
+  ProcessInfo process_info;
+  process_info.abort_msg_address = abort_msg_address;
+  engrave_tombstone(unique_fd(dup(tombstone_fd)), &unwinder, threads, tid, process_info, nullptr,
+                    nullptr);
 }
 
 void engrave_tombstone(unique_fd output_fd, unwindstack::Unwinder* unwinder,
                        const std::map<pid_t, ThreadInfo>& threads, pid_t target_thread,
-                       uint64_t abort_msg_address, OpenFilesList* open_files,
-                       std::string* amfd_data, uintptr_t gwp_asan_state_ptr,
-                       uintptr_t gwp_asan_metadata_ptr) {
+                       const ProcessInfo& process_info, OpenFilesList* open_files,
+                       std::string* amfd_data) {
   // don't copy log messages to tombstone unless this is a dev device
   bool want_logs = android::base::GetBoolProperty("ro.debuggable", false);
 
@@ -628,12 +634,7 @@
     LOG(FATAL) << "failed to find target thread";
   }
 
-  GwpAsanCrashData gwp_asan_crash_data(unwinder->GetProcessMemory().get(),
-                                       gwp_asan_state_ptr,
-                                       gwp_asan_metadata_ptr, it->second);
-
-  dump_thread(&log, unwinder, it->second, abort_msg_address, true,
-              gwp_asan_crash_data);
+  dump_thread(&log, unwinder, it->second, process_info, true);
 
   if (want_logs) {
     dump_logs(&log, it->second.pid, 50);
@@ -644,7 +645,7 @@
       continue;
     }
 
-    dump_thread(&log, unwinder, thread_info, 0, false, gwp_asan_crash_data);
+    dump_thread(&log, unwinder, thread_info, process_info, false);
   }
 
   if (open_files) {
diff --git a/fs_mgr/liblp/images.cpp b/fs_mgr/liblp/images.cpp
index e4d92ca..0b1e522 100644
--- a/fs_mgr/liblp/images.cpp
+++ b/fs_mgr/liblp/images.cpp
@@ -124,7 +124,7 @@
 }
 
 bool WriteToImageFile(const std::string& file, const LpMetadata& input) {
-    unique_fd fd(open(file.c_str(), O_CREAT | O_RDWR | O_TRUNC | O_CLOEXEC, 0644));
+    unique_fd fd(open(file.c_str(), O_CREAT | O_RDWR | O_TRUNC | O_CLOEXEC | O_BINARY, 0644));
     if (fd < 0) {
         PERROR << __PRETTY_FUNCTION__ << " open failed: " << file;
         return false;
@@ -184,7 +184,7 @@
 }
 
 bool ImageBuilder::Export(const std::string& file) {
-    unique_fd fd(open(file.c_str(), O_CREAT | O_RDWR | O_TRUNC | O_CLOEXEC, 0644));
+    unique_fd fd(open(file.c_str(), O_CREAT | O_RDWR | O_TRUNC | O_CLOEXEC | O_BINARY, 0644));
     if (fd < 0) {
         PERROR << "open failed: " << file;
         return false;
@@ -208,7 +208,7 @@
         std::string file_name = "super_" + name + ".img";
         std::string file_path = output_dir + "/" + file_name;
 
-        static const int kOpenFlags = O_CREAT | O_RDWR | O_TRUNC | O_CLOEXEC | O_NOFOLLOW;
+        static const int kOpenFlags = O_CREAT | O_RDWR | O_TRUNC | O_CLOEXEC | O_NOFOLLOW | O_BINARY;
         unique_fd fd(open(file_path.c_str(), kOpenFlags, 0644));
         if (fd < 0) {
             PERROR << "open failed: " << file_path;
@@ -443,7 +443,7 @@
 }
 
 int ImageBuilder::OpenImageFile(const std::string& file) {
-    android::base::unique_fd source_fd = GetControlFileOrOpen(file.c_str(), O_RDONLY | O_CLOEXEC);
+    unique_fd source_fd = GetControlFileOrOpen(file.c_str(), O_RDONLY | O_CLOEXEC | O_BINARY);
     if (source_fd < 0) {
         PERROR << "open image file failed: " << file;
         return -1;
diff --git a/fs_mgr/libsnapshot/Android.bp b/fs_mgr/libsnapshot/Android.bp
index 6c739fc..0a0a21d 100644
--- a/fs_mgr/libsnapshot/Android.bp
+++ b/fs_mgr/libsnapshot/Android.bp
@@ -235,7 +235,4 @@
         // TODO(b/148818798): remove when parent bug is fixed.
         "libutilscallstack",
     ],
-    init_rc: [
-        "snapshotctl.rc",
-    ],
 }
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/return.h b/fs_mgr/libsnapshot/include/libsnapshot/return.h
index 1f132fa..dedc445 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/return.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/return.h
@@ -30,7 +30,6 @@
     enum class ErrorCode : int32_t {
         SUCCESS = static_cast<int32_t>(FiemapStatus::ErrorCode::SUCCESS),
         ERROR = static_cast<int32_t>(FiemapStatus::ErrorCode::ERROR),
-        NEEDS_REBOOT = ERROR + 1,
         NO_SPACE = static_cast<int32_t>(FiemapStatus::ErrorCode::NO_SPACE),
     };
     ErrorCode error_code() const { return error_code_; }
@@ -43,7 +42,6 @@
     static Return Ok() { return Return(ErrorCode::SUCCESS); }
     static Return Error() { return Return(ErrorCode::ERROR); }
     static Return NoSpace(uint64_t size) { return Return(ErrorCode::NO_SPACE, size); }
-    static Return NeedsReboot() { return Return(ErrorCode::NEEDS_REBOOT); }
     // Does not set required_size_ properly even when status.error_code() == NO_SPACE.
     explicit Return(const FiemapStatus& status)
         : error_code_(FromFiemapStatusErrorCode(status.error_code())), required_size_(0) {}
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
index b440c71..32345d2 100644
--- a/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot.h
@@ -140,7 +140,6 @@
     // Before calling this function, all snapshots must be mapped.
     bool FinishedSnapshotWrites();
 
-  private:
     // Initiate a merge on all snapshot devices. This should only be used after an
     // update has been marked successful after booting.
     bool InitiateMerge();
@@ -149,7 +148,11 @@
     // /data is mounted.
     //
     // If a merge is in progress, this function will block until the merge is
-    // completed. If a merge or update was cancelled, this will clean up any
+    // completed.
+    //    - Callback is called periodically during the merge. If callback()
+    //      returns false during the merge, ProcessUpdateState() will pause
+    //      and returns Merging.
+    // If a merge or update was cancelled, this will clean up any
     // update artifacts and return.
     //
     // Note that after calling this, GetUpdateState() may still return that a
@@ -169,9 +172,9 @@
     //
     // The optional callback allows the caller to periodically check the
     // progress with GetUpdateState().
-    UpdateState ProcessUpdateState(const std::function<void()>& callback = {});
+    UpdateState ProcessUpdateState(const std::function<bool()>& callback = {},
+                                   const std::function<bool()>& before_cancel = {});
 
-  public:
     // Initiate the merge if necessary, then wait for the merge to finish.
     // See InitiateMerge() and ProcessUpdateState() for details.
     // Returns:
@@ -179,16 +182,8 @@
     //   - Unverified if called on the source slot
     //   - MergeCompleted if merge is completed
     //   - other states indicating an error has occurred
-    UpdateState InitiateMergeAndWait(SnapshotMergeReport* report = nullptr);
-
-    // Wait for the merge if rebooted into the new slot. Does NOT initiate a
-    // merge. If the merge has not been initiated (but should be), wait.
-    // Returns:
-    //   - Return::Ok(): there is no merge or merge finishes
-    //   - Return::NeedsReboot(): merge finishes but need a reboot before
-    //     applying the next update.
-    //   - Return::Error(): other irrecoverable errors
-    Return WaitForMerge();
+    UpdateState InitiateMergeAndWait(SnapshotMergeReport* report = nullptr,
+                                     const std::function<bool()>& before_cancel = {});
 
     // Find the status of the current update, if any.
     //
@@ -375,14 +370,23 @@
 
     // Check for a cancelled or rolled back merge, returning true if such a
     // condition was detected and handled.
-    bool HandleCancelledUpdate(LockedFile* lock);
+    bool HandleCancelledUpdate(LockedFile* lock, const std::function<bool()>& before_cancel);
 
     // Helper for HandleCancelledUpdate. Assumes booting from new slot.
     bool AreAllSnapshotsCancelled(LockedFile* lock);
 
+    // Determine whether partition names in |snapshots| have been flashed and
+    // store result to |out|.
+    // Return true if values are successfully retrieved and false on error
+    // (e.g. super partition metadata cannot be read). When it returns true,
+    // |out| stores true for partitions that have been flashed and false for
+    // partitions that have not been flashed.
+    bool GetSnapshotFlashingStatus(LockedFile* lock, const std::vector<std::string>& snapshots,
+                                   std::map<std::string, bool>* out);
+
     // Remove artifacts created by the update process, such as snapshots, and
     // set the update state to None.
-    bool RemoveAllUpdateState(LockedFile* lock);
+    bool RemoveAllUpdateState(LockedFile* lock, const std::function<bool()>& prolog = {});
 
     // Interact with /metadata/ota.
     std::unique_ptr<LockedFile> OpenLock(int lock_flags);
@@ -437,8 +441,8 @@
     //   UpdateState::MergeCompleted
     //   UpdateState::MergeFailed
     //   UpdateState::MergeNeedsReboot
-    UpdateState CheckMergeState();
-    UpdateState CheckMergeState(LockedFile* lock);
+    UpdateState CheckMergeState(const std::function<bool()>& before_cancel);
+    UpdateState CheckMergeState(LockedFile* lock, const std::function<bool()>& before_cancel);
     UpdateState CheckTargetMergeState(LockedFile* lock, const std::string& name);
 
     // Interact with status files under /metadata/ota/snapshots.
@@ -513,6 +517,11 @@
 
     std::string ReadUpdateSourceSlotSuffix();
 
+    // Helper for RemoveAllSnapshots.
+    // Check whether |name| should be deleted as a snapshot name.
+    bool ShouldDeleteSnapshot(LockedFile* lock, const std::map<std::string, bool>& flashing_status,
+                              Slot current_slot, const std::string& name);
+
     std::string gsid_dir_;
     std::string metadata_dir_;
     std::unique_ptr<IDeviceInfo> device_;
diff --git a/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h
new file mode 100644
index 0000000..91dd34f
--- /dev/null
+++ b/fs_mgr/libsnapshot/include/libsnapshot/snapshot_stats.h
@@ -0,0 +1,59 @@
+// Copyright (C) 2020 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <chrono>
+#include <memory>
+
+#include <android/snapshot/snapshot.pb.h>
+#include <libsnapshot/snapshot.h>
+
+namespace android {
+namespace snapshot {
+
+class SnapshotMergeStats {
+  public:
+    // Not thread safe.
+    static SnapshotMergeStats* GetInstance(SnapshotManager& manager);
+
+    // Called when merge starts or resumes.
+    bool Start();
+    void set_state(android::snapshot::UpdateState state);
+
+    // Called when merge ends. Properly clean up permanent storage.
+    class Result {
+      public:
+        virtual ~Result() {}
+        virtual const SnapshotMergeReport& report() const = 0;
+        // Time between successful Start() / Resume() to Finish().
+        virtual std::chrono::steady_clock::duration merge_time() const = 0;
+    };
+    std::unique_ptr<Result> Finish();
+
+  private:
+    bool ReadState();
+    bool WriteState();
+    bool DeleteState();
+    SnapshotMergeStats(const std::string& path);
+
+    std::string path_;
+    SnapshotMergeReport report_;
+    // Time of the last successful Start() / Resume() call.
+    std::chrono::time_point<std::chrono::steady_clock> start_time_;
+    bool running_{false};
+};
+
+}  // namespace snapshot
+}  // namespace android
diff --git a/fs_mgr/libsnapshot/return.cpp b/fs_mgr/libsnapshot/return.cpp
index 6559c12..cc64af5 100644
--- a/fs_mgr/libsnapshot/return.cpp
+++ b/fs_mgr/libsnapshot/return.cpp
@@ -24,8 +24,6 @@
     switch (error_code()) {
         case ErrorCode::ERROR:
             return "Error";
-        case ErrorCode::NEEDS_REBOOT:
-            return "Retry after reboot";
         case ErrorCode::SUCCESS:
             [[fallthrough]];
         case ErrorCode::NO_SPACE:
diff --git a/fs_mgr/libsnapshot/snapshot.cpp b/fs_mgr/libsnapshot/snapshot.cpp
index 187f24c..7f9c7f1 100644
--- a/fs_mgr/libsnapshot/snapshot.cpp
+++ b/fs_mgr/libsnapshot/snapshot.cpp
@@ -43,10 +43,10 @@
 #endif
 
 #include <android/snapshot/snapshot.pb.h>
+#include <libsnapshot/snapshot_stats.h>
 #include "device_info.h"
 #include "partition_cow_creator.h"
 #include "snapshot_metadata_updater.h"
-#include "snapshot_stats.h"
 #include "utility.h"
 
 namespace android {
@@ -219,7 +219,12 @@
     return true;
 }
 
-bool SnapshotManager::RemoveAllUpdateState(LockedFile* lock) {
+bool SnapshotManager::RemoveAllUpdateState(LockedFile* lock, const std::function<bool()>& prolog) {
+    if (prolog && !prolog()) {
+        LOG(WARNING) << "Can't RemoveAllUpdateState: prolog failed.";
+        return false;
+    }
+
     LOG(INFO) << "Removing all update state.";
 
 #ifdef LIBSNAPSHOT_USE_CALLSTACK
@@ -789,9 +794,10 @@
 // Note that when a merge fails, we will *always* try again to complete the
 // merge each time the device boots. There is no harm in doing so, and if
 // the problem was transient, we might manage to get a new outcome.
-UpdateState SnapshotManager::ProcessUpdateState(const std::function<void()>& callback) {
+UpdateState SnapshotManager::ProcessUpdateState(const std::function<bool()>& callback,
+                                                const std::function<bool()>& before_cancel) {
     while (true) {
-        UpdateState state = CheckMergeState();
+        UpdateState state = CheckMergeState(before_cancel);
         if (state == UpdateState::MergeFailed) {
             AcknowledgeMergeFailure();
         }
@@ -801,8 +807,8 @@
             return state;
         }
 
-        if (callback) {
-            callback();
+        if (callback && !callback()) {
+            return state;
         }
 
         // This wait is not super time sensitive, so we have a relatively
@@ -811,24 +817,27 @@
     }
 }
 
-UpdateState SnapshotManager::CheckMergeState() {
+UpdateState SnapshotManager::CheckMergeState(const std::function<bool()>& before_cancel) {
     auto lock = LockExclusive();
     if (!lock) {
         return UpdateState::MergeFailed;
     }
 
-    UpdateState state = CheckMergeState(lock.get());
+    UpdateState state = CheckMergeState(lock.get(), before_cancel);
     if (state == UpdateState::MergeCompleted) {
         // Do this inside the same lock. Failures get acknowledged without the
         // lock, because flock() might have failed.
         AcknowledgeMergeSuccess(lock.get());
     } else if (state == UpdateState::Cancelled) {
-        RemoveAllUpdateState(lock.get());
+        if (!RemoveAllUpdateState(lock.get(), before_cancel)) {
+            return ReadSnapshotUpdateStatus(lock.get()).state();
+        }
     }
     return state;
 }
 
-UpdateState SnapshotManager::CheckMergeState(LockedFile* lock) {
+UpdateState SnapshotManager::CheckMergeState(LockedFile* lock,
+                                             const std::function<bool()>& before_cancel) {
     UpdateState state = ReadUpdateState(lock);
     switch (state) {
         case UpdateState::None:
@@ -849,7 +858,7 @@
             // This is an edge case. Normally cancelled updates are detected
             // via the merge poll below, but if we never started a merge, we
             // need to also check here.
-            if (HandleCancelledUpdate(lock)) {
+            if (HandleCancelledUpdate(lock, before_cancel)) {
                 return UpdateState::Cancelled;
             }
             return state;
@@ -1169,7 +1178,8 @@
     return true;
 }
 
-bool SnapshotManager::HandleCancelledUpdate(LockedFile* lock) {
+bool SnapshotManager::HandleCancelledUpdate(LockedFile* lock,
+                                            const std::function<bool()>& before_cancel) {
     auto slot = GetCurrentSlot();
     if (slot == Slot::Unknown) {
         return false;
@@ -1177,15 +1187,30 @@
 
     // If all snapshots were reflashed, then cancel the entire update.
     if (AreAllSnapshotsCancelled(lock)) {
-        RemoveAllUpdateState(lock);
-        return true;
+        LOG(WARNING) << "Detected re-flashing, cancelling unverified update.";
+        return RemoveAllUpdateState(lock, before_cancel);
     }
 
-    // This unverified update might be rolled back, or it might not (b/147347110
-    // comment #77). Take no action, as update_engine is responsible for deciding
-    // whether to cancel.
-    LOG(ERROR) << "Update state is being processed before reboot, taking no action.";
-    return false;
+    // If update has been rolled back, then cancel the entire update.
+    // Client (update_engine) is responsible for doing additional cleanup work on its own states
+    // when ProcessUpdateState() returns UpdateState::Cancelled.
+    auto current_slot = GetCurrentSlot();
+    if (current_slot != Slot::Source) {
+        LOG(INFO) << "Update state is being processed while booting at " << current_slot
+                  << " slot, taking no action.";
+        return false;
+    }
+
+    // current_slot == Source. Attempt to detect rollbacks.
+    if (access(GetRollbackIndicatorPath().c_str(), F_OK) != 0) {
+        // This unverified update is not attempted. Take no action.
+        PLOG(INFO) << "Rollback indicator not detected. "
+                   << "Update state is being processed before reboot, taking no action.";
+        return false;
+    }
+
+    LOG(WARNING) << "Detected rollback, cancelling unverified update.";
+    return RemoveAllUpdateState(lock, before_cancel);
 }
 
 std::unique_ptr<LpMetadata> SnapshotManager::ReadCurrentMetadata() {
@@ -1219,6 +1244,28 @@
         return true;
     }
 
+    std::map<std::string, bool> flashing_status;
+
+    if (!GetSnapshotFlashingStatus(lock, snapshots, &flashing_status)) {
+        LOG(WARNING) << "Failed to determine whether partitions have been flashed. Not"
+                     << "removing update states.";
+        return false;
+    }
+
+    bool all_snapshots_cancelled = std::all_of(flashing_status.begin(), flashing_status.end(),
+                                               [](const auto& pair) { return pair.second; });
+
+    if (all_snapshots_cancelled) {
+        LOG(WARNING) << "All partitions are re-flashed after update, removing all update states.";
+    }
+    return all_snapshots_cancelled;
+}
+
+bool SnapshotManager::GetSnapshotFlashingStatus(LockedFile* lock,
+                                                const std::vector<std::string>& snapshots,
+                                                std::map<std::string, bool>* out) {
+    CHECK(lock);
+
     auto source_slot_suffix = ReadUpdateSourceSlotSuffix();
     if (source_slot_suffix.empty()) {
         return false;
@@ -1244,20 +1291,17 @@
         return false;
     }
 
-    bool all_snapshots_cancelled = true;
     for (const auto& snapshot_name : snapshots) {
         if (GetMetadataPartitionState(*metadata, snapshot_name) ==
             MetadataPartitionState::Updated) {
-            all_snapshots_cancelled = false;
-            continue;
+            out->emplace(snapshot_name, false);
+        } else {
+            // Delete snapshots for partitions that are re-flashed after the update.
+            LOG(WARNING) << "Detected re-flashing of partition " << snapshot_name << ".";
+            out->emplace(snapshot_name, true);
         }
-        // Delete snapshots for partitions that are re-flashed after the update.
-        LOG(WARNING) << "Detected re-flashing of partition " << snapshot_name << ".";
     }
-    if (all_snapshots_cancelled) {
-        LOG(WARNING) << "All partitions are re-flashed after update, removing all update states.";
-    }
-    return all_snapshots_cancelled;
+    return true;
 }
 
 bool SnapshotManager::RemoveAllSnapshots(LockedFile* lock) {
@@ -1267,10 +1311,38 @@
         return false;
     }
 
+    std::map<std::string, bool> flashing_status;
+    if (!GetSnapshotFlashingStatus(lock, snapshots, &flashing_status)) {
+        LOG(WARNING) << "Failed to get flashing status";
+    }
+
+    auto current_slot = GetCurrentSlot();
     bool ok = true;
     bool has_mapped_cow_images = false;
     for (const auto& name : snapshots) {
-        if (!UnmapPartitionWithSnapshot(lock, name) || !DeleteSnapshot(lock, name)) {
+        // If booting off source slot, it is okay to unmap and delete all the snapshots.
+        // If boot indicator is missing, update state is None or Initiated, so
+        //   it is also okay to unmap and delete all the snapshots.
+        // If booting off target slot,
+        //  - should not unmap because:
+        //    - In Android mode, snapshots are not mapped, but
+        //      filesystems are mounting off dm-linear targets directly.
+        //    - In recovery mode, assume nothing is mapped, so it is optional to unmap.
+        //  - If partition is flashed or unknown, it is okay to delete snapshots.
+        //    Otherwise (UPDATED flag), only delete snapshots if they are not mapped
+        //    as dm-snapshot (for example, after merge completes).
+        bool should_unmap = current_slot != Slot::Target;
+        bool should_delete = ShouldDeleteSnapshot(lock, flashing_status, current_slot, name);
+
+        bool partition_ok = true;
+        if (should_unmap && !UnmapPartitionWithSnapshot(lock, name)) {
+            partition_ok = false;
+        }
+        if (partition_ok && should_delete && !DeleteSnapshot(lock, name)) {
+            partition_ok = false;
+        }
+
+        if (!partition_ok) {
             // Remember whether or not we were able to unmap the cow image.
             auto cow_image_device = GetCowImageDeviceName(name);
             has_mapped_cow_images |=
@@ -1293,6 +1365,34 @@
     return ok;
 }
 
+// See comments in RemoveAllSnapshots().
+bool SnapshotManager::ShouldDeleteSnapshot(LockedFile* lock,
+                                           const std::map<std::string, bool>& flashing_status,
+                                           Slot current_slot, const std::string& name) {
+    if (current_slot != Slot::Target) {
+        return true;
+    }
+    auto it = flashing_status.find(name);
+    if (it == flashing_status.end()) {
+        LOG(WARNING) << "Can't determine flashing status for " << name;
+        return true;
+    }
+    if (it->second) {
+        // partition flashed, okay to delete obsolete snapshots
+        return true;
+    }
+    // partition updated, only delete if not dm-snapshot
+    SnapshotStatus status;
+    if (!ReadSnapshotStatus(lock, name, &status)) {
+        LOG(WARNING) << "Unable to read snapshot status for " << name
+                     << ", guessing snapshot device name";
+        auto extra_name = GetSnapshotExtraDeviceName(name);
+        return !IsSnapshotDevice(name) && !IsSnapshotDevice(extra_name);
+    }
+    auto dm_name = GetSnapshotDeviceName(name, status);
+    return !IsSnapshotDevice(dm_name);
+}
+
 UpdateState SnapshotManager::GetUpdateState(double* progress) {
     // If we've never started an update, the state file won't exist.
     auto state_file = GetStateFilePath();
@@ -1379,12 +1479,14 @@
     auto slot = GetCurrentSlot();
 
     if (slot != Slot::Target) {
-        if (slot == Slot::Source && !device_->IsRecovery()) {
+        if (slot == Slot::Source) {
             // Device is rebooting into the original slot, so mark this as a
             // rollback.
             auto path = GetRollbackIndicatorPath();
             if (!android::base::WriteStringToFile("1", path)) {
                 PLOG(ERROR) << "Unable to write rollback indicator: " << path;
+            } else {
+                LOG(INFO) << "Rollback detected, writing rollback indicator to " << path;
             }
         }
         LOG(INFO) << "Not booting from new slot. Will not mount snapshots.";
@@ -2391,7 +2493,8 @@
     return AutoUnmountDevice::New(device_->GetMetadataDir());
 }
 
-UpdateState SnapshotManager::InitiateMergeAndWait(SnapshotMergeReport* stats_report) {
+UpdateState SnapshotManager::InitiateMergeAndWait(SnapshotMergeReport* stats_report,
+                                                  const std::function<bool()>& before_cancel) {
     {
         auto lock = LockExclusive();
         // Sync update state from file with bootloader.
@@ -2401,23 +2504,24 @@
         }
     }
 
-    SnapshotMergeStats merge_stats(*this);
+    auto merge_stats = SnapshotMergeStats::GetInstance(*this);
 
     unsigned int last_progress = 0;
-    auto callback = [&]() -> void {
+    auto callback = [&]() -> bool {
         double progress;
         GetUpdateState(&progress);
         if (last_progress < static_cast<unsigned int>(progress)) {
             last_progress = progress;
             LOG(INFO) << "Waiting for merge to complete: " << last_progress << "%.";
         }
+        return true;  // continue
     };
 
     LOG(INFO) << "Waiting for any previous merge request to complete. "
               << "This can take up to several minutes.";
-    merge_stats.Resume();
-    auto state = ProcessUpdateState(callback);
-    merge_stats.set_state(state);
+    merge_stats->Start();
+    auto state = ProcessUpdateState(callback, before_cancel);
+    merge_stats->set_state(state);
     if (state == UpdateState::None) {
         LOG(INFO) << "Can't find any snapshot to merge.";
         return state;
@@ -2428,10 +2532,6 @@
             return state;
         }
 
-        // This is the first snapshot merge that is requested after OTA. We can
-        // initialize the merge duration statistics.
-        merge_stats.Start();
-
         if (!InitiateMerge()) {
             LOG(ERROR) << "Failed to initiate merge.";
             return state;
@@ -2439,43 +2539,22 @@
         // All other states can be handled by ProcessUpdateState.
         LOG(INFO) << "Waiting for merge to complete. This can take up to several minutes.";
         last_progress = 0;
-        state = ProcessUpdateState(callback);
-        merge_stats.set_state(state);
+        state = ProcessUpdateState(callback, before_cancel);
+        merge_stats->set_state(state);
     }
 
     LOG(INFO) << "Merge finished with state \"" << state << "\".";
     if (stats_report) {
-        *stats_report = merge_stats.GetReport();
+        auto result = merge_stats->Finish();
+        if (result) {
+            *stats_report = result->report();
+        } else {
+            LOG(WARNING) << "SnapshotMergeStatus::Finish failed.";
+        }
     }
     return state;
 }
 
-Return SnapshotManager::WaitForMerge() {
-    LOG(INFO) << "Waiting for any previous merge request to complete. "
-              << "This can take up to several minutes.";
-    while (true) {
-        auto state = ProcessUpdateState();
-        if (state == UpdateState::Unverified && GetCurrentSlot() == Slot::Target) {
-            LOG(INFO) << "Wait for merge to be initiated.";
-            std::this_thread::sleep_for(kUpdateStateCheckInterval);
-            continue;
-        }
-        LOG(INFO) << "Wait for merge exits with state " << state;
-        switch (state) {
-            case UpdateState::None:
-                [[fallthrough]];
-            case UpdateState::MergeCompleted:
-                [[fallthrough]];
-            case UpdateState::Cancelled:
-                return Return::Ok();
-            case UpdateState::MergeNeedsReboot:
-                return Return::NeedsReboot();
-            default:
-                return Return::Error();
-        }
-    }
-}
-
 bool SnapshotManager::HandleImminentDataWipe(const std::function<void()>& callback) {
     if (!device_->IsRecovery()) {
         LOG(ERROR) << "Data wipes are only allowed in recovery.";
@@ -2506,7 +2585,10 @@
         return false;
     }
 
-    UpdateState state = ProcessUpdateState(callback);
+    UpdateState state = ProcessUpdateState([&]() -> bool {
+        callback();
+        return true;
+    });
     LOG(INFO) << "Update state in recovery: " << state;
     switch (state) {
         case UpdateState::MergeFailed:
diff --git a/fs_mgr/libsnapshot/snapshot_stats.cpp b/fs_mgr/libsnapshot/snapshot_stats.cpp
index 635b47d..5da7b98 100644
--- a/fs_mgr/libsnapshot/snapshot_stats.cpp
+++ b/fs_mgr/libsnapshot/snapshot_stats.cpp
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "snapshot_stats.h"
+#include <libsnapshot/snapshot_stats.h>
 
 #include <sstream>
 
@@ -23,22 +23,17 @@
 namespace android {
 namespace snapshot {
 
-SnapshotMergeStats::SnapshotMergeStats(SnapshotManager& parent) : parent_(parent) {
-    init_time_ = std::chrono::steady_clock::now();
+SnapshotMergeStats* SnapshotMergeStats::GetInstance(SnapshotManager& parent) {
+    static SnapshotMergeStats g_instance(parent.GetMergeStateFilePath());
+    CHECK(g_instance.path_ == parent.GetMergeStateFilePath());
+    return &g_instance;
 }
 
-SnapshotMergeStats::~SnapshotMergeStats() {
-    std::string error;
-    auto file_path = parent_.GetMergeStateFilePath();
-    if (!android::base::RemoveFileIfExists(file_path, &error)) {
-        LOG(ERROR) << "Failed to remove merge statistics file " << file_path << ": " << error;
-        return;
-    }
-}
+SnapshotMergeStats::SnapshotMergeStats(const std::string& path) : path_(path), running_(false) {}
 
 bool SnapshotMergeStats::ReadState() {
     std::string contents;
-    if (!android::base::ReadFileToString(parent_.GetMergeStateFilePath(), &contents)) {
+    if (!android::base::ReadFileToString(path_, &contents)) {
         PLOG(INFO) << "Read merge statistics file failed";
         return false;
     }
@@ -55,34 +50,73 @@
         LOG(ERROR) << "Unable to serialize SnapshotMergeStats.";
         return false;
     }
-    auto file_path = parent_.GetMergeStateFilePath();
-    if (!WriteStringToFileAtomic(contents, file_path)) {
+    if (!WriteStringToFileAtomic(contents, path_)) {
         PLOG(ERROR) << "Could not write to merge statistics file";
         return false;
     }
     return true;
 }
 
-void SnapshotMergeStats::Start() {
-    report_.set_resume_count(0);
-    report_.set_state(UpdateState::None);
-    WriteState();
+bool SnapshotMergeStats::DeleteState() {
+    std::string error;
+    if (!android::base::RemoveFileIfExists(path_, &error)) {
+        LOG(ERROR) << "Failed to remove merge statistics file " << path_ << ": " << error;
+        return false;
+    }
+    return true;
 }
 
-void SnapshotMergeStats::Resume() {
-    if (!ReadState()) {
-        return;
+bool SnapshotMergeStats::Start() {
+    if (running_) {
+        LOG(ERROR) << "SnapshotMergeStats running_ == " << running_;
+        return false;
     }
-    report_.set_resume_count(report_.resume_count() + 1);
-    WriteState();
+    running_ = true;
+
+    start_time_ = std::chrono::steady_clock::now();
+    if (ReadState()) {
+        report_.set_resume_count(report_.resume_count() + 1);
+    } else {
+        report_.set_resume_count(0);
+        report_.set_state(UpdateState::None);
+    }
+
+    return WriteState();
 }
 
 void SnapshotMergeStats::set_state(android::snapshot::UpdateState state) {
     report_.set_state(state);
 }
 
-SnapshotMergeReport SnapshotMergeStats::GetReport() {
-    return report_;
+class SnapshotMergeStatsResultImpl : public SnapshotMergeStats::Result {
+  public:
+    SnapshotMergeStatsResultImpl(const SnapshotMergeReport& report,
+                                 std::chrono::steady_clock::duration merge_time)
+        : report_(report), merge_time_(merge_time) {}
+    const SnapshotMergeReport& report() const override { return report_; }
+    std::chrono::steady_clock::duration merge_time() const override { return merge_time_; }
+
+  private:
+    SnapshotMergeReport report_;
+    std::chrono::steady_clock::duration merge_time_;
+};
+
+std::unique_ptr<SnapshotMergeStats::Result> SnapshotMergeStats::Finish() {
+    if (!running_) {
+        LOG(ERROR) << "SnapshotMergeStats running_ == " << running_;
+        return nullptr;
+    }
+    running_ = false;
+
+    auto result = std::make_unique<SnapshotMergeStatsResultImpl>(
+            report_, std::chrono::steady_clock::now() - start_time_);
+
+    // We still want to report result if state is not deleted. Just leave
+    // it there and move on. A side effect is that it may be reported over and
+    // over again in the future, but there is nothing we can do.
+    (void)DeleteState();
+
+    return result;
 }
 
 }  // namespace snapshot
diff --git a/fs_mgr/libsnapshot/snapshot_stats.h b/fs_mgr/libsnapshot/snapshot_stats.h
deleted file mode 100644
index 60109a4..0000000
--- a/fs_mgr/libsnapshot/snapshot_stats.h
+++ /dev/null
@@ -1,45 +0,0 @@
-// Copyright (C) 2020 The Android Open Source Project
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#pragma once
-
-#include <chrono>
-
-#include <android/snapshot/snapshot.pb.h>
-#include <libsnapshot/snapshot.h>
-
-namespace android {
-namespace snapshot {
-
-class SnapshotMergeStats {
-  public:
-    SnapshotMergeStats(SnapshotManager& parent);
-    ~SnapshotMergeStats();
-    void Start();
-    void Resume();
-    void set_state(android::snapshot::UpdateState state);
-    SnapshotMergeReport GetReport();
-
-  private:
-    bool ReadState();
-    bool WriteState();
-
-    const SnapshotManager& parent_;
-    SnapshotMergeReport report_;
-    std::chrono::time_point<std::chrono::steady_clock> init_time_;
-    std::chrono::time_point<std::chrono::steady_clock> end_time_;
-};
-
-}  // namespace snapshot
-}  // namespace android
diff --git a/fs_mgr/libsnapshot/snapshot_test.cpp b/fs_mgr/libsnapshot/snapshot_test.cpp
index 0ffaa71..7d16ec2 100644
--- a/fs_mgr/libsnapshot/snapshot_test.cpp
+++ b/fs_mgr/libsnapshot/snapshot_test.cpp
@@ -1570,45 +1570,6 @@
             << "FinishedSnapshotWrites should detect overflow of CoW device.";
 }
 
-TEST_F(SnapshotUpdateTest, WaitForMerge) {
-    AddOperationForPartitions();
-
-    // Execute the update.
-    ASSERT_TRUE(sm->BeginUpdate());
-    ASSERT_TRUE(sm->CreateUpdateSnapshots(manifest_));
-
-    // Write some data to target partitions.
-    for (const auto& name : {"sys_b", "vnd_b", "prd_b"}) {
-        ASSERT_TRUE(WriteSnapshotAndHash(name));
-    }
-
-    ASSERT_TRUE(sm->FinishedSnapshotWrites());
-
-    // Simulate shutting down the device.
-    ASSERT_TRUE(UnmapAll());
-
-    // After reboot, init does first stage mount.
-    {
-        auto init = SnapshotManager::NewForFirstStageMount(new TestDeviceInfo(fake_super, "_b"));
-        ASSERT_NE(nullptr, init);
-        ASSERT_TRUE(init->CreateLogicalAndSnapshotPartitions("super", snapshot_timeout_));
-    }
-
-    auto new_sm = SnapshotManager::New(new TestDeviceInfo(fake_super, "_b"));
-    ASSERT_NE(nullptr, new_sm);
-
-    auto waiter = std::async(std::launch::async, [&new_sm] { return new_sm->WaitForMerge(); });
-    ASSERT_EQ(std::future_status::timeout, waiter.wait_for(1s))
-            << "WaitForMerge should block when not initiated";
-
-    auto merger =
-            std::async(std::launch::async, [&new_sm] { return new_sm->InitiateMergeAndWait(); });
-    // Small images, so should be merged pretty quickly.
-    ASSERT_EQ(std::future_status::ready, waiter.wait_for(3s)) << "WaitForMerge did not finish";
-    ASSERT_TRUE(waiter.get());
-    ASSERT_THAT(merger.get(), AnyOf(UpdateState::None, UpdateState::MergeCompleted));
-}
-
 TEST_F(SnapshotUpdateTest, LowSpace) {
     static constexpr auto kMaxFree = 10_MiB;
     auto userdata = std::make_unique<LowSpaceUserdata>();
diff --git a/fs_mgr/libsnapshot/snapshotctl.cpp b/fs_mgr/libsnapshot/snapshotctl.cpp
index 34d3d69..aa5e9c1 100644
--- a/fs_mgr/libsnapshot/snapshotctl.cpp
+++ b/fs_mgr/libsnapshot/snapshotctl.cpp
@@ -26,9 +26,9 @@
 #include <android-base/unique_fd.h>
 #include <android/snapshot/snapshot.pb.h>
 #include <libsnapshot/snapshot.h>
+#include <libsnapshot/snapshot_stats.h>
 #include <statslog.h>
 
-#include "snapshot_stats.h"
 #include "utility.h"
 
 using namespace std::string_literals;
@@ -178,6 +178,7 @@
     }
 
     LOG(ERROR) << "Snapshot failed to merge with state \"" << state << "\".";
+
     return false;
 }
 
diff --git a/fs_mgr/libsnapshot/snapshotctl.rc b/fs_mgr/libsnapshot/snapshotctl.rc
deleted file mode 100644
index ccb2c41..0000000
--- a/fs_mgr/libsnapshot/snapshotctl.rc
+++ /dev/null
@@ -1,2 +0,0 @@
-on property:sys.boot_completed=1
-    exec_background - root root -- /system/bin/snapshotctl merge --logcat --log-to-file --report
diff --git a/init/init.cpp b/init/init.cpp
index a530d18..63aefc1 100644
--- a/init/init.cpp
+++ b/init/init.cpp
@@ -778,11 +778,10 @@
     if (false) DumpState();
 
     // Make the GSI status available before scripts start running.
-    if (android::gsi::IsGsiRunning()) {
-        SetProperty("ro.gsid.image_running", "1");
-    } else {
-        SetProperty("ro.gsid.image_running", "0");
-    }
+    auto is_running = android::gsi::IsGsiRunning() ? "1" : "0";
+    SetProperty(gsi::kGsiBootedProp, is_running);
+    auto is_installed = android::gsi::IsGsiInstalled() ? "1" : "0";
+    SetProperty(gsi::kGsiInstalledProp, is_installed);
 
     am.QueueBuiltinAction(SetupCgroupsAction, "SetupCgroups");
     am.QueueBuiltinAction(SetKptrRestrictAction, "SetKptrRestrict");
diff --git a/liblog/include/android/log.h b/liblog/include/android/log.h
index c98455d..b9839d6 100644
--- a/liblog/include/android/log.h
+++ b/liblog/include/android/log.h
@@ -156,7 +156,7 @@
  * Let the logging function choose the best log target.
  * This is not part of the enum since adding either -1 or 0xFFFFFFFF forces the enum to be signed or
  * unsigned, which breaks unfortunately common arithmetic against LOG_ID_MIN and LOG_ID_MAX. */
-#define LOG_ID_DEFAULT -1
+#define LOG_ID_DEFAULT (-1)
 
 /**
  * Writes the constant string `text` to the log buffer `id`,
diff --git a/rootdir/init.rc b/rootdir/init.rc
index 276cc1e..01b7ca3 100644
--- a/rootdir/init.rc
+++ b/rootdir/init.rc
@@ -701,9 +701,7 @@
 
     # A tmpfs directory, which will contain all apps CE DE data directory that
     # bind mount from the original source.
-    chown root root /data_mirror
-    chmod 0700 /data_mirror
-    mount tmpfs tmpfs /data_mirror mode=0700,uid=0,gid=1000 nodev noexec nosuid
+    mount tmpfs tmpfs /data_mirror nodev noexec nosuid mode=0700,uid=0,gid=1000
     restorecon /data_mirror
     mkdir /data_mirror/data_ce 0700 root root
     mkdir /data_mirror/data_de 0700 root root
@@ -732,6 +730,12 @@
     # Create root dir for Incremental Service
     mkdir /data/incremental 0771 system system encryption=Require
 
+    # Create directories for statsd
+    mkdir /data/misc/stats-data/ 0770 statsd system
+    mkdir /data/misc/stats-service/ 0770 statsd system
+    mkdir /data/misc/stats-active-metric/ 0770 statsd system
+    mkdir /data/misc/train-info/ 0770 statsd system
+
     # Wait for apexd to finish activating APEXes before starting more processes.
     wait_for_prop apexd.status activated
     perform_apex_config
diff --git a/run-as/run-as.cpp b/run-as/run-as.cpp
index 432c434..cc92c68 100644
--- a/run-as/run-as.cpp
+++ b/run-as/run-as.cpp
@@ -247,12 +247,12 @@
 
   std::string seinfo = std::string(info.seinfo) + ":fromRunAs";
   if (selinux_android_setcontext(uid, 0, seinfo.c_str(), pkgname) < 0) {
-    error(1, errno, "couldn't set SELinux security context");
+    error(1, errno, "couldn't set SELinux security context '%s'", seinfo.c_str());
   }
 
   // cd into the data directory, and set $HOME correspondingly.
   if (TEMP_FAILURE_RETRY(chdir(info.data_dir)) == -1) {
-    error(1, errno, "couldn't chdir to package's data directory");
+    error(1, errno, "couldn't chdir to package's data directory '%s'", info.data_dir);
   }
   setenv("HOME", info.data_dir, 1);