diff --git a/fs_mgr/fs_mgr.cpp b/fs_mgr/fs_mgr.cpp
index 78d0c07..28da9db 100644
--- a/fs_mgr/fs_mgr.cpp
+++ b/fs_mgr/fs_mgr.cpp
@@ -47,8 +47,9 @@
 #include <logwrap/logwrap.h>
 #include <private/android_logger.h>  // for __android_log_is_debuggable()
 
+#include "fs_mgr.h"
+#include "fs_mgr_avb.h"
 #include "fs_mgr_priv.h"
-#include "fs_mgr_priv_avb.h"
 
 #define KEY_LOC_PROP   "ro.crypto.keyfile.userdata"
 #define KEY_IN_FOOTER  "footer"
@@ -819,7 +820,7 @@
                     return -1;
                 }
             }
-            if (!avb_handle->SetUpAvb(&fstab->recs[i])) {
+            if (!avb_handle->SetUpAvb(&fstab->recs[i], true /* wait_for_verity_dev */)) {
                 LERROR << "Failed to set up AVB on partition: "
                        << fstab->recs[i].mount_point << ", skipping!";
                 /* Skips mounting the device. */
@@ -1031,7 +1032,7 @@
                     return -1;
                 }
             }
-            if (!avb_handle->SetUpAvb(&fstab->recs[i])) {
+            if (!avb_handle->SetUpAvb(&fstab->recs[i], true /* wait_for_verity_dev */)) {
                 LERROR << "Failed to set up AVB on partition: "
                        << fstab->recs[i].mount_point << ", skipping!";
                 /* Skips mounting the device. */
diff --git a/fs_mgr/fs_mgr_avb.cpp b/fs_mgr/fs_mgr_avb.cpp
index 76ec236..7c82bb1 100644
--- a/fs_mgr/fs_mgr_avb.cpp
+++ b/fs_mgr/fs_mgr_avb.cpp
@@ -38,9 +38,9 @@
 #include <utils/Compat.h>
 
 #include "fs_mgr.h"
+#include "fs_mgr_avb.h"
 #include "fs_mgr_avb_ops.h"
 #include "fs_mgr_priv.h"
-#include "fs_mgr_priv_avb.h"
 #include "fs_mgr_priv_dm_ioctl.h"
 #include "fs_mgr_priv_sha.h"
 
@@ -336,7 +336,8 @@
 
 static bool hashtree_dm_verity_setup(struct fstab_rec* fstab_entry,
                                      const AvbHashtreeDescriptor& hashtree_desc,
-                                     const std::string& salt, const std::string& root_digest) {
+                                     const std::string& salt, const std::string& root_digest,
+                                     bool wait_for_verity_dev) {
     // Gets the device mapper fd.
     android::base::unique_fd fd(open("/dev/device-mapper", O_RDWR));
     if (fd < 0) {
@@ -375,13 +376,12 @@
     // Marks the underlying block device as read-only.
     fs_mgr_set_blk_ro(fstab_entry->blk_device);
 
-    // TODO(bowgotsai): support verified all partition at boot.
     // Updates fstab_rec->blk_device to verity device name.
     free(fstab_entry->blk_device);
     fstab_entry->blk_device = strdup(verity_blk_name.c_str());
 
     // Makes sure we've set everything up properly.
-    if (fs_mgr_test_access(verity_blk_name.c_str()) < 0) {
+    if (wait_for_verity_dev && fs_mgr_test_access(verity_blk_name.c_str()) < 0) {
         return false;
     }
 
@@ -519,7 +519,7 @@
     return nullptr;
 }
 
-bool FsManagerAvbHandle::SetUpAvb(struct fstab_rec* fstab_entry) {
+bool FsManagerAvbHandle::SetUpAvb(struct fstab_rec* fstab_entry, bool wait_for_verity_dev) {
     if (!fstab_entry) return false;
     if (!avb_slot_data_ || avb_slot_data_->num_vbmeta_images < 1) {
         return false;
@@ -545,7 +545,8 @@
     }
 
     // Converts HASHTREE descriptor to verity_table_params.
-    if (!hashtree_dm_verity_setup(fstab_entry, hashtree_descriptor, salt, root_digest)) {
+    if (!hashtree_dm_verity_setup(fstab_entry, hashtree_descriptor, salt, root_digest,
+                                  wait_for_verity_dev)) {
         return false;
     }
     return true;
diff --git a/fs_mgr/fs_mgr_fstab.cpp b/fs_mgr/fs_mgr_fstab.cpp
index c231a23..0a694c1 100644
--- a/fs_mgr/fs_mgr_fstab.cpp
+++ b/fs_mgr/fs_mgr_fstab.cpp
@@ -814,6 +814,11 @@
     return fstab->fs_mgr_flags & MF_VERIFY;
 }
 
+int fs_mgr_is_avb(const struct fstab_rec *fstab)
+{
+    return fstab->fs_mgr_flags & MF_AVB;
+}
+
 int fs_mgr_is_verifyatboot(const struct fstab_rec *fstab)
 {
     return fstab->fs_mgr_flags & MF_VERIFYATBOOT;
diff --git a/fs_mgr/include/fs_mgr.h b/fs_mgr/include/fs_mgr.h
index ec6a415..458176b 100644
--- a/fs_mgr/include/fs_mgr.h
+++ b/fs_mgr/include/fs_mgr.h
@@ -124,6 +124,7 @@
 int fs_mgr_is_nonremovable(const struct fstab_rec *fstab);
 int fs_mgr_is_verified(const struct fstab_rec *fstab);
 int fs_mgr_is_verifyatboot(const struct fstab_rec *fstab);
+int fs_mgr_is_avb(const struct fstab_rec *fstab);
 int fs_mgr_is_encryptable(const struct fstab_rec *fstab);
 int fs_mgr_is_file_encrypted(const struct fstab_rec *fstab);
 void fs_mgr_get_file_encryption_modes(const struct fstab_rec *fstab,
diff --git a/fs_mgr/fs_mgr_priv_avb.h b/fs_mgr/include/fs_mgr_avb.h
similarity index 83%
rename from fs_mgr/fs_mgr_priv_avb.h
rename to fs_mgr/include/fs_mgr_avb.h
index 99a033e..526a5ce 100644
--- a/fs_mgr/fs_mgr_priv_avb.h
+++ b/fs_mgr/include/fs_mgr_avb.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2016 The Android Open Source Project
+ * Copyright (C) 2017 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.
@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-#ifndef __CORE_FS_MGR_PRIV_AVB_H
-#define __CORE_FS_MGR_PRIV_AVB_H
+#ifndef __CORE_FS_MGR_AVB_H
+#define __CORE_FS_MGR_AVB_H
 
 #include <memory>
 #include <string>
@@ -63,18 +63,22 @@
     static FsManagerAvbUniquePtr Open(const std::string& device_file_by_name_prefix);
 
     // Sets up dm-verity on the given fstab entry.
+    // The 'wait_for_verity_dev' parameter makes this function wait for the
+    // verity device to get created before return.
     // Returns true if the mount point is eligible to mount, it includes:
     //   - status_ is kFsMgrAvbHandleHashtreeDisabled or
     //   - status_ is kFsMgrAvbHandleSuccess and sending ioctl DM_TABLE_LOAD
     //     to load verity table is success.
     // Otherwise, returns false.
-    bool SetUpAvb(fstab_rec* fstab_entry);
+    bool SetUpAvb(fstab_rec* fstab_entry, bool wait_for_verity_dev);
 
-    FsManagerAvbHandle(const FsManagerAvbHandle&) = delete; // no copy
-    FsManagerAvbHandle& operator=(const FsManagerAvbHandle&) = delete; // no assignment
+    bool AvbHashtreeDisabled() { return status_ == kFsManagerAvbHandleHashtreeDisabled; }
 
-    FsManagerAvbHandle(FsManagerAvbHandle&&) noexcept = delete;  // no move
-    FsManagerAvbHandle& operator=(FsManagerAvbHandle&&) noexcept = delete; // no move assignment
+    FsManagerAvbHandle(const FsManagerAvbHandle&) = delete;             // no copy
+    FsManagerAvbHandle& operator=(const FsManagerAvbHandle&) = delete;  // no assignment
+
+    FsManagerAvbHandle(FsManagerAvbHandle&&) noexcept = delete;             // no move
+    FsManagerAvbHandle& operator=(FsManagerAvbHandle&&) noexcept = delete;  // no move assignment
 
     ~FsManagerAvbHandle() {
         if (avb_slot_data_) {
@@ -90,4 +94,4 @@
     FsManagerAvbHandleStatus status_;
 };
 
-#endif /* __CORE_FS_MGR_PRIV_AVB_H */
+#endif /* __CORE_FS_MGR_AVB_H */