diff --git a/compositor/DrmDisplayCompositor.cpp b/compositor/DrmDisplayCompositor.cpp
index 08d998c..5ad2de2 100644
--- a/compositor/DrmDisplayCompositor.cpp
+++ b/compositor/DrmDisplayCompositor.cpp
@@ -49,20 +49,6 @@
   return str << flattenting_state_str[static_cast<int>(state)];
 }
 
-class CompositorVsyncCallback : public VsyncCallback {
- public:
-  explicit CompositorVsyncCallback(DrmDisplayCompositor *compositor)
-      : compositor_(compositor) {
-  }
-
-  void Callback(int display, int64_t timestamp) override {
-    compositor_->Vsync(display, timestamp);
-  }
-
- private:
-  DrmDisplayCompositor *compositor_;
-};
-
 DrmDisplayCompositor::DrmDisplayCompositor()
     : resource_manager_(nullptr),
       display_(-1),
@@ -103,7 +89,10 @@
   pthread_mutex_destroy(&lock_);
 }
 
-int DrmDisplayCompositor::Init(ResourceManager *resource_manager, int display) {
+auto DrmDisplayCompositor::Init(ResourceManager *resource_manager, int display,
+                                std::function<void()> client_refresh_callback)
+    -> int {
+  client_refresh_callback_ = std::move(client_refresh_callback);
   resource_manager_ = resource_manager;
   display_ = display;
   DrmDevice *drm = resource_manager_->GetDrmDevice(display);
@@ -118,9 +107,19 @@
   }
   planner_ = Planner::CreateInstance(drm);
 
-  vsync_worker_.Init(drm, display_);
-  auto callback = std::make_shared<CompositorVsyncCallback>(this);
-  vsync_worker_.RegisterCallback(callback);
+  vsync_worker_.Init(drm, display_, [this](int64_t timestamp) {
+    AutoLock lock(&lock_, "DrmDisplayCompositor::Init()");
+    if (lock.Lock())
+      return;
+    flatten_countdown_--;
+    if (!CountdownExpired())
+      return;
+    lock.Unlock();
+    int ret = FlattenActiveComposition();
+    ALOGV("scene flattening triggered for display %d at timestamp %" PRIu64
+          " result = %d \n",
+          display_, timestamp, ret);
+  });
 
   initialized_ = true;
   return 0;
@@ -445,9 +444,7 @@
 }
 
 int DrmDisplayCompositor::FlattenOnClient() {
-  const std::lock_guard<std::mutex> lock(refresh_callback_lock);
-
-  if (refresh_callback_hook_ && refresh_callback_data_) {
+  if (client_refresh_callback_) {
     {
       AutoLock lock(&lock_, __func__);
       if (!IsFlatteningNeeded()) {
@@ -463,7 +460,7 @@
         "No writeback connector available, "
         "falling back to client composition");
     SetFlattening(FlatteningState::kClientRequested);
-    refresh_callback_hook_(refresh_callback_data_, display_);
+    client_refresh_callback_();
     return 0;
   }
 
@@ -479,20 +476,6 @@
   return flatten_countdown_ <= 0;
 }
 
-void DrmDisplayCompositor::Vsync(int display, int64_t timestamp) {
-  AutoLock lock(&lock_, __func__);
-  if (lock.Lock())
-    return;
-  flatten_countdown_--;
-  if (!CountdownExpired())
-    return;
-  lock.Unlock();
-  int ret = FlattenActiveComposition();
-  ALOGV("scene flattening triggered for display %d at timestamp %" PRIu64
-        " result = %d \n",
-        display, timestamp, ret);
-}
-
 void DrmDisplayCompositor::Dump(std::ostringstream *out) const {
   int ret = pthread_mutex_lock(&lock_);
   if (ret)
diff --git a/compositor/DrmDisplayCompositor.h b/compositor/DrmDisplayCompositor.h
index c0eed0c..afdb79e 100644
--- a/compositor/DrmDisplayCompositor.h
+++ b/compositor/DrmDisplayCompositor.h
@@ -52,25 +52,14 @@
   DrmDisplayCompositor();
   ~DrmDisplayCompositor();
 
-  int Init(ResourceManager *resource_manager, int display);
-
-  hwc2_callback_data_t refresh_callback_data_ = NULL;
-  HWC2_PFN_REFRESH refresh_callback_hook_ = NULL;
-  std::mutex refresh_callback_lock;
-
-  void SetRefreshCallback(hwc2_callback_data_t data,
-                          hwc2_function_pointer_t hook) {
-    const std::lock_guard<std::mutex> lock(refresh_callback_lock);
-    refresh_callback_data_ = data;
-    refresh_callback_hook_ = reinterpret_cast<HWC2_PFN_REFRESH>(hook);
-  }
+  auto Init(ResourceManager *resource_manager, int display,
+            std::function<void()> client_refresh_callback) -> int;
 
   std::unique_ptr<DrmDisplayComposition> CreateInitializedComposition() const;
   int ApplyComposition(std::unique_ptr<DrmDisplayComposition> composition);
   int TestComposition(DrmDisplayComposition *composition);
   int Composite();
   void Dump(std::ostringstream *out) const;
-  void Vsync(int display, int64_t timestamp);
   void ClearDisplay();
   UniqueFd TakeOutFence() {
     if (!active_composition_) {
@@ -86,6 +75,7 @@
   std::tuple<uint32_t, uint32_t, int> GetActiveModeResolution();
 
  private:
+  std::function<void()> client_refresh_callback_;
   struct ModeState {
     bool needs_modeset = false;
     DrmMode mode;