IUpdateEngine.cleanupSuccessfulUpdate is async

UpdateAttempterAndroid::CleanupSuccessfulUpdate uses
CleanupPreviousUpdateAction, which is asynchronous. Hence,
the binder function needs to be asynchronous too.

Test: update_attempter_android --merge after update
Bug: 147696014

Change-Id: I0ce43db6d75a5a13a105c25645824612bd4d6be3
Merged-In: I0ce43db6d75a5a13a105c25645824612bd4d6be3

8 files changed, 134 insertions, 27 deletions

diff --git a/binder_bindings/android/os/IUpdateEngine.aidl b/binder_bindings/android/os/IUpdateEngine.aidl
index bbb86a3f..c9580da4 100644
--- a/binder_bindings/android/os/IUpdateEngine.aidl
+++ b/binder_bindings/android/os/IUpdateEngine.aidl
@@ -62,8 +62,13 @@ interface IUpdateEngine {
    *
    * Wait for merge to finish, and clean up necessary files.
    *
-   * @return SUCCESS if successful. ERROR if transient errors (e.g. merged but
-   * needs reboot). DEVICE_CORRUPTED for permanent errors.
+   * @param callback Report status updates in callback (not the one previously
+   * bound with {@link #bind()}).
+   * {@link IUpdateEngineCallback#onStatusUpdate} is called with
+   * CLEANUP_PREVIOUS_UPDATE and a progress value during the cleanup.
+   * {@link IUpdateEngineCallback#onPayloadApplicationComplete} is called at
+   * the end with SUCCESS if successful. ERROR if transient errors (e.g. merged
+   * but needs reboot). DEVICE_CORRUPTED for permanent errors.
    */
-  int cleanupSuccessfulUpdate();
+  void cleanupSuccessfulUpdate(IUpdateEngineCallback callback);
 }
diff --git a/binder_service_android.cc b/binder_service_android.cc
index c376f4ea..6b8a5529 100644
--- a/binder_service_android.cc
+++ b/binder_service_android.cc
@@ -16,6 +16,8 @@
 
 #include "update_engine/binder_service_android.h"
 
+#include <memory>
+
 #include <base/bind.h>
 #include <base/logging.h>
 #include <binderwrapper/binder_wrapper.h>
@@ -217,10 +219,37 @@ Status BinderUpdateEngineAndroidService::allocateSpaceForPayload(
   return Status::ok();
 }
 
+class CleanupSuccessfulUpdateCallback
+    : public CleanupSuccessfulUpdateCallbackInterface {
+ public:
+  CleanupSuccessfulUpdateCallback(
+      const android::sp<IUpdateEngineCallback>& callback)
+      : callback_(callback) {}
+  void OnCleanupComplete(int32_t error_code) {
+    ignore_result(callback_->onPayloadApplicationComplete(error_code));
+  }
+  void OnCleanupProgressUpdate(double progress) {
+    ignore_result(callback_->onStatusUpdate(
+        static_cast<int32_t>(
+            update_engine::UpdateStatus::CLEANUP_PREVIOUS_UPDATE),
+        progress));
+  }
+  void RegisterForDeathNotifications(base::Closure unbind) {
+    const android::sp<android::IBinder>& callback_binder =
+        IUpdateEngineCallback::asBinder(callback_);
+    auto binder_wrapper = android::BinderWrapper::Get();
+    binder_wrapper->RegisterForDeathNotifications(callback_binder, unbind);
+  }
+
+ private:
+  android::sp<IUpdateEngineCallback> callback_;
+};
+
 Status BinderUpdateEngineAndroidService::cleanupSuccessfulUpdate(
-    int32_t* return_value) {
+    const android::sp<IUpdateEngineCallback>& callback) {
   brillo::ErrorPtr error;
-  *return_value = service_delegate_->CleanupSuccessfulUpdate(&error);
+  service_delegate_->CleanupSuccessfulUpdate(
+      std::make_unique<CleanupSuccessfulUpdateCallback>(callback), &error);
   if (error != nullptr)
     return ErrorPtrToStatus(error);
   return Status::ok();
diff --git a/binder_service_android.h b/binder_service_android.h
index 1c38d2be..5f282252 100644
--- a/binder_service_android.h
+++ b/binder_service_android.h
@@ -75,7 +75,7 @@ class BinderUpdateEngineAndroidService : public android::os::BnUpdateEngine,
       const std::vector<android::String16>& header_kv_pairs,
       int64_t* return_value) override;
   android::binder::Status cleanupSuccessfulUpdate(
-      int32_t* return_value) override;
+      const android::sp<android::os::IUpdateEngineCallback>& callback) override;
 
  private:
   // Remove the passed |callback| from the list of registered callbacks. Called
diff --git a/common/action.h b/common/action.h
index c93e73cb..fd82c2d5 100644
--- a/common/action.h
+++ b/common/action.h
@@ -229,7 +229,8 @@ class NoOpAction : public AbstractAction {
   void PerformAction() override {
     processor_->ActionComplete(this, ErrorCode::kSuccess);
   }
-  std::string Type() const override { return "NoOpAction"; }
+  static std::string StaticType() { return "NoOpAction"; }
+  std::string Type() const override { return StaticType(); }
 };
 
 };  // namespace chromeos_update_engine
diff --git a/service_delegate_android_interface.h b/service_delegate_android_interface.h
index a12f1e8c..34a97123 100644
--- a/service_delegate_android_interface.h
+++ b/service_delegate_android_interface.h
@@ -19,6 +19,7 @@
 
 #include <inttypes.h>
 
+#include <memory>
 #include <string>
 #include <vector>
 
@@ -26,6 +27,18 @@
 
 namespace chromeos_update_engine {
 
+// See ServiceDelegateAndroidInterface.CleanupSuccessfulUpdate
+// Wraps a IUpdateEngineCallback binder object used specifically for
+// CleanupSuccessfulUpdate.
+class CleanupSuccessfulUpdateCallbackInterface {
+ public:
+  virtual ~CleanupSuccessfulUpdateCallbackInterface() {}
+  virtual void OnCleanupProgressUpdate(double progress) = 0;
+  virtual void OnCleanupComplete(int32_t error_code) = 0;
+  // Call RegisterForDeathNotifications on the internal binder object.
+  virtual void RegisterForDeathNotifications(base::Closure unbind) = 0;
+};
+
 // This class defines the interface exposed by the Android version of the
 // daemon service. This interface only includes the method calls that such
 // daemon exposes. For asynchronous events initiated by a class implementing
@@ -99,9 +112,11 @@ class ServiceDelegateAndroidInterface {
   // Wait for merge to complete, then clean up merge after an update has been
   // successful.
   //
-  // This function returns immediately if no merge is needed, but may block
-  // for a long time (up to several minutes) in the worst case.
-  virtual int32_t CleanupSuccessfulUpdate(brillo::ErrorPtr* error) = 0;
+  // This function returns immediately. Progress updates are provided in
+  // |callback|.
+  virtual void CleanupSuccessfulUpdate(
+      std::unique_ptr<CleanupSuccessfulUpdateCallbackInterface> callback,
+      brillo::ErrorPtr* error) = 0;
 
  protected:
   ServiceDelegateAndroidInterface() = default;
diff --git a/update_attempter_android.cc b/update_attempter_android.cc
index 60c3b34b..eecd2daf 100644
--- a/update_attempter_android.cc
+++ b/update_attempter_android.cc
@@ -31,6 +31,7 @@
 #include <brillo/strings/string_utils.h>
 #include <log/log_safetynet.h>
 
+#include "update_engine/cleanup_previous_update_action.h"
 #include "update_engine/common/constants.h"
 #include "update_engine/common/error_code_utils.h"
 #include "update_engine/common/file_fetcher.h"
@@ -552,6 +553,12 @@ void UpdateAttempterAndroid::ActionCompleted(ActionProcessor* processor,
   // Reset download progress regardless of whether or not the download
   // action succeeded.
   const string type = action->Type();
+  if (type == CleanupPreviousUpdateAction::StaticType() ||
+      (type == NoOpAction::StaticType() &&
+       status_ == UpdateStatus::CLEANUP_PREVIOUS_UPDATE)) {
+    cleanup_previous_update_code_ = code;
+    NotifyCleanupPreviousUpdateCallbacksAndClear();
+  }
   if (type == DownloadAction::StaticType()) {
     download_progress_ = 0;
   }
@@ -950,17 +957,27 @@ uint64_t UpdateAttempterAndroid::AllocateSpaceForPayload(
   return 0;
 }
 
-int32_t UpdateAttempterAndroid::CleanupSuccessfulUpdate(
+void UpdateAttempterAndroid::CleanupSuccessfulUpdate(
+    std::unique_ptr<CleanupSuccessfulUpdateCallbackInterface> callback,
     brillo::ErrorPtr* error) {
-  ErrorCode error_code =
-      boot_control_->GetDynamicPartitionControl()->CleanupSuccessfulUpdate();
-  if (error_code == ErrorCode::kSuccess) {
-    LOG(INFO) << "Previous update is merged and cleaned up successfully.";
-  } else {
-    LOG(ERROR) << "CleanupSuccessfulUpdate failed with "
-               << utils::ErrorCodeToString(error_code);
+  if (cleanup_previous_update_code_.has_value()) {
+    LOG(INFO) << "CleanupSuccessfulUpdate has previously completed with "
+              << utils::ErrorCodeToString(*cleanup_previous_update_code_);
+    if (callback) {
+      callback->OnCleanupComplete(
+          static_cast<int32_t>(*cleanup_previous_update_code_));
+    }
+    return;
+  }
+  if (callback) {
+    auto callback_ptr = callback.get();
+    cleanup_previous_update_callbacks_.emplace_back(std::move(callback));
+    callback_ptr->RegisterForDeathNotifications(
+        base::Bind(&UpdateAttempterAndroid::RemoveCleanupPreviousUpdateCallback,
+                   base::Unretained(this),
+                   base::Unretained(callback_ptr)));
   }
-  return static_cast<int32_t>(error_code);
+  ScheduleCleanupPreviousUpdate();
 }
 
 void UpdateAttempterAndroid::ScheduleCleanupPreviousUpdate() {
@@ -981,6 +998,29 @@ void UpdateAttempterAndroid::ScheduleCleanupPreviousUpdate() {
   processor_->StartProcessing();
 }
 
-void UpdateAttempterAndroid::OnCleanupProgressUpdate(double progress) {}
+void UpdateAttempterAndroid::OnCleanupProgressUpdate(double progress) {
+  for (auto&& callback : cleanup_previous_update_callbacks_) {
+    callback->OnCleanupProgressUpdate(progress);
+  }
+}
+
+void UpdateAttempterAndroid::NotifyCleanupPreviousUpdateCallbacksAndClear() {
+  CHECK(cleanup_previous_update_code_.has_value());
+  for (auto&& callback : cleanup_previous_update_callbacks_) {
+    callback->OnCleanupComplete(
+        static_cast<int32_t>(*cleanup_previous_update_code_));
+  }
+  cleanup_previous_update_callbacks_.clear();
+}
+
+void UpdateAttempterAndroid::RemoveCleanupPreviousUpdateCallback(
+    CleanupSuccessfulUpdateCallbackInterface* callback) {
+  auto end_it =
+      std::remove_if(cleanup_previous_update_callbacks_.begin(),
+                     cleanup_previous_update_callbacks_.end(),
+                     [&](const auto& e) { return e.get() == callback; });
+  cleanup_previous_update_callbacks_.erase(
+      end_it, cleanup_previous_update_callbacks_.end());
+}
 
 }  // namespace chromeos_update_engine
diff --git a/update_attempter_android.h b/update_attempter_android.h
index 106e1171..f8c78de1 100644
--- a/update_attempter_android.h
+++ b/update_attempter_android.h
@@ -82,7 +82,9 @@ class UpdateAttempterAndroid
       const std::string& metadata_filename,
       const std::vector<std::string>& key_value_pair_headers,
       brillo::ErrorPtr* error) override;
-  int32_t CleanupSuccessfulUpdate(brillo::ErrorPtr* error) override;
+  void CleanupSuccessfulUpdate(
+      std::unique_ptr<CleanupSuccessfulUpdateCallbackInterface> callback,
+      brillo::ErrorPtr* error) override;
 
   // ActionProcessorDelegate methods:
   void ProcessingDone(const ActionProcessor* processor,
@@ -184,6 +186,13 @@ class UpdateAttempterAndroid
   // Enqueue and run a CleanupPreviousUpdateAction.
   void ScheduleCleanupPreviousUpdate();
 
+  // Notify and clear |cleanup_previous_update_callbacks_|.
+  void NotifyCleanupPreviousUpdateCallbacksAndClear();
+
+  // Remove |callback| from |cleanup_previous_update_callbacks_|.
+  void RemoveCleanupPreviousUpdateCallback(
+      CleanupSuccessfulUpdateCallbackInterface* callback);
+
   DaemonStateInterface* daemon_state_;
 
   // DaemonStateAndroid pointers.
@@ -221,6 +230,12 @@ class UpdateAttempterAndroid
 
   ::android::base::unique_fd payload_fd_;
 
+  std::vector<std::unique_ptr<CleanupSuccessfulUpdateCallbackInterface>>
+      cleanup_previous_update_callbacks_;
+  // Result of previous CleanupPreviousUpdateAction. Nullopt If
+  // CleanupPreviousUpdateAction has not been executed.
+  std::optional<ErrorCode> cleanup_previous_update_code_{std::nullopt};
+
   DISALLOW_COPY_AND_ASSIGN(UpdateAttempterAndroid);
 };
 
diff --git a/update_engine_client_android.cc b/update_engine_client_android.cc
index 7d9bc3de..1a68cf49 100644
--- a/update_engine_client_android.cc
+++ b/update_engine_client_android.cc
@@ -80,6 +80,7 @@ class UpdateEngineClientAndroid : public brillo::Daemon {
 
   android::sp<android::os::IUpdateEngine> service_;
   android::sp<android::os::BnUpdateEngineCallback> callback_;
+  android::sp<android::os::BnUpdateEngineCallback> cleanup_callback_;
 
   brillo::BinderWatcher binder_watcher_;
 };
@@ -226,13 +227,14 @@ int UpdateEngineClientAndroid::OnInit() {
   }
 
   if (FLAGS_merge) {
-    int32_t ret = 0;
-    Status status = service_->cleanupSuccessfulUpdate(&ret);
-    if (status.isOk()) {
-      LOG(INFO) << "CleanupSuccessfulUpdate exits with "
-                << utils::ErrorCodeToString(static_cast<ErrorCode>(ret));
+    // Register a callback object with the service.
+    cleanup_callback_ = new UECallback(this);
+    Status status = service_->cleanupSuccessfulUpdate(cleanup_callback_);
+    if (!status.isOk()) {
+      LOG(ERROR) << "Failed to call cleanupSuccessfulUpdate.";
+      return ExitWhenIdle(status);
     }
-    return ExitWhenIdle(status);
+    keep_running = true;
   }
 
   if (FLAGS_follow) {