diff --git a/vsyncworker.h b/vsyncworker.h
index 98ac546..a1ba1a5 100644
--- a/vsyncworker.h
+++ b/vsyncworker.h
@@ -28,13 +28,20 @@
 
 namespace android {
 
+class VsyncCallback {
+ public:
+  virtual ~VsyncCallback() {
+  }
+  virtual void Callback(int display, int64_t timestamp) = 0;
+};
+
 class VSyncWorker : public Worker {
  public:
   VSyncWorker();
   ~VSyncWorker() override;
 
   int Init(DrmResources *drm, int display);
-  int SetProcs(hwc_procs_t const *procs);
+  int RegisterCallback(std::shared_ptr<VsyncCallback> callback);
 
   int VSyncControl(bool enabled);
 
@@ -46,7 +53,11 @@
   int SyntheticWaitVBlank(int64_t *timestamp);
 
   DrmResources *drm_;
-  hwc_procs_t const *procs_;
+
+  // 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;
 
   int display_;
   bool enabled_;