1 files changed, 10 insertions, 0 deletions
diff --git a/lib/xray/xray_log_interface.cc b/lib/xray/xray_log_interface.cc
index ffed601c0..ee14ae4b1 100644
--- a/lib/xray/xray_log_interface.cc
+++ b/lib/xray/xray_log_interface.cc
@@ -27,12 +27,22 @@ void __xray_set_log_impl(XRayLogImpl Impl) XRAY_NEVER_INSTRUMENT {
       Impl.handle_arg0 == nullptr || Impl.flush_log == nullptr) {
     __sanitizer::SpinMutexLock Guard(&XRayImplMutex);
     GlobalXRayImpl.reset();
+    __xray_remove_handler();
+    __xray_remove_handler_arg1();
     return;
   }
 
   __sanitizer::SpinMutexLock Guard(&XRayImplMutex);
   GlobalXRayImpl.reset(new XRayLogImpl);
   *GlobalXRayImpl = Impl;
+  __xray_set_handler(Impl.handle_arg0);
+}
+
+void __xray_remove_log_impl() XRAY_NEVER_INSTRUMENT {
+  __sanitizer::SpinMutexLock Guard(&XRayImplMutex);
+  GlobalXRayImpl.reset();
+  __xray_remove_handler();
+  __xray_remove_handler_arg1();
 }
 
 XRayLogInitStatus __xray_log_init(size_t BufferSize, size_t MaxBuffers,