diff --git a/DrmHwcTwo.cpp b/DrmHwcTwo.cpp
index 65a317c..81bb96d 100644
--- a/DrmHwcTwo.cpp
+++ b/DrmHwcTwo.cpp
@@ -32,22 +32,6 @@
 
 namespace android {
 
-class DrmVsyncCallback : public VsyncCallback {
- public:
-  DrmVsyncCallback(hwc2_callback_data_t data, hwc2_function_pointer_t hook)
-      : data_(data), hook_(hook) {
-  }
-
-  void Callback(int display, int64_t timestamp) {
-    auto hook = reinterpret_cast<HWC2_PFN_VSYNC>(hook_);
-    hook(data_, display, timestamp);
-  }
-
- private:
-  hwc2_callback_data_t data_;
-  hwc2_function_pointer_t hook_;
-};
-
 DrmHwcTwo::DrmHwcTwo() {
   common.tag = HARDWARE_DEVICE_TAG;
   common.version = HWC_DEVICE_API_VERSION_2_0;
@@ -194,17 +178,10 @@
                                         hwc2_callback_data_t data,
                                         hwc2_function_pointer_t function) {
   supported(__func__);
-  auto callback = static_cast<HWC2::Callback>(descriptor);
 
-  if (!function) {
-    callbacks_.erase(callback);
-    return HWC2::Error::None;
-  }
-
-  callbacks_.emplace(callback, HwcCallback(data, function));
-
-  switch (callback) {
+  switch (static_cast<HWC2::Callback>(descriptor)) {
     case HWC2::Callback::Hotplug: {
+      SetHotplugCallback(data, function);
       auto &drmDevices = resource_manager_.getDrmDevices();
       for (auto &device : drmDevices)
         HandleInitialHotplugState(device.get());
@@ -317,21 +294,16 @@
   return SetActiveConfig(connector_->get_preferred_mode_id());
 }
 
-HWC2::Error DrmHwcTwo::HwcDisplay::RegisterVsyncCallback(
+void DrmHwcTwo::HwcDisplay::RegisterVsyncCallback(
     hwc2_callback_data_t data, hwc2_function_pointer_t func) {
   supported(__func__);
-  auto callback = std::make_shared<DrmVsyncCallback>(data, func);
-  vsync_worker_.RegisterCallback(std::move(callback));
-  return HWC2::Error::None;
+  vsync_worker_.RegisterClientCallback(data, func);
 }
 
 void DrmHwcTwo::HwcDisplay::RegisterRefreshCallback(
     hwc2_callback_data_t data, hwc2_function_pointer_t func) {
   supported(__func__);
-  auto hook = reinterpret_cast<HWC2_PFN_REFRESH>(func);
-  compositor_.SetRefreshCallback([data, hook](int display) {
-    hook(data, static_cast<hwc2_display_t>(display));
-  });
+  compositor_.SetRefreshCallback(data, func);
 }
 
 HWC2::Error DrmHwcTwo::HwcDisplay::AcceptDisplayChanges() {
@@ -1110,14 +1082,13 @@
 }
 
 void DrmHwcTwo::HandleDisplayHotplug(hwc2_display_t displayid, int state) {
-  auto cb = callbacks_.find(HWC2::Callback::Hotplug);
-  if (cb == callbacks_.end())
-    return;
+  const std::lock_guard<std::mutex> lock(hotplug_callback_lock);
 
-  auto hotplug = reinterpret_cast<HWC2_PFN_HOTPLUG>(cb->second.func);
-  hotplug(cb->second.data, displayid,
-          (state == DRM_MODE_CONNECTED ? HWC2_CONNECTION_CONNECTED
-                                       : HWC2_CONNECTION_DISCONNECTED));
+  if (hotplug_callback_hook_ && hotplug_callback_data_)
+    hotplug_callback_hook_(hotplug_callback_data_, displayid,
+                           state == DRM_MODE_CONNECTED
+                               ? HWC2_CONNECTION_CONNECTED
+                               : HWC2_CONNECTION_DISCONNECTED);
 }
 
 void DrmHwcTwo::HandleInitialHotplugState(DrmDevice *drmDevice) {
diff --git a/DrmHwcTwo.h b/DrmHwcTwo.h
index ca7a00b..7c3b856 100644
--- a/DrmHwcTwo.h
+++ b/DrmHwcTwo.h
@@ -42,6 +42,17 @@
 
   HWC2::Error Init();
 
+  hwc2_callback_data_t hotplug_callback_data_ = NULL;
+  HWC2_PFN_HOTPLUG hotplug_callback_hook_ = NULL;
+  std::mutex hotplug_callback_lock;
+
+  void SetHotplugCallback(hwc2_callback_data_t data,
+                          hwc2_function_pointer_t hook) {
+    const std::lock_guard<std::mutex> lock(hotplug_callback_lock);
+    hotplug_callback_data_ = data;
+    hotplug_callback_hook_ = reinterpret_cast<HWC2_PFN_HOTPLUG>(hook);
+  }
+
   class HwcLayer {
    public:
     HWC2::Composition sf_type() const {
@@ -149,14 +160,6 @@
     android_dataspace_t dataspace_ = HAL_DATASPACE_UNKNOWN;
   };
 
-  struct HwcCallback {
-    HwcCallback(hwc2_callback_data_t d, hwc2_function_pointer_t f)
-        : data(d), func(f) {
-    }
-    hwc2_callback_data_t data;
-    hwc2_function_pointer_t func;
-  };
-
   class HwcDisplay {
    public:
     HwcDisplay(ResourceManager *resource_manager, DrmDevice *drm,
@@ -165,8 +168,8 @@
     HwcDisplay(const HwcDisplay &) = delete;
     HWC2::Error Init(std::vector<DrmPlane *> *planes);
 
-    HWC2::Error RegisterVsyncCallback(hwc2_callback_data_t data,
-                                      hwc2_function_pointer_t func);
+    void RegisterVsyncCallback(hwc2_callback_data_t data,
+                               hwc2_function_pointer_t func);
     void RegisterRefreshCallback(hwc2_callback_data_t data,
                                  hwc2_function_pointer_t func);
     HWC2::Error CreateComposition(bool test);
@@ -422,7 +425,6 @@
 
   ResourceManager resource_manager_;
   std::map<hwc2_display_t, HwcDisplay> displays_;
-  std::map<HWC2::Callback, HwcCallback> callbacks_;
 
   std::string mDumpString;
 };
diff --git a/compositor/DrmDisplayCompositor.cpp b/compositor/DrmDisplayCompositor.cpp
index 467f8ba..ba0d56b 100644
--- a/compositor/DrmDisplayCompositor.cpp
+++ b/compositor/DrmDisplayCompositor.cpp
@@ -813,7 +813,9 @@
 }
 
 int DrmDisplayCompositor::FlattenOnClient() {
-  if (refresh_display_cb_) {
+  const std::lock_guard<std::mutex> lock(refresh_callback_lock);
+
+  if (refresh_callback_hook_ && refresh_callback_data_) {
     {
       AutoLock lock(&lock_, __func__);
       if (!IsFlatteningNeeded()) {
@@ -829,7 +831,7 @@
         "No writeback connector available, "
         "falling back to client composition");
     SetFlattening(FlatteningState::kClientRequested);
-    refresh_display_cb_(display_);
+    refresh_callback_hook_(refresh_callback_data_, display_);
     return 0;
   } else {
     ALOGV("No writeback connector available");
diff --git a/compositor/DrmDisplayCompositor.h b/compositor/DrmDisplayCompositor.h
index e500c9e..29afc66 100644
--- a/compositor/DrmDisplayCompositor.h
+++ b/compositor/DrmDisplayCompositor.h
@@ -59,9 +59,15 @@
 
   int Init(ResourceManager *resource_manager, int display);
 
-  template <typename Fn>
-  void SetRefreshCallback(Fn &&refresh_cb) {
-    refresh_display_cb_ = std::forward<Fn>(refresh_cb);
+  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);
   }
 
   std::unique_ptr<DrmDisplayComposition> CreateComposition() const;
diff --git a/drm/VSyncWorker.cpp b/drm/VSyncWorker.cpp
index dfbf8ce..b2f7e5f 100644
--- a/drm/VSyncWorker.cpp
+++ b/drm/VSyncWorker.cpp
@@ -50,6 +50,14 @@
   Unlock();
 }
 
+void VSyncWorker::RegisterClientCallback(hwc2_callback_data_t data,
+                                         hwc2_function_pointer_t hook) {
+  Lock();
+  vsync_callback_data_ = data;
+  vsync_callback_hook_ = reinterpret_cast<HWC2_PFN_VSYNC>(hook);
+  Unlock();
+}
+
 void VSyncWorker::VSyncControl(bool enabled) {
   Lock();
   enabled_ = enabled;
@@ -151,37 +159,17 @@
                 (int64_t)vblank.reply.tval_usec * 1000;
   }
 
-  /*
-   * VSync could be disabled during routine execution so it could potentially
-   * lead to crash since callback's inner hook could be invalid anymore. We have
-   * no control over lifetime of this hook, therefore we can't rely that it'll
-   * be valid after vsync disabling.
-   *
-   * Blocking VSyncControl to wait until routine
-   * will finish execution is logically correct way to fix this issue, but it
-   * creates visible lags and stutters, so we have to resort to other ways of
-   * mitigating this issue.
-   *
-   * Doing check before attempt to invoke callback drastically shortens the
-   * window when such situation could happen and that allows us to practically
-   * avoid this issue.
-   *
-   * Please note that issue described below is different one and it is related
-   * to RegisterCallback, not to disabling vsync via VSyncControl.
-   */
   if (!enabled_)
     return;
-  /*
-   * There's a race here where a change in callback_ will not take effect until
-   * the next subsequent requested vsync. This is unavoidable since we can't
-   * call the vsync hook while holding the thread lock.
-   *
-   * We could shorten the race window by caching callback_ right before calling
-   * the hook. However, in practice, callback_ is only updated once, so it's not
-   * worth the overhead.
-   */
+
   if (callback)
     callback->Callback(display, timestamp);
+
+  Lock();
+  if (enabled_ && vsync_callback_hook_ && vsync_callback_data_)
+    vsync_callback_hook_(vsync_callback_data_, display, timestamp);
+  Unlock();
+
   last_timestamp_ = timestamp;
 }
 }  // namespace android
diff --git a/drm/VSyncWorker.h b/drm/VSyncWorker.h
index 0121a6c..7454b51 100644
--- a/drm/VSyncWorker.h
+++ b/drm/VSyncWorker.h
@@ -19,6 +19,7 @@
 
 #include <hardware/hardware.h>
 #include <hardware/hwcomposer.h>
+#include <hardware/hwcomposer2.h>
 #include <stdint.h>
 
 #include <map>
@@ -42,6 +43,8 @@
 
   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);
 
   void VSyncControl(bool enabled);
 
@@ -62,6 +65,9 @@
   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