11 files changed, 560 insertions, 167 deletions

diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index de8fafd30..6f3a3a2a7 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -828,3 +828,230 @@ TEST_F(UpdaterTest, last_command_verify) {
   RunBlockImageUpdate(true, entries, image_file_, "t");
   ASSERT_EQ(-1, access(last_command_file_.c_str(), R_OK));
 }
+
+class ResumableUpdaterTest : public testing::TestWithParam<size_t> {
+ protected:
+  void SetUp() override {
+    RegisterBuiltins();
+    RegisterInstallFunctions();
+    RegisterBlockImageFunctions();
+
+    Paths::Get().set_cache_temp_source(temp_saved_source_.path);
+    Paths::Get().set_last_command_file(temp_last_command_.path);
+    Paths::Get().set_stash_directory_base(temp_stash_base_.path);
+
+    index_ = GetParam();
+    image_file_ = image_temp_file_.path;
+    last_command_file_ = temp_last_command_.path;
+  }
+
+  void TearDown() override {
+    // Clean up the last_command_file if any.
+    ASSERT_TRUE(android::base::RemoveFileIfExists(last_command_file_));
+
+    // Clear partition updated marker if any.
+    std::string updated_marker{ temp_stash_base_.path };
+    updated_marker += "/" + get_sha1(image_temp_file_.path) + ".UPDATED";
+    ASSERT_TRUE(android::base::RemoveFileIfExists(updated_marker));
+  }
+
+  TemporaryFile temp_saved_source_;
+  TemporaryDir temp_stash_base_;
+  std::string last_command_file_;
+  std::string image_file_;
+  size_t index_;
+
+ private:
+  TemporaryFile temp_last_command_;
+  TemporaryFile image_temp_file_;
+};
+
+static std::string g_source_image;
+static std::string g_target_image;
+static PackageEntries g_entries;
+
+static std::vector<std::string> GenerateTransferList() {
+  std::string a(4096, 'a');
+  std::string b(4096, 'b');
+  std::string c(4096, 'c');
+  std::string d(4096, 'd');
+  std::string e(4096, 'e');
+  std::string f(4096, 'f');
+  std::string g(4096, 'g');
+  std::string h(4096, 'h');
+  std::string i(4096, 'i');
+  std::string zero(4096, '\0');
+
+  std::string a_hash = get_sha1(a);
+  std::string b_hash = get_sha1(b);
+  std::string c_hash = get_sha1(c);
+  std::string e_hash = get_sha1(e);
+
+  auto loc = [](const std::string& range_text) {
+    std::vector<std::string> pieces = android::base::Split(range_text, "-");
+    size_t left;
+    size_t right;
+    if (pieces.size() == 1) {
+      CHECK(android::base::ParseUint(pieces[0], &left));
+      right = left + 1;
+    } else {
+      CHECK_EQ(2u, pieces.size());
+      CHECK(android::base::ParseUint(pieces[0], &left));
+      CHECK(android::base::ParseUint(pieces[1], &right));
+      right++;
+    }
+    return android::base::StringPrintf("2,%zu,%zu", left, right);
+  };
+
+  // patch 1: "b d c" -> "g"
+  TemporaryFile patch_file_bdc_g;
+  std::string bdc = b + d + c;
+  std::string bdc_hash = get_sha1(bdc);
+  std::string g_hash = get_sha1(g);
+  CHECK_EQ(0, bsdiff::bsdiff(reinterpret_cast<const uint8_t*>(bdc.data()), bdc.size(),
+                             reinterpret_cast<const uint8_t*>(g.data()), g.size(),
+                             patch_file_bdc_g.path, nullptr));
+  std::string patch_bdc_g;
+  CHECK(android::base::ReadFileToString(patch_file_bdc_g.path, &patch_bdc_g));
+
+  // patch 2: "a b c d" -> "d c b"
+  TemporaryFile patch_file_abcd_dcb;
+  std::string abcd = a + b + c + d;
+  std::string abcd_hash = get_sha1(abcd);
+  std::string dcb = d + c + b;
+  std::string dcb_hash = get_sha1(dcb);
+  CHECK_EQ(0, bsdiff::bsdiff(reinterpret_cast<const uint8_t*>(abcd.data()), abcd.size(),
+                             reinterpret_cast<const uint8_t*>(dcb.data()), dcb.size(),
+                             patch_file_abcd_dcb.path, nullptr));
+  std::string patch_abcd_dcb;
+  CHECK(android::base::ReadFileToString(patch_file_abcd_dcb.path, &patch_abcd_dcb));
+
+  std::vector<std::string> transfer_list{
+    "4",
+    "10",  // total blocks written
+    "2",   // maximum stash entries
+    "2",   // maximum number of stashed blocks
+
+    // a b c d e a b c d e
+    "stash " + b_hash + " " + loc("1"),
+    // a b c d e a b c d e    [b(1)]
+    "stash " + c_hash + " " + loc("2"),
+    // a b c d e a b c d e    [b(1)][c(2)]
+    "new " + loc("1-2"),
+    // a i h d e a b c d e    [b(1)][c(2)]
+    "zero " + loc("0"),
+    // 0 i h d e a b c d e    [b(1)][c(2)]
+
+    // bsdiff "b d c" (from stash, 3, stash) to get g(3)
+    android::base::StringPrintf(
+        "bsdiff 0 %zu %s %s %s 3 %s %s %s:%s %s:%s",
+        patch_bdc_g.size(),                  // patch start (0), patch length
+        bdc_hash.c_str(),                    // source hash
+        g_hash.c_str(),                      // target hash
+        loc("3").c_str(),                    // target range
+        loc("3").c_str(), loc("1").c_str(),  // load "d" from block 3, into buffer at offset 1
+        b_hash.c_str(), loc("0").c_str(),    // load "b" from stash, into buffer at offset 0
+        c_hash.c_str(), loc("2").c_str()),   // load "c" from stash, into buffer at offset 2
+
+    // 0 i h g e a b c d e    [b(1)][c(2)]
+    "free " + b_hash,
+    // 0 i h g e a b c d e    [c(2)]
+    "free " + a_hash,
+    // 0 i h g e a b c d e
+    "stash " + a_hash + " " + loc("5"),
+    // 0 i h g e a b c d e    [a(5)]
+    "move " + e_hash + " " + loc("5") + " 1 " + loc("4"),
+    // 0 i h g e e b c d e    [a(5)]
+
+    // bsdiff "a b c d" (from stash, 6-8) to "d c b" (6-8)
+    android::base::StringPrintf(  //
+        "bsdiff %zu %zu %s %s %s 4 %s %s %s:%s",
+        patch_bdc_g.size(),                          // patch start
+        patch_bdc_g.size() + patch_abcd_dcb.size(),  // patch length
+        abcd_hash.c_str(),                           // source hash
+        dcb_hash.c_str(),                            // target hash
+        loc("6-8").c_str(),                          // target range
+        loc("6-8").c_str(),                          // load "b c d" from blocks 6-8
+        loc("1-3").c_str(),                          //   into buffer at offset 1-3
+        a_hash.c_str(),                              // load "a" from stash
+        loc("0").c_str()),                           //   into buffer at offset 0
+
+    // 0 i h g e e d c b e    [a(5)]
+    "new " + loc("4"),
+    // 0 i h g f e d c b e    [a(5)]
+    "move " + a_hash + " " + loc("9") + " 1 - " + a_hash + ":" + loc("0"),
+    // 0 i h g f e d c b a    [a(5)]
+    "free " + a_hash,
+    // 0 i h g f e d c b a
+  };
+
+  std::string new_data = i + h + f;
+  std::string patch_data = patch_bdc_g + patch_abcd_dcb;
+
+  g_entries = {
+    { "new_data", new_data },
+    { "patch_data", patch_data },
+  };
+  g_source_image = a + b + c + d + e + a + b + c + d + e;
+  g_target_image = zero + i + h + g + f + e + d + c + b + a;
+
+  return transfer_list;
+}
+
+static const std::vector<std::string> g_transfer_list = GenerateTransferList();
+
+INSTANTIATE_TEST_CASE_P(InterruptAfterEachCommand, ResumableUpdaterTest,
+                        ::testing::Range(static_cast<size_t>(0),
+                                         g_transfer_list.size() - kTransferListHeaderLines));
+
+TEST_P(ResumableUpdaterTest, InterruptVerifyResume) {
+  ASSERT_TRUE(android::base::WriteStringToFile(g_source_image, image_file_));
+
+  LOG(INFO) << "Interrupting at line " << index_ << " ("
+            << g_transfer_list[kTransferListHeaderLines + index_] << ")";
+
+  std::vector<std::string> transfer_list_copy{ g_transfer_list };
+  transfer_list_copy[kTransferListHeaderLines + index_] = "fail";
+
+  g_entries["transfer_list"] = android::base::Join(transfer_list_copy, '\n');
+
+  // Run update that's expected to fail.
+  RunBlockImageUpdate(false, g_entries, image_file_, "");
+
+  std::string last_command_expected;
+
+  // Assert the last_command_file.
+  if (index_ == 0) {
+    ASSERT_EQ(-1, access(last_command_file_.c_str(), R_OK));
+  } else {
+    last_command_expected =
+        std::to_string(index_ - 1) + "\n" + g_transfer_list[kTransferListHeaderLines + index_ - 1];
+    std::string last_command_actual;
+    ASSERT_TRUE(android::base::ReadFileToString(last_command_file_, &last_command_actual));
+    ASSERT_EQ(last_command_expected, last_command_actual);
+  }
+
+  g_entries["transfer_list"] = android::base::Join(g_transfer_list, '\n');
+
+  // Resume the interrupted update, by doing verification first.
+  RunBlockImageUpdate(true, g_entries, image_file_, "t");
+
+  // last_command_file should remain intact.
+  if (index_ == 0) {
+    ASSERT_EQ(-1, access(last_command_file_.c_str(), R_OK));
+  } else {
+    std::string last_command_actual;
+    ASSERT_TRUE(android::base::ReadFileToString(last_command_file_, &last_command_actual));
+    ASSERT_EQ(last_command_expected, last_command_actual);
+  }
+
+  // Resume the update.
+  RunBlockImageUpdate(false, g_entries, image_file_, "t");
+
+  // last_command_file should be gone after successful update.
+  ASSERT_EQ(-1, access(last_command_file_.c_str(), R_OK));
+
+  std::string updated_image_actual;
+  ASSERT_TRUE(android::base::ReadFileToString(image_file_, &updated_image_actual));
+  ASSERT_EQ(g_target_image, updated_image_actual);
+}
diff --git a/updater/blockimg.cpp b/updater/blockimg.cpp
index 4adb974cb..bdb64636b 100644
--- a/updater/blockimg.cpp
+++ b/updater/blockimg.cpp
@@ -1498,6 +1498,7 @@ static Value* PerformBlockImageUpdate(const char* name, State* state,
                                       const std::vector<std::unique_ptr<Expr>>& argv,
                                       const CommandMap& command_map, bool dryrun) {
   CommandParameters params = {};
+  stash_map.clear();
   params.canwrite = !dryrun;
 
   LOG(INFO) << "performing " << (dryrun ? "verification" : "update");
diff --git a/updater_sample/README.md b/updater_sample/README.md
index f6c63a7b6..8ec43d3c6 100644
--- a/updater_sample/README.md
+++ b/updater_sample/README.md
@@ -91,6 +91,31 @@ If they doesn't match, sample app calls `applyPayload` again with the same
 parameters, and handles update completion properly using `onPayloadApplicationCompleted`
 callback. The second problem is solved by adding `PAUSED` updater state.
 
+
+## Sample App UI
+
+### Text fields
+
+- `Current Build:` - shows current active build.
+- `Updater state:` - SystemUpdaterSample app state.
+- `Engine status:` - last reported update_engine status.
+- `Engine error:` - last reported payload application error.
+
+### Buttons
+
+- `Reload` - reloads update configs from device storage.
+- `View config` - shows selected update config.
+- `Apply` - applies selected update config.
+- `Stop` - cancel running update, calls `UpdateEngine#cancel`.
+- `Reset` - reset update, calls `UpdateEngine#resetStatus`, can be called
+            only when update is not running.
+- `Suspend` - suspend running update, uses `UpdateEngine#cancel`.
+- `Resume` - resumes suspended update, uses `UpdateEngine#applyPayload`.
+- `Switch Slot` - if `ab_config.force_switch_slot` config set true,
+            this button will be enabled after payload is applied,
+            to switch A/B slot on next reboot.
+
+
 ## Sending HTTP headers from UpdateEngine
 
 Sometimes OTA package server might require some HTTP headers to be present,
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
index 2fe04bdde..e4c09346b 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
@@ -64,8 +64,8 @@ public class UpdateManager {
 
     private AtomicBoolean mManualSwitchSlotRequired = new AtomicBoolean(true);
 
-    /** Validate state only once when app binds to UpdateEngine. */
-    private AtomicBoolean mStateValidityEnsured = new AtomicBoolean(false);
+    /** Synchronize state with engine status only once when app binds to UpdateEngine. */
+    private AtomicBoolean mStateSynchronized = new AtomicBoolean(false);
 
     @GuardedBy("mLock")
     private UpdateData mLastUpdateData = null;
@@ -90,10 +90,12 @@ public class UpdateManager {
     }
 
     /**
-     * Binds to {@link UpdateEngine}.
+     * Binds to {@link UpdateEngine}. Invokes onStateChangeCallback if present.
      */
     public void bind() {
-        mStateValidityEnsured.set(false);
+        getOnStateChangeCallback().ifPresent(callback -> callback.accept(mUpdaterState.get()));
+
+        mStateSynchronized.set(false);
         this.mUpdateEngine.bind(mUpdateEngineCallback);
     }
 
@@ -104,11 +106,8 @@ public class UpdateManager {
         this.mUpdateEngine.unbind();
     }
 
-    /**
-     * @return a number from {@code 0.0} to {@code 1.0}.
-     */
-    public float getProgress() {
-        return (float) this.mProgress.get();
+    public int getUpdaterState() {
+        return mUpdaterState.get();
     }
 
     /**
@@ -202,21 +201,40 @@ public class UpdateManager {
      * Updates {@link this.mState} and if state is changed,
      * it also notifies {@link this.mOnStateChangeCallback}.
      */
-    private void setUpdaterState(int updaterState) {
+    private void setUpdaterState(int newUpdaterState)
+            throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "setUpdaterState invoked newState=" + newUpdaterState);
         int previousState = mUpdaterState.get();
+        mUpdaterState.set(newUpdaterState);
+        if (previousState != newUpdaterState) {
+            getOnStateChangeCallback().ifPresent(callback -> callback.accept(newUpdaterState));
+        }
+    }
+
+    /**
+     * Same as {@link this.setUpdaterState}. Logs the error if new state
+     * cannot be set.
+     */
+    private void setUpdaterStateSilent(int newUpdaterState) {
         try {
-            mUpdaterState.set(updaterState);
+            setUpdaterState(newUpdaterState);
         } catch (UpdaterState.InvalidTransitionException e) {
-            // Note: invalid state transitions should be handled properly,
-            //       but to make sample app simple, we just throw runtime exception.
-            throw new RuntimeException("Can't set state " + updaterState, e);
-        }
-        if (previousState != updaterState) {
-            getOnStateChangeCallback().ifPresent(callback -> callback.accept(updaterState));
+            // Most likely UpdateEngine status and UpdaterSample state got de-synchronized.
+            // To make sample app simple, we don't handle it properly.
+            Log.e(TAG, "Failed to set updater state", e);
         }
     }
 
     /**
+     * Creates new UpdaterState, assigns it to {@link this.mUpdaterState},
+     * and notifies callbacks.
+     */
+    private void initializeUpdateState(int state) {
+        this.mUpdaterState = new UpdaterState(state);
+        getOnStateChangeCallback().ifPresent(callback -> callback.accept(state));
+    }
+
+    /**
      * Requests update engine to stop any ongoing update. If an update has been applied,
      * leave it as is.
      *
@@ -224,13 +242,10 @@ public class UpdateManager {
      * update engine would throw an error when the method is called, and the only way to
      * handle it is to catch the exception.</p>
      */
-    public void cancelRunningUpdate() {
-        try {
-            mUpdateEngine.cancel();
-            setUpdaterState(UpdaterState.IDLE);
-        } catch (Exception e) {
-            Log.w(TAG, "UpdateEngine failed to stop the ongoing update", e);
-        }
+    public synchronized void cancelRunningUpdate() throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "cancelRunningUpdate invoked");
+        setUpdaterState(UpdaterState.IDLE);
+        mUpdateEngine.cancel();
     }
 
     /**
@@ -240,13 +255,10 @@ public class UpdateManager {
      * update engine would throw an error when the method is called, and the only way to
      * handle it is to catch the exception.</p>
      */
-    public void resetUpdate() {
-        try {
-            mUpdateEngine.resetStatus();
-            setUpdaterState(UpdaterState.IDLE);
-        } catch (Exception e) {
-            Log.w(TAG, "UpdateEngine failed to reset the update", e);
-        }
+    public synchronized void resetUpdate() throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "resetUpdate invoked");
+        setUpdaterState(UpdaterState.IDLE);
+        mUpdateEngine.resetStatus();
     }
 
     /**
@@ -255,7 +267,8 @@ public class UpdateManager {
      * <p>UpdateEngine works asynchronously. This method doesn't wait until
      * end of the update.</p>
      */
