diff --git a/Checkpoint.cpp b/Checkpoint.cpp
index 1020ad7..08f0fce 100644
--- a/Checkpoint.cpp
+++ b/Checkpoint.cpp
@@ -371,9 +371,10 @@
 
 }  // namespace
 
-Status cp_restoreCheckpoint(const std::string& blockDevice) {
+Status cp_restoreCheckpoint(const std::string& blockDevice, int restore_limit) {
     bool validating = true;
     std::string action = "Validating";
+    int restore_count = 0;
 
     for (;;) {
         Relocations relocations;
@@ -449,6 +450,12 @@
                 } else {
                     lseek64(device_fd, le->source * kSectorSize, SEEK_SET);
                     write(device_fd, &buffer[0], le->size);
+                    restore_count++;
+                    if (restore_limit && restore_count >= restore_limit) {
+                        LOG(WARNING) << "Hit the test limit";
+                        status = Status::fromExceptionCode(EAGAIN, "Hit the test limit");
+                        break;
+                    }
                 }
             }
         }
diff --git a/Checkpoint.h b/Checkpoint.h
index 64ceed3..70dad8a 100644
--- a/Checkpoint.h
+++ b/Checkpoint.h
@@ -37,7 +37,7 @@
 
 android::binder::Status cp_prepareCheckpoint();
 
-android::binder::Status cp_restoreCheckpoint(const std::string& mountPoint);
+android::binder::Status cp_restoreCheckpoint(const std::string& mountPoint, int count = 0);
 
 android::binder::Status cp_markBootAttempt();
 
diff --git a/VoldNativeService.cpp b/VoldNativeService.cpp
index 1001d2b..69f8a8c 100644
--- a/VoldNativeService.cpp
+++ b/VoldNativeService.cpp
@@ -850,6 +850,14 @@
     return cp_restoreCheckpoint(mountPoint);
 }
 
+binder::Status VoldNativeService::restoreCheckpointPart(const std::string& mountPoint, int count) {
+    ENFORCE_UID(AID_SYSTEM);
+    CHECK_ARGUMENT_PATH(mountPoint);
+    ACQUIRE_LOCK;
+
+    return cp_restoreCheckpoint(mountPoint, count);
+}
+
 binder::Status VoldNativeService::markBootAttempt() {
     ENFORCE_UID(AID_SYSTEM);
     ACQUIRE_LOCK;
diff --git a/VoldNativeService.h b/VoldNativeService.h
index 7db3e5c..954b8ae 100644
--- a/VoldNativeService.h
+++ b/VoldNativeService.h
@@ -127,6 +127,7 @@
     binder::Status commitChanges();
     binder::Status prepareCheckpoint();
     binder::Status restoreCheckpoint(const std::string& mountPoint);
+    binder::Status restoreCheckpointPart(const std::string& mountPoint, int count);
     binder::Status markBootAttempt();
     binder::Status abortChanges();
     binder::Status supportsCheckpoint(bool* _aidl_return);
diff --git a/binder/android/os/IVold.aidl b/binder/android/os/IVold.aidl
index 4b21078..83ee116 100644
--- a/binder/android/os/IVold.aidl
+++ b/binder/android/os/IVold.aidl
@@ -103,6 +103,7 @@
     void commitChanges();
     void prepareCheckpoint();
     void restoreCheckpoint(@utf8InCpp String device);
+    void restoreCheckpointPart(@utf8InCpp String device, int count);
     void markBootAttempt();
     boolean supportsCheckpoint();
 
diff --git a/vdc.cpp b/vdc.cpp
index 35775a7..d01fb49 100644
--- a/vdc.cpp
+++ b/vdc.cpp
@@ -127,6 +127,10 @@
         checkStatus(vold->prepareCheckpoint());
     } else if (args[0] == "checkpoint" && args[1] == "restoreCheckpoint" && args.size() == 3) {
         checkStatus(vold->restoreCheckpoint(args[2]));
+    } else if (args[0] == "checkpoint" && args[1] == "restoreCheckpointPart" && args.size() == 4) {
+        int count;
+        if (!android::base::ParseInt(args[3], &count)) exit(EINVAL);
+        checkStatus(vold->restoreCheckpointPart(args[2], count));
     } else if (args[0] == "checkpoint" && args[1] == "markBootAttempt" && args.size() == 2) {
         checkStatus(vold->markBootAttempt());
     } else if (args[0] == "checkpoint" && args[1] == "abortChanges" && args.size() == 2) {