Handle retry count correctly

Test: vdc startCheckpoint 2 then reboot 3 times checking state

Change-Id: I4eeda7f73d82a7c8b2469571fa558df2fac47354
diff --git a/fs_mgr/fs_mgr.cpp b/fs_mgr/fs_mgr.cpp
index 3ab9732..07d25f6 100644
--- a/fs_mgr/fs_mgr.cpp
+++ b/fs_mgr/fs_mgr.cpp
@@ -850,56 +850,115 @@
     return true;
 }
 
-bool fs_mgr_update_checkpoint_partition(struct fstab_rec* rec) {
-    if (fs_mgr_is_checkpoint_fs(rec)) {
-        if (!strcmp(rec->fs_type, "f2fs")) {
-            std::string opts(rec->fs_options);
+class CheckpointManager {
+  public:
+    CheckpointManager(int needs_checkpoint = -1) : needs_checkpoint_(needs_checkpoint) {}
 
-            opts += ",checkpoint=disable";
-            free(rec->fs_options);
-            rec->fs_options = strdup(opts.c_str());
-        } else {
-            LERROR << rec->fs_type << " does not implement checkpoints.";
+    bool Update(struct fstab_rec* rec) {
+        if (!fs_mgr_is_checkpoint(rec)) {
+            return true;
         }
-    } else if (fs_mgr_is_checkpoint_blk(rec)) {
-        call_vdc({"checkpoint", "restoreCheckpoint", rec->blk_device});
 
-        android::base::unique_fd fd(
-                TEMP_FAILURE_RETRY(open(rec->blk_device, O_RDONLY | O_CLOEXEC)));
-        if (!fd) {
-            PERROR << "Cannot open device " << rec->blk_device;
+        if (fs_mgr_is_checkpoint_blk(rec)) {
+            call_vdc({"checkpoint", "restoreCheckpoint", rec->blk_device});
+        }
+
+        if (needs_checkpoint_ == UNKNOWN &&
+            !call_vdc_ret({"checkpoint", "needsCheckpoint"}, &needs_checkpoint_)) {
+            LERROR << "Failed to find if checkpointing is needed. Assuming no.";
+            needs_checkpoint_ = NO;
+        }
+
+        if (needs_checkpoint_ != YES) {
+            return true;
+        }
+
+        if (!UpdateCheckpointPartition(rec)) {
+            LERROR << "Could not set up checkpoint partition, skipping!";
             return false;
         }
 
-        uint64_t size = get_block_device_size(fd) / 512;
-        if (!size) {
-            PERROR << "Cannot get device size";
-            return false;
+        return true;
+    }
+
+    bool Revert(struct fstab_rec* rec) {
+        if (!fs_mgr_is_checkpoint(rec)) {
+            return true;
         }
 
-        android::dm::DmTable table;
-        if (!table.AddTarget(
-                    std::make_unique<android::dm::DmTargetBow>(0, size, rec->blk_device))) {
-            LERROR << "Failed to add Bow target";
-            return false;
+        if (device_map_.find(rec->blk_device) == device_map_.end()) {
+            return true;
         }
 
+        std::string bow_device = rec->blk_device;
+        free(rec->blk_device);
+        rec->blk_device = strdup(device_map_[bow_device].c_str());
+        device_map_.erase(bow_device);
+
         DeviceMapper& dm = DeviceMapper::Instance();
-        if (!dm.CreateDevice("bow", table)) {
-            PERROR << "Failed to create bow device";
-            return false;
+        if (!dm.DeleteDevice("bow")) {
+            PERROR << "Failed to remove bow device";
         }
 
-        std::string name;
-        if (!dm.GetDmDevicePathByName("bow", &name)) {
-            PERROR << "Failed to get bow device name";
-            return false;
-        }
-
-        rec->blk_device = strdup(name.c_str());
+        return true;
     }
-    return true;
-}
+
+  private:
+    bool UpdateCheckpointPartition(struct fstab_rec* rec) {
+        if (fs_mgr_is_checkpoint_fs(rec)) {
+            if (!strcmp(rec->fs_type, "f2fs")) {
+                std::string opts(rec->fs_options);
+
+                opts += ",checkpoint=disable";
+                free(rec->fs_options);
+                rec->fs_options = strdup(opts.c_str());
+            } else {
+                LERROR << rec->fs_type << " does not implement checkpoints.";
+            }
+        } else if (fs_mgr_is_checkpoint_blk(rec)) {
+            android::base::unique_fd fd(
+                    TEMP_FAILURE_RETRY(open(rec->blk_device, O_RDONLY | O_CLOEXEC)));
+            if (!fd) {
+                PERROR << "Cannot open device " << rec->blk_device;
+                return false;
+            }
+
+            uint64_t size = get_block_device_size(fd) / 512;
+            if (!size) {
+                PERROR << "Cannot get device size";
+                return false;
+            }
+
+            android::dm::DmTable table;
+            if (!table.AddTarget(
+                        std::make_unique<android::dm::DmTargetBow>(0, size, rec->blk_device))) {
+                LERROR << "Failed to add bow target";
+                return false;
+            }
+
+            DeviceMapper& dm = DeviceMapper::Instance();
+            if (!dm.CreateDevice("bow", table)) {
+                PERROR << "Failed to create bow device";
+                return false;
+            }
+
+            std::string name;
+            if (!dm.GetDmDevicePathByName("bow", &name)) {
+                PERROR << "Failed to get bow device name";
+                return false;
+            }
+
+            device_map_[name] = rec->blk_device;
+            free(rec->blk_device);
+            rec->blk_device = strdup(name.c_str());
+        }
+        return true;
+    }
+
+    enum { UNKNOWN = -1, NO = 0, YES = 1 };
+    int needs_checkpoint_;
+    std::map<std::string, std::string> device_map_;
+};
 
 /* When multiple fstab records share the same mount_point, it will
  * try to mount each one in turn, and ignore any duplicates after a
@@ -913,7 +972,7 @@
     int mret = -1;
     int mount_errno = 0;
     int attempted_idx = -1;
-    int need_checkpoint = -1;
+    CheckpointManager checkpoint_manager;
     FsManagerAvbUniquePtr avb_handle(nullptr);
 
     if (!fstab) {
@@ -960,16 +1019,8 @@
             }
         }
 
-        if (fs_mgr_is_checkpoint(&fstab->recs[i])) {
-            if (need_checkpoint == -1 &&
-                !call_vdc_ret({"checkpoint", "needsCheckpoint"}, &need_checkpoint)) {
-                LERROR << "Failed to find if checkpointing is needed. Assuming no.";
-                need_checkpoint = 0;
-            }
-            if (need_checkpoint == 1 && !fs_mgr_update_checkpoint_partition(&fstab->recs[i])) {
-                LERROR << "Could not set up checkpoint partition, skipping!";
-                continue;
-            }
+        if (!checkpoint_manager.Update(&fstab->recs[i])) {
+            continue;
         }
 
         if (fstab->recs[i].fs_mgr_flags & MF_WAIT &&
@@ -1052,6 +1103,9 @@
                    << " is wiped and " << fstab->recs[top_idx].mount_point
                    << " " << fstab->recs[top_idx].fs_type
                    << " is formattable. Format it.";
+
+            checkpoint_manager.Revert(&fstab->recs[top_idx]);
+
             if (fs_mgr_is_encryptable(&fstab->recs[top_idx]) &&
                 strcmp(fstab->recs[top_idx].key_loc, KEY_IN_FOOTER)) {
                 int fd = open(fstab->recs[top_idx].key_loc, O_WRONLY);
@@ -1172,11 +1226,12 @@
  * in turn, and stop on 1st success, or no more match.
  */
 static int fs_mgr_do_mount_helper(fstab* fstab, const char* n_name, char* n_blk_device,
-                                  char* tmp_mount_point, int need_checkpoint) {
+                                  char* tmp_mount_point, int needs_checkpoint) {
     int i = 0;
     int mount_errors = 0;
     int first_mount_errno = 0;
     char* mount_point;
+    CheckpointManager checkpoint_manager(needs_checkpoint);
     FsManagerAvbUniquePtr avb_handle(nullptr);
 
     if (!fstab) {
@@ -1205,16 +1260,9 @@
             }
         }
 
-        if (fs_mgr_is_checkpoint(&fstab->recs[i])) {
-            if (need_checkpoint == -1 &&
-                !call_vdc_ret({"checkpoint", "needsCheckpoint"}, &need_checkpoint)) {
-                LERROR << "Failed to find if checkpointing is needed. Assuming no.";
-                need_checkpoint = 0;
-            }
-            if (need_checkpoint == 1 && !fs_mgr_update_checkpoint_partition(&fstab->recs[i])) {
-                LERROR << "Could not set up checkpoint partition, skipping!";
-                continue;
-            }
+        if (!checkpoint_manager.Update(&fstab->recs[i])) {
+            LERROR << "Could not set up checkpoint partition, skipping!";
+            continue;
         }
 
         /* First check the filesystem if requested */
@@ -1291,8 +1339,8 @@
 }
 
 int fs_mgr_do_mount(fstab* fstab, const char* n_name, char* n_blk_device, char* tmp_mount_point,
-                    bool needs_cp) {
-    return fs_mgr_do_mount_helper(fstab, n_name, n_blk_device, tmp_mount_point, needs_cp);
+                    bool needs_checkpoint) {
+    return fs_mgr_do_mount_helper(fstab, n_name, n_blk_device, tmp_mount_point, needs_checkpoint);
 }
 
 /*
diff --git a/rootdir/init.rc b/rootdir/init.rc
index 6a6a8f9..9aaad8f 100644
--- a/rootdir/init.rc
+++ b/rootdir/init.rc
@@ -398,6 +398,10 @@
     class_start early_hal
 
 on post-fs-data
+    # Start checkpoint before we touch data
+    start vold
+    exec - system system -- /system/bin/vdc checkpoint prepareCheckpoint
+
     # We chown/chmod /data again so because mount is run as root + defaults
     chown system system /data
     chmod 0771 /data
@@ -405,8 +409,6 @@
     restorecon /data
 
     # Make sure we have the device encryption key.
-    start vold
-    exec - system system -- /system/bin/vdc checkpoint prepareDriveForCheckpoint /data
     installkey /data
 
     # Start bootcharting as soon as possible after the data partition is