-    public void applyUpdate(Context context, UpdateConfig config) {
+    public synchronized void applyUpdate(Context context, UpdateConfig config)
+            throws UpdaterState.InvalidTransitionException {
         mEngineErrorCode.set(UpdateEngineErrorCodes.UNKNOWN);
         setUpdaterState(UpdaterState.RUNNING);
 
@@ -277,7 +290,8 @@ public class UpdateManager {
         }
     }
 
-    private void applyAbNonStreamingUpdate(UpdateConfig config) {
+    private void applyAbNonStreamingUpdate(UpdateConfig config)
+            throws UpdaterState.InvalidTransitionException {
         UpdateData.Builder builder = UpdateData.builder()
                 .setExtraProperties(prepareExtraProperties(config));
 
@@ -306,7 +320,7 @@ public class UpdateManager {
                 updateEngineApplyPayload(builder.build());
             } else {
                 Log.e(TAG, "PrepareStreamingService failed, result code is " + code);
-                setUpdaterState(UpdaterState.ERROR);
+                setUpdaterStateSilent(UpdaterState.ERROR);
             }
         });
     }
@@ -333,6 +347,8 @@ public class UpdateManager {
      * with the given id.</p>
      */
     private void updateEngineApplyPayload(UpdateData update) {
+        Log.d(TAG, "updateEngineApplyPayload invoked with url " + update.mPayload.getUrl());
+
         synchronized (mLock) {
             mLastUpdateData = update;
         }
@@ -348,11 +364,15 @@ public class UpdateManager {
                     properties.toArray(new String[0]));
         } catch (Exception e) {
             Log.e(TAG, "UpdateEngine failed to apply the update", e);
-            setUpdaterState(UpdaterState.ERROR);
+            setUpdaterStateSilent(UpdaterState.ERROR);
         }
     }
 
