diff --git a/drivers/gpu/nvgpu/include/nvgpu/posix/bug.h b/drivers/gpu/nvgpu/include/nvgpu/posix/bug.h
index 10518652a..e81297b2d 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/posix/bug.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/posix/bug.h
@@ -94,25 +94,7 @@ void nvgpu_posix_bug(const char *fmt, ...) __attribute__ ((noreturn));
 bool nvgpu_posix_warn(bool cond, const char *fmt, ...);
 
 #ifdef __NVGPU_UNIT_TEST__
-/* Provide a simple API for BUG() handling */
-
-/**
- * @brief Bug handler register.
- *
- * @param handler [in]	Bug handler
- *
- * Registers a handler for handling bug.  Used in unit testing.
- */
-void bug_handler_register(jmp_buf *handler);
-
-/**
- * @brief Cancels the bug handler.
- *
- * @param void.
- *
- * Cancels the handler for handling bug.  Used in unit testing.
- */
-void bug_handler_cancel(void);
+void nvgpu_bug_cb_longjmp(void *arg);
 
 /**
  * Macro to indicate that a BUG() call is expected when executing
@@ -120,18 +102,23 @@ void bug_handler_cancel(void);
  * and the setjmp API to set a long jump point that gets called by the BUG()
  * function if enabled. This allows the macro to simply expand as true if
  * BUG() was called, and false otherwise.
+ * Note: it is safe to call nvgpu_bug_unregister_cb for a callback
+ * that was already invoked/unregistered.
  */
-#define EXPECT_BUG(code_to_run)				\
-	({						\
-		jmp_buf handler;			\
-		volatile bool bug_result = true;	\
-		if (setjmp(handler) == 0) {		\
-			bug_handler_register(&handler);	\
-			code_to_run;			\
-			bug_handler_cancel();		\
-			bug_result = false;		\
-		}					\
-		bug_result;				\
+#define EXPECT_BUG(code_to_run)					\
+	({							\
+		jmp_buf handler;				\
+		volatile bool bug_result = true;		\
+		struct nvgpu_bug_cb callback;			\
+		callback.cb = nvgpu_bug_cb_longjmp;		\
+		callback.arg = &handler;			\
+		nvgpu_bug_register_cb(&callback);		\
+		if (setjmp(handler) == 0) {			\
+			code_to_run;				\
+			bug_result = false;			\
+		}						\
+		nvgpu_bug_unregister_cb(&callback);		\
+		bug_result;					\
 	})
 #endif
 
diff --git a/drivers/gpu/nvgpu/libnvgpu-drv_safe.export b/drivers/gpu/nvgpu/libnvgpu-drv_safe.export
index fd30f94d4..84a33e469 100644
--- a/drivers/gpu/nvgpu/libnvgpu-drv_safe.export
+++ b/drivers/gpu/nvgpu/libnvgpu-drv_safe.export
@@ -1,8 +1,6 @@
 # Copyright (c) 2019-2020, NVIDIA CORPORATION.  All rights reserved.
 
 bitmap_find_next_zero_area_off
-bug_handler_cancel
-bug_handler_register
 fb_gv11b_write_mmu_fault_buffer_get
 find_first_bit
 find_first_zero_bit
@@ -272,6 +270,7 @@ nvgpu_big_pages_possible
 nvgpu_bitmap_clear
 nvgpu_bitmap_set
 nvgpu_bsearch
+nvgpu_bug_cb_longjmp
 nvgpu_bug_register_cb
 nvgpu_bug_unregister_cb
 nvgpu_can_busy
diff --git a/drivers/gpu/nvgpu/os/posix/bug.c b/drivers/gpu/nvgpu/os/posix/bug.c
index 61d490d26..4c08debdd 100644
--- a/drivers/gpu/nvgpu/os/posix/bug.c
+++ b/drivers/gpu/nvgpu/os/posix/bug.c
@@ -37,20 +37,24 @@
 
 #define BACKTRACE_MAXSIZE 1024
 
+struct nvgpu_bug_desc {
+	bool in_use;
+	pthread_once_t once;
+	struct nvgpu_spinlock lock;
+	struct nvgpu_list_node head;
+};
+
+struct nvgpu_bug_desc bug = {
+	.once = PTHREAD_ONCE_INIT
+};
+
 #ifdef __NVGPU_UNIT_TEST__
-static _Thread_local bool expect_bug;
-static _Thread_local jmp_buf *jmp_handler;
-
-void bug_handler_register(jmp_buf *handler)
+void nvgpu_bug_cb_longjmp(void *arg)
 {
-	expect_bug = true;
-	jmp_handler = handler;
-}
+	nvgpu_info(NULL, "Expected BUG detected!");
 
-void bug_handler_cancel(void)
-{
-	expect_bug = false;
-	jmp_handler = NULL;
+	jmp_buf *jmp_handler = arg;
+	longjmp(*jmp_handler, 1);
 }
 #endif
 
@@ -77,17 +81,6 @@ void dump_stack(void)
 	nvgpu_posix_dump_stack(2);
 }
 
-struct nvgpu_bug_desc {
-	bool in_use;
-	pthread_once_t once;
-	struct nvgpu_spinlock lock;
-	struct nvgpu_list_node head;
-};
-
-struct nvgpu_bug_desc bug = {
-	.once = PTHREAD_ONCE_INIT
-};
-
 static void nvgpu_bug_init(void)
 {
 	nvgpu_err(NULL, "doing init for bug cb");
@@ -119,14 +112,6 @@ void nvgpu_posix_bug(const char *fmt, ...)
 {
 	struct nvgpu_bug_cb *cb;
 
-#ifdef __NVGPU_UNIT_TEST__
-	if (expect_bug) {
-		nvgpu_info(NULL, "Expected BUG detected!");
-		expect_bug = false;
-		/* Perform a long jump to where "setjmp()" was called. */
-		longjmp(*jmp_handler, 1);
-	}
-#endif
 	/*
 	 * If BUG was unexpected, raise a SIGSEGV signal, dump the stack and
 	 * kill the thread.