diff --git a/Benchmark.cpp b/Benchmark.cpp
index b0a3b85..0770da7 100644
--- a/Benchmark.cpp
+++ b/Benchmark.cpp
@@ -23,8 +23,8 @@
 #include <android-base/logging.h>
 
 #include <cutils/iosched_policy.h>
-#include <hardware_legacy/power.h>
 #include <private/android_filesystem_config.h>
+#include <wakelock/wakelock.h>
 
 #include <thread>
 
@@ -181,7 +181,7 @@
 void Benchmark(const std::string& path,
                const android::sp<android::os::IVoldTaskListener>& listener) {
     std::lock_guard<std::mutex> lock(kBenchmarkLock);
-    acquire_wake_lock(PARTIAL_WAKE_LOCK, kWakeLock);
+    android::wakelock::WakeLock wl{kWakeLock};
 
     PerformanceBoost boost;
     android::os::PersistableBundle extras;
@@ -190,8 +190,6 @@
     if (listener) {
         listener->onFinished(res, extras);
     }
-
-    release_wake_lock(kWakeLock);
 }
 
 }  // namespace vold