+    /**
+     * Re-applies {@link this.mLastUpdateData} to update_engine.
+     */
     private void updateEngineReApplyPayload() {
+        Log.d(TAG, "updateEngineReApplyPayload invoked");
         UpdateData lastUpdate;
         synchronized (mLock) {
             // mLastPayloadSpec might be empty in some cases.
@@ -377,8 +397,14 @@ public class UpdateManager {
      * {@link UpdateEngine#applyPayload} might take several seconds to finish, and it will
      * invoke callbacks {@link this#onStatusUpdate} and {@link this#onPayloadApplicationComplete)}.
      */
-    public void setSwitchSlotOnReboot() {
+    public synchronized void setSwitchSlotOnReboot() {
         Log.d(TAG, "setSwitchSlotOnReboot invoked");
+
+        // When mManualSwitchSlotRequired set false, next time
+        // onApplicationPayloadComplete is called,
+        // it will set updater state to REBOOT_REQUIRED.
+        mManualSwitchSlotRequired.set(false);
+
         UpdateData.Builder builder;
         synchronized (mLock) {
             // To make sample app simple, we don't handle it.
@@ -396,65 +422,62 @@ public class UpdateManager {
     }
 
     /**
-     * Verifies if mUpdaterState matches mUpdateEngineStatus.
-     * If they don't match, runs applyPayload to trigger onPayloadApplicationComplete
-     * callback, which updates mUpdaterState.
+     * Synchronize UpdaterState with UpdateEngine status.
+     * Apply necessary UpdateEngine operation if status are out of sync.
+     *
+     * It's expected to be called once when sample app binds itself to UpdateEngine.
      */
-    private void ensureCorrectUpdaterState() {
-        // When mUpdaterState is one of IDLE, PAUSED, ERROR, SLOT_SWITCH_REQUIRED
-        //    then mUpdateEngineStatus must be IDLE.
-        // When mUpdaterState is RUNNING,
-        //    then mUpdateEngineStatus must not be IDLE or UPDATED_NEED_REBOOT.
-        // When mUpdaterState is REBOOT_REQUIRED,
-        //    then mUpdateEngineStatus must be UPDATED_NEED_REBOOT.
-        int state = mUpdaterState.get();
-        int updateEngineStatus = mUpdateEngineStatus.get();
-        if (state == UpdaterState.IDLE
-                || state == UpdaterState.ERROR
-                || state == UpdaterState.PAUSED
-                || state == UpdaterState.SLOT_SWITCH_REQUIRED) {
-            ensureUpdateEngineStatusIdle(state, updateEngineStatus);
-        } else if (state == UpdaterState.RUNNING) {
-            ensureUpdateEngineStatusRunning(state, updateEngineStatus);
-        } else if (state == UpdaterState.REBOOT_REQUIRED) {
-            ensureUpdateEngineStatusReboot(state, updateEngineStatus);
-        }
-    }
+    private void synchronizeUpdaterStateWithUpdateEngineStatus() {
+        Log.d(TAG, "synchronizeUpdaterStateWithUpdateEngineStatus is invoked.");
 
-    private void ensureUpdateEngineStatusIdle(int state, int updateEngineStatus) {
-        if (updateEngineStatus == UpdateEngine.UpdateStatusConstants.IDLE) {
-            return;
-        }
-        // It might happen when update is started not from the sample app.
-        // To make the sample app simple, we won't handle this case.
-        throw new RuntimeException("When mUpdaterState is " + state
-                + " mUpdateEngineStatus expected to be "
-                + UpdateEngine.UpdateStatusConstants.IDLE
-                + ", but it is " + updateEngineStatus);
-    }
+        int state = mUpdaterState.get();
+        int engineStatus = mUpdateEngineStatus.get();
 
-    private void ensureUpdateEngineStatusRunning(int state, int updateEngineStatus) {
-        if (updateEngineStatus != UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT
-                && updateEngineStatus != UpdateEngine.UpdateStatusConstants.IDLE) {
+        if (engineStatus == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT) {
+            // If update has been installed before running the sample app,
+            // set state to REBOOT_REQUIRED.
+            initializeUpdateState(UpdaterState.REBOOT_REQUIRED);
             return;
         }
-        // Re-apply latest update. It makes update_engine to invoke
-        // onPayloadApplicationComplete callback. The callback notifies
-        // if update was successful or not.
-        updateEngineReApplyPayload();
-    }
 
-    private void ensureUpdateEngineStatusReboot(int state, int updateEngineStatus) {
-        if (updateEngineStatus == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT) {
-            return;
+        switch (state) {
+            case UpdaterState.IDLE:
+            case UpdaterState.ERROR:
+            case UpdaterState.PAUSED:
+            case UpdaterState.SLOT_SWITCH_REQUIRED:
+                // It might happen when update is started not from the sample app.
+                // To make the sample app simple, we won't handle this case.
+                Preconditions.checkState(
+                        engineStatus == UpdateEngine.UpdateStatusConstants.IDLE,
+                        "When mUpdaterState is %s, mUpdateEngineStatus "
+                                + "must be 0/IDLE, but it is %s",
+                        state,
+                        engineStatus);
+                break;
+            case UpdaterState.RUNNING:
+                if (engineStatus == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT
+                        || engineStatus == UpdateEngine.UpdateStatusConstants.IDLE) {
+                    Log.i(TAG, "ensureUpdateEngineStatusIsRunning - re-applying last payload");
+                    // Re-apply latest update. It makes update_engine to invoke
+                    // onPayloadApplicationComplete callback. The callback notifies
+                    // if update was successful or not.
+                    updateEngineReApplyPayload();
+                }
+                break;
+            case UpdaterState.REBOOT_REQUIRED:
+                // This might happen when update is installed by other means,
+                // and sample app is not aware of it.
+                // To make the sample app simple, we won't handle this case.
+                Preconditions.checkState(
+                        engineStatus == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT,
+                        "When mUpdaterState is %s, mUpdateEngineStatus "
+                                + "must be 6/UPDATED_NEED_REBOOT, but it is %s",
+                        state,
+                        engineStatus);
+                break;
+            default:
+                throw new IllegalStateException("This block should not be reached.");
         }
-        // This might happen when update is installed by other means,
-        // and sample app is not aware of it. To make the sample app simple,
-        // we won't handle this case.
-        throw new RuntimeException("When mUpdaterState is " + state
-                + " mUpdateEngineStatus expected to be "
-                + UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT
-                + ", but it is " + updateEngineStatus);
     }
 
     /**
@@ -468,13 +491,19 @@ public class UpdateManager {
      * @param progress a number from 0.0 to 1.0.
      */
     private void onStatusUpdate(int status, float progress) {
+        Log.d(TAG, String.format(
+                        "onStatusUpdate invoked, status=%s, progress=%.2f",
+                        status,
+                        progress));
+
         int previousStatus = mUpdateEngineStatus.get();
         mUpdateEngineStatus.set(status);
         mProgress.set(progress);
 
-        if (!mStateValidityEnsured.getAndSet(true)) {
-            // We ensure correct state once only when sample app is bound to UpdateEngine.
-            ensureCorrectUpdaterState();
+        if (!mStateSynchronized.getAndSet(true)) {
+            // We synchronize state with engine status once
+            // only when sample app is bound to UpdateEngine.
+            synchronizeUpdaterStateWithUpdateEngineStatus();
         }
 
         getOnProgressUpdateCallback().ifPresent(callback -> callback.accept(progress));
@@ -489,11 +518,11 @@ public class UpdateManager {
         mEngineErrorCode.set(errorCode);
         if (errorCode == UpdateEngine.ErrorCodeConstants.SUCCESS
                 || errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
-            setUpdaterState(isManualSwitchSlotRequired()
+            setUpdaterStateSilent(isManualSwitchSlotRequired()
                     ? UpdaterState.SLOT_SWITCH_REQUIRED
                     : UpdaterState.REBOOT_REQUIRED);
         } else if (errorCode != UpdateEngineErrorCodes.USER_CANCELLED) {
-            setUpdaterState(UpdaterState.ERROR);
+            setUpdaterStateSilent(UpdaterState.ERROR);
         }
 
         getOnEngineCompleteCallback()
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
index 36a90982e..573d336e9 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
@@ -51,12 +51,15 @@ public class UpdaterState {
      * are allowed to transition to from key.
      */
     private static final ImmutableMap<Integer, ImmutableSet<Integer>> TRANSITIONS =
-            ImmutableMap.of(
-                    IDLE, ImmutableSet.of(RUNNING),
-                    RUNNING, ImmutableSet.of(ERROR, PAUSED, REBOOT_REQUIRED, SLOT_SWITCH_REQUIRED),
-                    PAUSED, ImmutableSet.of(RUNNING),
-                    SLOT_SWITCH_REQUIRED, ImmutableSet.of(ERROR)
-            );
+            ImmutableMap.<Integer, ImmutableSet<Integer>>builder()
+                    .put(IDLE, ImmutableSet.of(ERROR, RUNNING))
+                    .put(RUNNING, ImmutableSet.of(
+                            ERROR, PAUSED, REBOOT_REQUIRED, SLOT_SWITCH_REQUIRED))
+                    .put(PAUSED, ImmutableSet.of(ERROR, RUNNING, IDLE))
+                    .put(SLOT_SWITCH_REQUIRED, ImmutableSet.of(ERROR, IDLE))
+                    .put(ERROR, ImmutableSet.of(IDLE))
+                    .put(REBOOT_REQUIRED, ImmutableSet.of(IDLE))
+                    .build();
 
     private AtomicInteger mState;
 
diff --git a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
index 1de72c2d6..6c71cb6f4 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -88,7 +88,7 @@ public class MainActivity extends Activity {
 
         this.mTextViewConfigsDirHint.setText(UpdateConfigs.getConfigsRoot(this));
 
-        uiReset();
+        uiResetWidgets();
         loadUpdateConfigs();
 
         this.mUpdateManager.setOnStateChangeCallback(this::onUpdaterStateChange);
@@ -108,7 +108,6 @@ public class MainActivity extends Activity {
     @Override
     protected void onResume() {
         super.onResume();
-        // TODO(zhomart) load saved states
         // Binding to UpdateEngine invokes onStatusUpdate callback,
         // persisted updater state has to be loaded and prepared beforehand.
         this.mUpdateManager.bind();
@@ -117,7 +116,6 @@ public class MainActivity extends Activity {
     @Override
     protected void onPause() {
         this.mUpdateManager.unbind();
-        // TODO(zhomart) save state
         super.onPause();
     }
 
@@ -149,14 +147,22 @@ public class MainActivity extends Activity {
                 .setMessage("Do you really want to apply this update?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    uiSetUpdating();
+                    uiResetWidgets();
                     uiResetEngineText();
-                    mUpdateManager.applyUpdate(this, getSelectedConfig());
+                    applyUpdate(getSelectedConfig());
                 })
                 .setNegativeButton(android.R.string.cancel, null)
                 .show();
     }
 
+    private void applyUpdate(UpdateConfig config) {
+        try {
+            mUpdateManager.applyUpdate(this, config);
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to apply update " + config.getName(), e);
+        }
+    }
+
     /**
      * stop button clicked
      */
@@ -166,11 +172,19 @@ public class MainActivity extends Activity {
                 .setMessage("Do you really want to cancel running update?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    mUpdateManager.cancelRunningUpdate();
+                    cancelRunningUpdate();
                 })
                 .setNegativeButton(android.R.string.cancel, null).show();
     }
 
+    private void cancelRunningUpdate() {
+        try {
+            mUpdateManager.cancelRunningUpdate();
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to cancel running update", e);
+        }
+    }
+
     /**
      * reset button clicked
      */
@@ -181,11 +195,19 @@ public class MainActivity extends Activity {
                         + " and restore old version?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    mUpdateManager.resetUpdate();
+                    resetUpdate();
                 })
                 .setNegativeButton(android.R.string.cancel, null).show();
     }
 
+    private void resetUpdate() {
+        try {
+            mUpdateManager.resetUpdate();
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to reset update", e);
+        }
+    }
+
     /**
      * switch slot button clicked
      */
@@ -199,9 +221,25 @@ public class MainActivity extends Activity {
      * values from {@link UpdaterState}.
      */
     private void onUpdaterStateChange(int state) {
-        Log.i(TAG, "onUpdaterStateChange invoked state=" + state);
+        Log.i(TAG, "UpdaterStateChange state="
+                + UpdaterState.getStateText(state)
+                + "/" + state);
         runOnUiThread(() -> {
             setUiUpdaterState(state);
+
+            if (state == UpdaterState.IDLE) {
+                uiStateIdle();
+            } else if (state == UpdaterState.RUNNING) {
+                uiStateRunning();
+            } else if (state == UpdaterState.PAUSED) {
+                uiStatePaused();
+            } else if (state == UpdaterState.ERROR) {
+                uiStateError();
+            } else if (state == UpdaterState.SLOT_SWITCH_REQUIRED) {
+                uiStateSlotSwitchRequired();
+            } else if (state == UpdaterState.REBOOT_REQUIRED) {
+                uiStateRebootRequired();
+            }
         });
     }
 
@@ -210,17 +248,10 @@ public class MainActivity extends Activity {
      * be one of the values from {@link UpdateEngine.UpdateStatusConstants}.
      */
     private void onEngineStatusUpdate(int status) {
+        Log.i(TAG, "StatusUpdate - status="
+                + UpdateEngineStatuses.getStatusText(status)
+                + "/" + status);
         runOnUiThread(() -> {
-            Log.e(TAG, "StatusUpdate - status="
-                    + UpdateEngineStatuses.getStatusText(status)
-                    + "/" + status);
-            if (status == UpdateEngine.UpdateStatusConstants.IDLE) {
-                Log.d(TAG, "status changed, resetting ui");
-                uiReset();
-            } else {
-                Log.d(TAG, "status changed, setting ui to updating mode");
-                uiSetUpdating();
-            }
             setUiEngineStatus(status);
         });
     }
@@ -234,19 +265,12 @@ public class MainActivity extends Activity {
         final String completionState = UpdateEngineErrorCodes.isUpdateSucceeded(errorCode)
                 ? "SUCCESS"
                 : "FAILURE";
+        Log.i(TAG,
+                "PayloadApplicationCompleted - errorCode="
+                        + UpdateEngineErrorCodes.getCodeName(errorCode) + "/" + errorCode
+                        + " " + completionState);
         runOnUiThread(() -> {
-            Log.i(TAG,
-                    "Completed - errorCode="
-                            + UpdateEngineErrorCodes.getCodeName(errorCode) + "/" + errorCode
-                            + " " + completionState);
             setUiEngineErrorCode(errorCode);
-            if (errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
-                // if update was successfully applied.
-                if (mUpdateManager.isManualSwitchSlotRequired()) {
-                    // Show "Switch Slot" button.
-                    uiShowSwitchSlotInfo();
-                }
-            }
         });
     }
 
@@ -258,45 +282,66 @@ public class MainActivity extends Activity {
     }
 
     /** resets ui */
-    private void uiReset() {
+    private void uiResetWidgets() {
         mTextViewBuild.setText(Build.DISPLAY);
-        mSpinnerConfigs.setEnabled(true);
-        mButtonReload.setEnabled(true);
-        mButtonApplyConfig.setEnabled(true);
+        mSpinnerConfigs.setEnabled(false);
+        mButtonReload.setEnabled(false);
+        mButtonApplyConfig.setEnabled(false);
         mButtonStop.setEnabled(false);
         mButtonReset.setEnabled(false);
-        mProgressBar.setProgress(0);
         mProgressBar.setEnabled(false);
         mProgressBar.setVisibility(ProgressBar.INVISIBLE);
-        uiHideSwitchSlotInfo();
+        mButtonSwitchSlot.setEnabled(false);
+        mTextViewUpdateInfo.setTextColor(Color.parseColor("#aaaaaa"));
     }
 
     private void uiResetEngineText() {
         mTextViewEngineStatus.setText(R.string.unknown);
         mTextViewEngineErrorCode.setText(R.string.unknown);
-        // Note: Do not reset mTextViewUpdaterState; UpdateManager notifies properly.
+        // Note: Do not reset mTextViewUpdaterState; UpdateManager notifies updater state properly.
     }
 
-    /** sets ui updating mode */
-    private void uiSetUpdating() {
-        mTextViewBuild.setText(Build.DISPLAY);
-        mSpinnerConfigs.setEnabled(false);
-        mButtonReload.setEnabled(false);
-        mButtonApplyConfig.setEnabled(false);
-        mButtonStop.setEnabled(true);
+    private void uiStateIdle() {
+        uiResetWidgets();
+        mSpinnerConfigs.setEnabled(true);
+        mButtonReload.setEnabled(true);
+        mButtonApplyConfig.setEnabled(true);
+        mProgressBar.setProgress(0);
+    }
+
+    private void uiStateRunning() {
+        uiResetWidgets();
         mProgressBar.setEnabled(true);
+        mProgressBar.setVisibility(ProgressBar.VISIBLE);
+        mButtonStop.setEnabled(true);
+    }
+
+    private void uiStatePaused() {
+        uiResetWidgets();
         mButtonReset.setEnabled(true);
+        mProgressBar.setEnabled(true);
         mProgressBar.setVisibility(ProgressBar.VISIBLE);
     }
 
-    private void uiShowSwitchSlotInfo() {
+    private void uiStateSlotSwitchRequired() {
+        uiResetWidgets();
+        mButtonReset.setEnabled(true);
+        mProgressBar.setEnabled(true);
+        mProgressBar.setVisibility(ProgressBar.VISIBLE);
         mButtonSwitchSlot.setEnabled(true);
         mTextViewUpdateInfo.setTextColor(Color.parseColor("#777777"));
     }
 
-    private void uiHideSwitchSlotInfo() {
-        mTextViewUpdateInfo.setTextColor(Color.parseColor("#AAAAAA"));
-        mButtonSwitchSlot.setEnabled(false);
+    private void uiStateError() {
+        uiResetWidgets();
+        mButtonReset.setEnabled(true);
+        mProgressBar.setEnabled(true);
+        mProgressBar.setVisibility(ProgressBar.VISIBLE);
+    }
+
+    private void uiStateRebootRequired() {
+        uiResetWidgets();
+        mButtonReset.setEnabled(true);
     }
 
     /**
diff --git a/updater_sample/tests/res/raw/update_config_stream_001.json b/updater_sample/tests/res/raw/update_config_001_stream.json
index be51b7c95..be51b7c95 100644
--- a/updater_sample/tests/res/raw/update_config_stream_001.json
+++ b/updater_sample/tests/res/raw/update_config_001_stream.json
diff --git a/updater_sample/tests/res/raw/update_config_stream_002.json b/updater_sample/tests/res/raw/update_config_002_stream.json
index 5d7874cdb..5d7874cdb 100644
--- a/updater_sample/tests/res/raw/update_config_stream_002.json
+++ b/updater_sample/tests/res/raw/update_config_002_stream.json
diff --git a/updater_sample/tests/res/raw/update_config_003_nonstream.json b/updater_sample/tests/res/raw/update_config_003_nonstream.json
new file mode 100644
index 000000000..4175c35ea
--- /dev/null
+++ b/updater_sample/tests/res/raw/update_config_003_nonstream.json
@@ -0,0 +1,9 @@
+{
+    "__": "*** Generated using tools/gen_update_config.py ***",
+    "ab_config": {
+        "force_switch_slot": false
+    },
+    "ab_install_type": "NON_STREAMING",
+    "name": "S ota_002_package",
+    "url": "file:///data/sample-ota-packages/ota_003_package.zip"
+}
\ No newline at end of file
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
index 000f5663b..1cbd8601e 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
@@ -60,7 +60,7 @@ public class UpdateConfigTest {
     public void setUp() throws Exception {
         mContext = InstrumentationRegistry.getContext();
         mTargetContext = InstrumentationRegistry.getTargetContext();
-        mJsonStreaming001 = readResource(R.raw.update_config_stream_001);
+        mJsonStreaming001 = readResource(R.raw.update_config_001_stream);
     }
 
     @Test
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
index 0657a5eb6..e05ad290c 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
@@ -18,19 +18,22 @@ package com.example.android.systemupdatersample;
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.content.Context;
 import android.os.UpdateEngine;
 import android.os.UpdateEngineCallback;
+import android.support.test.InstrumentationRegistry;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
 
+import com.example.android.systemupdatersample.tests.R;
 import com.example.android.systemupdatersample.util.PayloadSpecs;
+import com.google.common.collect.ImmutableList;
+import com.google.common.io.CharStreams;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -40,7 +43,9 @@ import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
-import java.util.function.IntConsumer;
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStreamReader;
 
 /**
  * Tests for {@link UpdateManager}
@@ -56,37 +61,86 @@ public class UpdateManagerTest {
     private UpdateEngine mUpdateEngine;
     @Mock
     private PayloadSpecs mPayloadSpecs;
-    private UpdateManager mUpdateManager;
+    private UpdateManager mSubject;
+    private Context mContext;
+    private UpdateConfig mNonStreamingUpdate003;
 
     @Before
-    public void setUp() {
-        mUpdateManager = new UpdateManager(mUpdateEngine, mPayloadSpecs);
+    public void setUp() throws Exception {
+        mContext = InstrumentationRegistry.getContext();
+        mSubject = new UpdateManager(mUpdateEngine, mPayloadSpecs);
+        mNonStreamingUpdate003 =
+                UpdateConfig.fromJson(readResource(R.raw.update_config_003_nonstream));
     }
 
     @Test
-    public void storesProgressThenInvokesCallbacks() {
-        IntConsumer statusUpdateCallback = mock(IntConsumer.class);
-
-        // When UpdateManager is bound to update_engine, it passes
-        // UpdateManager.UpdateEngineCallbackImpl as a callback to update_engine.
+    public void applyUpdate_appliesPayloadToUpdateEngine() throws Exception {
+        PayloadSpec payload = buildMockPayloadSpec();
+        when(mPayloadSpecs.forNonStreaming(any(File.class))).thenReturn(payload);
         when(mUpdateEngine.bind(any(UpdateEngineCallback.class))).thenAnswer(answer -> {
+            // When UpdateManager is bound to update_engine, it passes
+            // UpdateEngineCallback as a callback to update_engine.
             UpdateEngineCallback callback = answer.getArgument(0);
-            callback.onStatusUpdate(/*engineStatus*/ 4, /*engineProgress*/ 0.2f);
+            callback.onStatusUpdate(
+                    UpdateEngine.UpdateStatusConstants.IDLE,
+                    /*engineProgress*/ 0.0f);
             return null;
         });
 
-        mUpdateManager.setOnEngineStatusUpdateCallback(statusUpdateCallback);
+        mSubject.bind();
+        mSubject.applyUpdate(null, mNonStreamingUpdate003);
+
+        verify(mUpdateEngine).applyPayload(
+                "file://blah",
+                120,
+                340,
+                new String[] {
+                        "SWITCH_SLOT_ON_REBOOT=0" // ab_config.force_switch_slot = false
+                });
+    }
 
-        // Making sure that manager.getProgress() returns correct progress
-        // in "onEngineStatusUpdate" callback.
-        doAnswer(answer -> {
-            assertEquals(0.2f, mUpdateManager.getProgress(), 1E-5);
+    @Test
+    public void stateIsRunningAndEngineStatusIsIdle_reApplyLastUpdate() throws Exception {
+        PayloadSpec payload = buildMockPayloadSpec();
+        when(mPayloadSpecs.forNonStreaming(any(File.class))).thenReturn(payload);
+        when(mUpdateEngine.bind(any(UpdateEngineCallback.class))).thenAnswer(answer -> {
+            // When UpdateManager is bound to update_engine, it passes
+            // UpdateEngineCallback as a callback to update_engine.
+            UpdateEngineCallback callback = answer.getArgument(0);
+            callback.onStatusUpdate(
+                    UpdateEngine.UpdateStatusConstants.IDLE,
+                    /*engineProgress*/ 0.0f);
             return null;
-        }).when(statusUpdateCallback).accept(anyInt());
+        });
 
-        mUpdateManager.bind();
+        mSubject.bind();
+        mSubject.applyUpdate(null, mNonStreamingUpdate003);
+        mSubject.unbind();
+        mSubject.bind(); // re-bind - now it should re-apply last update
+
+        assertEquals(mSubject.getUpdaterState(), UpdaterState.RUNNING);
+        // it should be called 2 times
+        verify(mUpdateEngine, times(2)).applyPayload(
+                "file://blah",
+                120,
+                340,
+                new String[] {
+                        "SWITCH_SLOT_ON_REBOOT=0" // ab_config.force_switch_slot = false
+                });
+    }
+
+    private PayloadSpec buildMockPayloadSpec() {
+        PayloadSpec payload = mock(PayloadSpec.class);
+        when(payload.getUrl()).thenReturn("file://blah");
+        when(payload.getOffset()).thenReturn(120L);
+        when(payload.getSize()).thenReturn(340L);
+        when(payload.getProperties()).thenReturn(ImmutableList.of());
+        return payload;
+    }
 
-        verify(statusUpdateCallback, times(1)).accept(4);
+    private String readResource(int id) throws IOException {
+        return CharStreams.toString(new InputStreamReader(
+                mContext.getResources().openRawResource(id)));
     }
 
 }