diff --git a/drm/VSyncWorker.cpp b/drm/VSyncWorker.cpp
index 1c0de21..6e92838 100644
--- a/drm/VSyncWorker.cpp
+++ b/drm/VSyncWorker.cpp
@@ -37,27 +37,16 @@
       last_timestamp_(-1) {
 }
 
-int VSyncWorker::Init(DrmDevice *drm, int display) {
+auto VSyncWorker::Init(DrmDevice *drm, int display,
+                       std::function<void(uint64_t /*timestamp*/)> callback)
+    -> int {
   drm_ = drm;
   display_ = display;
+  callback_ = std::move(callback);
 
   return InitWorker();
 }
 
-void VSyncWorker::RegisterCallback(std::shared_ptr<VsyncCallback> callback) {
-  Lock();
-  callback_ = std::move(callback);
-  Unlock();
-}
-
-void VSyncWorker::RegisterClientCallback(hwc2_callback_data_t data,
-                                         hwc2_function_pointer_t hook) {
-  Lock();
-  vsync_callback_data_ = data;
-  vsync_callback_hook_ = (HWC2_PFN_VSYNC)hook;
-  Unlock();
-}
-
 void VSyncWorker::VSyncControl(bool enabled) {
   Lock();
   enabled_ = enabled;
@@ -133,7 +122,6 @@
   }
 
   int display = display_;
-  std::shared_ptr<VsyncCallback> callback(callback_);
   Unlock();
 
   DrmCrtc *crtc = drm_->GetCrtcForDisplay(display);
@@ -145,8 +133,9 @@
 
   drmVBlank vblank;
   memset(&vblank, 0, sizeof(vblank));
-  vblank.request.type = (drmVBlankSeqType)(
-      DRM_VBLANK_RELATIVE | (high_crtc & DRM_VBLANK_HIGH_CRTC_MASK));
+  vblank.request.type = (drmVBlankSeqType)(DRM_VBLANK_RELATIVE |
+                                           (high_crtc &
+                                            DRM_VBLANK_HIGH_CRTC_MASK));
   vblank.request.sequence = 1;
 
   int64_t timestamp = 0;
@@ -166,13 +155,9 @@
   if (!enabled_)
     return;
 
-  if (callback)
-    callback->Callback(display, timestamp);
-
-  Lock();
-  if (enabled_ && vsync_callback_hook_ && vsync_callback_data_)
-    vsync_callback_hook_(vsync_callback_data_, display, timestamp);
-  Unlock();
+  if (callback_) {
+    callback_(timestamp);
+  }
 
   last_timestamp_ = timestamp;
 }
diff --git a/drm/VSyncWorker.h b/drm/VSyncWorker.h
index b43918c..74ff487 100644
--- a/drm/VSyncWorker.h
+++ b/drm/VSyncWorker.h
@@ -23,6 +23,7 @@
 #include <stdint.h>
 
 #include <atomic>
+#include <functional>
 #include <map>
 
 #include "DrmDevice.h"
@@ -30,21 +31,13 @@
 
 namespace android {
 
-class VsyncCallback {
- public:
-  virtual ~VsyncCallback() = default;
-  virtual void Callback(int display, int64_t timestamp) = 0;
-};
-
 class VSyncWorker : public Worker {
  public:
   VSyncWorker();
   ~VSyncWorker() override = default;
 
-  int Init(DrmDevice *drm, int display);
-  void RegisterCallback(std::shared_ptr<VsyncCallback> callback);
-  void RegisterClientCallback(hwc2_callback_data_t data,
-                              hwc2_function_pointer_t hook);
+  auto Init(DrmDevice *drm, int display,
+            std::function<void(uint64_t /*timestamp*/)> callback) -> int;
 
   void VSyncControl(bool enabled);
 
@@ -57,17 +50,11 @@
 
   DrmDevice *drm_;
 
-  // shared_ptr since we need to use this outside of the thread lock (to
-  // actually call the hook) and we don't want the memory freed until we're
-  // done
-  std::shared_ptr<VsyncCallback> callback_ = NULL;
+  std::function<void(uint64_t /*timestamp*/)> callback_;
 
   int display_;
   std::atomic_bool enabled_;
   int64_t last_timestamp_;
-
-  hwc2_callback_data_t vsync_callback_data_ = NULL;
-  HWC2_PFN_VSYNC vsync_callback_hook_ = NULL;
 };
 }  // namespace android