Merge pull request #20891 from JasonFengJ9/debugonrestore

CRIU DebugOnRestore mode is disabled when debug events are hooked
eclipse-openj9 · Feb 19, 2025 · e2516e2 · e2516e2
2 parents 8fd2473 + ab9186f
commit e2516e2
Show file tree

Hide file tree

Showing 9 changed files with 190 additions and 101 deletions.
diff --git a/runtime/jvmti/j9jvmti.tdf b/runtime/jvmti/j9jvmti.tdf
@@ -694,7 +694,11 @@ TraceEvent=Trc_JVMTI_criuAddCapabilities_invoked Noenv Overhead=1 Level=3 Templa
 TraceEvent=Trc_JVMTI_createAgentLibraryWithOption_OOM Noenv Overhead=1 Level=3 Template="createAgentLibraryWithOption j9mem_allocate_memory OOM"
 TraceEvent=Trc_JVMTI_createAgentLibraryWithOption_Xrunjdwp_result Noenv Overhead=1 Level=3 Template="createAgentLibraryWithOption Xrunjdwp optionsPtr (%s) optionsLengthTmp (%zu) agentLibrary (%p) result (%d)"
 TraceEvent=Trc_JVMTI_createAgentLibraryWithOption_agentlib_result Noenv Overhead=1 Level=3 Template="createAgentLibraryWithOption optionsPtr (%s) libraryLength (%zu) options (%s) optionsLength (%zu) agentLibrary (%p) result (%d)"
-TraceEntry=Trc_JVMTI_jvmtiHookVMRestoreCRIUInit_Entry Overhead=1 Level=3 Noenv Template="jvmtiHookVMRestoreCRIUInit"
-TraceExit=Trc_JVMTI_jvmtiHookVMRestoreCRIUInit_Exit Overhead=1 Level=3 Noenv Template="jvmtiHookVMRestoreCRIUInit"
-TraceEntry=Trc_JVMTI_jvmtiHookVMRestoreStartAgent_Entry Overhead=1 Level=3 Noenv Template="jvmtiHookVMRestoreStartAgent"
-TraceExit=Trc_JVMTI_jvmtiHookVMRestoreStartAgent_Exit Overhead=1 Level=3 Noenv Template="jvmtiHookVMRestoreStartAgent"
+TraceEntry=Trc_JVMTI_jvmtiHookVMRestoreCRIUInit_Entry Obsolete Overhead=1 Level=3 Noenv Template="jvmtiHookVMRestoreCRIUInit"
+TraceExit=Trc_JVMTI_jvmtiHookVMRestoreCRIUInit_Exit Obsolete Overhead=1 Level=3 Noenv Template="jvmtiHookVMRestoreCRIUInit"
+TraceEntry=Trc_JVMTI_jvmtiHookVMRestoreStartAgent_Entry Obsolete Overhead=1 Level=3 Noenv Template="jvmtiHookVMRestoreStartAgent"
+TraceExit=Trc_JVMTI_jvmtiHookVMRestoreStartAgent_Exit Obsolete Overhead=1 Level=3 Noenv Template="jvmtiHookVMRestoreStartAgent"
+TraceEntry=Trc_JVMTI_jvmtiHookVMPreparingForRestore_Entry Overhead=1 Level=3 Noenv Template="jvmtiHookVMPreparingForRestore"
+TraceExit=Trc_JVMTI_jvmtiHookVMPreparingForRestore_Exit Overhead=1 Level=3 Noenv Template="jvmtiHookVMPreparingForRestore"
+TraceEntry=Trc_JVMTI_jvmtiHookVMCRIURestore_Entry Overhead=1 Level=3 Noenv Template="jvmtiHookVMCRIURestore"
+TraceExit=Trc_JVMTI_jvmtiHookVMCRIURestore_Exit Overhead=1 Level=3 Noenv Template="jvmtiHookVMCRIURestore"
diff --git a/runtime/jvmti/jvmtiCapability.c b/runtime/jvmti/jvmtiCapability.c
@@ -115,7 +115,8 @@ jvmtiGetPotentialCapabilities(jvmtiEnv* env, jvmtiCapabilities* capabilities_ptr
 	J9JavaVM * vm = j9env->vm;
 	J9JVMTIData * jvmtiData = J9JVMTI_DATA_FROM_VM(vm);
 	jvmtiCapabilities rv_capabilities;
-	J9HookInterface ** vmHook = vm->internalVMFunctions->getVMHookInterface(vm);
+	J9InternalVMFunctions *vmFuncs = vm->internalVMFunctions;
+	J9HookInterface **vmHook = vmFuncs->getVMHookInterface(vm);
 	jvmtiError rc;
 
 	Trc_JVMTI_jvmtiGetPotentialCapabilities_Entry(env);
@@ -206,8 +207,13 @@ jvmtiGetPotentialCapabilities(jvmtiEnv* env, jvmtiCapabilities* capabilities_ptr
 	}
 
 	if ((*vmHook)->J9HookIsEnabled(vmHook, J9HOOK_VM_POP_FRAMES_INTERRUPT)) {
-		rv_capabilities.can_pop_frame = 1;
-		rv_capabilities.can_force_early_return = 1;
+#if defined(J9VM_OPT_CRIU_SUPPORT)
+		if (!vmFuncs->isDebugAgentDisabled(vm))
+#endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
+		{
+			rv_capabilities.can_pop_frame = 1;
+			rv_capabilities.can_force_early_return = 1;
+		}
 	}
 
 	if ((*vmHook)->J9HookIsEnabled(vmHook, J9HOOK_VM_REQUIRED_DEBUG_ATTRIBUTES) ||
@@ -231,7 +237,12 @@ jvmtiGetPotentialCapabilities(jvmtiEnv* env, jvmtiCapabilities* capabilities_ptr
 	if ((*vmHook)->J9HookIsEnabled(vmHook, J9HOOK_VM_REQUIRED_DEBUG_ATTRIBUTES) ||
 		(vm->requiredDebugAttributes & J9VM_DEBUG_ATTRIBUTE_CAN_ACCESS_LOCALS)
 	) {
-		rv_capabilities.can_access_local_variables = 1;
+#if defined(J9VM_OPT_CRIU_SUPPORT)
+		if (!vmFuncs->isDebugAgentDisabled(vm))
+#endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
+		{
+			rv_capabilities.can_access_local_variables = 1;
+		}
 	}
 
 	rv_capabilities.can_tag_objects = 1;
@@ -548,16 +559,13 @@ mapCapabilitiesToEvents(J9JVMTIEnv * j9env, jvmtiCapabilities * capabilities, J9
 
 #if defined(J9VM_OPT_CRIU_SUPPORT)
 	J9JavaVM *vm = j9env->vm;
-	J9InternalVMFunctions *vmFuncs = vm->internalVMFunctions;
-	BOOLEAN skipHookReserve = vmFuncs->isCheckpointAllowed(vm)
-			&& vmFuncs->isDebugOnRestoreEnabled(vm);
-	/* Skip J9HookReserve for the events required by JDWP agent pre-checkpoint when DebugOnRestore is enabled,
+	/* Skip J9HookReserve for the events required by JDWP agent pre-checkpoint when isDebugAgentDisabled() returns TRUE,
 	 * these events will be registered post-restore if a JDWP agent is specified in the restore option file,
 	 * otherwise they are going to be unregistered by J9HookUnregister() which only clears J9HOOK_FLAG_HOOKED,
 	 * but not J9HOOK_FLAG_RESERVED.
 	 * J9HookUnreserve() might clear the flag set by other callers.
 	 */
-	if (!skipHookReserve)
+	if (!vm->internalVMFunctions->isDebugAgentDisabled(vm))
 #endif /* defined(J9VM_OPT_CRIU_SUPPORT)*/
 	{
 		if (capabilities->can_generate_single_step_events) {

diff --git a/runtime/jvmti/jvmtiHook.c b/runtime/jvmti/jvmtiHook.c
@@ -155,8 +155,6 @@ static BOOLEAN shouldPostEvent(J9VMThread *currentThread, J9Method *method);
 #if defined(J9VM_OPT_CRIU_SUPPORT)
 static void jvmtiHookVMCheckpoint(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData);
 static void jvmtiHookVMRestore(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData);
-static void jvmtiHookVMRestoreCRIUInit(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData);
-static void jvmtiHookVMRestoreStartAgent(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData);
 static void hookDisableHelper(J9JavaVM *vm, J9HookInterface **vmHook, UDATA eventNum, J9HookFunction function, BOOLEAN unreserve, void *userData);
 
 static void
@@ -557,37 +555,6 @@ jvmtiHookVMCheckpoint(J9HookInterface **hook, UDATA eventNum, void *eventData, v
 	TRACE_JVMTI_EVENT_RETURN(jvmtiHookVMCheckpoint);
 }
 
-static void
-jvmtiHookVMRestoreCRIUInit(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData)
-{
-	Trc_JVMTI_jvmtiHookVMRestoreCRIUInit_Entry();
-	criuRestoreInitializeLib(((J9RestoreEvent *)eventData)->currentThread->javaVM, (J9JVMTIEnv *)userData);
-	TRACE_JVMTI_EVENT_RETURN(jvmtiHookVMRestoreCRIUInit);
-}
-
-static void
-jvmtiHookVMRestoreStartAgent(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData)
-{
-	J9VMThread *currentThread = ((J9RestoreEvent *)eventData)->currentThread;
-	J9JavaVM *vm = currentThread->javaVM;
-	J9InternalVMFunctions const * const vmFuncs = vm->internalVMFunctions;
-
-	Trc_JVMTI_jvmtiHookVMRestoreStartAgent_Entry();
-	vmFuncs->internalExitVMToJNI(currentThread);
-	if (J9_ARE_ANY_BITS_SET(vm->checkpointState.flags, J9VM_CRIU_IS_JDWP_ENABLED)) {
-		criuRestoreStartAgent(vm);
-	} else {
-		/* Last part of cleanup if there was no JDWP agent specified.
-		 * This releases VM access hence can't be invoked within criuDisableHooks() from
-		 * J9HOOK_VM_PREPARING_FOR_RESTORE.
-		 */
-		jvmtiEnv *jvmti_env = vm->checkpointState.jvmtienv;
-		(*jvmti_env)->DisposeEnvironment(jvmti_env);
-	}
-	vmFuncs->internalEnterVMFromJNI(currentThread);
-	TRACE_JVMTI_EVENT_RETURN(jvmtiHookVMRestoreStartAgent);
-}
-
 static void
 jvmtiHookVMRestore(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData)
 {
@@ -1075,6 +1042,21 @@ jvmtiHookClassLoad(J9HookInterface** hook, UDATA eventNum, void* eventData, void
 UDATA
 isEventHookable(J9JVMTIEnv * j9env, jvmtiEvent event)
 {
+#if defined(J9VM_OPT_CRIU_SUPPORT)
+	J9JavaVM *vm = j9env->vm;
+	if (vm->internalVMFunctions->isDebugAgentDisabled(vm)) {
+		switch(event) {
+			case JVMTI_EVENT_FRAME_POP: /* fall through */
+			case JVMTI_EVENT_FIELD_ACCESS: /* fall through */
+			case JVMTI_EVENT_FIELD_MODIFICATION: /* fall through */
+			case JVMTI_EVENT_SINGLE_STEP: /* fall through */
+			case JVMTI_EVENT_BREAKPOINT:
+				return FALSE;
+			default:
+				break;
+		}
+	}
+#endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
 	return processEvent(j9env, event, hookIsDisabled) == 0;
 }
 
@@ -1136,6 +1118,18 @@ hookNonEventCapabilities(J9JVMTIEnv * j9env, jvmtiCapabilities * capabilities)
 	J9JVMTIHookInterfaceWithID * gcOmrHook = &j9env->gcOmrHook;
 	J9JVMTIData * jvmtiData = J9JVMTI_DATA_FROM_VM(vm);
 
+#if defined(J9VM_OPT_CRIU_SUPPORT)
+	if (vm->internalVMFunctions->isDebugAgentDisabled(vm)) {
+		if (capabilities->can_pop_frame
+			|| capabilities->can_force_early_return
+			|| capabilities->can_access_local_variables
+			|| capabilities->can_generate_breakpoint_events
+		) {
+			return 1;
+		}
+	}
+#endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
+
 	if (capabilities->can_generate_breakpoint_events) {
 		if (hookRegister(vmHook, J9HOOK_VM_BREAKPOINT, jvmtiHookBreakpoint, OMR_GET_CALLSITE(), j9env)) {
 			return 1;
@@ -2046,17 +2040,6 @@ hookGlobalEvents(J9JVMTIData * jvmtiData)
 		return 1;
 	}
 
-#if defined(J9VM_OPT_CRIU_SUPPORT)
-	if (vm->internalVMFunctions->isDebugOnRestoreEnabled(vm)) {
-		if ((*vmHook)->J9HookRegisterWithCallSite(vmHook, J9HOOK_TAG_AGENT_ID | J9HOOK_VM_PREPARING_FOR_RESTORE, jvmtiHookVMRestoreCRIUInit, OMR_GET_CALLSITE(), jvmtiData, J9HOOK_AGENTID_FIRST)) {
-			return 1;
-		}
-		if ((*vmHook)->J9HookRegisterWithCallSite(vmHook, J9HOOK_TAG_AGENT_ID | J9HOOK_VM_CRIU_RESTORE, jvmtiHookVMRestoreStartAgent, OMR_GET_CALLSITE(), jvmtiData, J9HOOK_AGENTID_FIRST)) {
-			return 1;
-		}
-	}
-#endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
-
 	if ((*vmHook)->J9HookRegisterWithCallSite(vmHook, J9HOOK_TAG_AGENT_ID | J9HOOK_VM_SHUTTING_DOWN, jvmtiHookVMShutdownLast, OMR_GET_CALLSITE(), jvmtiData, J9HOOK_AGENTID_LAST)) {
 		return 1;
 	}

diff --git a/runtime/jvmti/jvmtiStartup.c b/runtime/jvmti/jvmtiStartup.c
@@ -65,6 +65,10 @@ I_32 JNICALL loadAgentLibraryOnAttach(struct J9JavaVM *vm, const char *library,
 static BOOLEAN isAgentLibraryLoaded(J9JavaVM *vm, const char *library, BOOLEAN decorate);
 static jint createAgentLibraryWithOption(J9JavaVM *vm, J9VMInitArgs *argsList, IDATA agentIndex, J9JVMTIAgentLibrary **agentLibrary, CreateAgentOption createAgentOption, BOOLEAN *isJDWPagent);
 static BOOLEAN processAgentLibraryFromArgsList(J9JavaVM *vm, J9VMInitArgs *argsList, BOOLEAN loadLibrary, CreateAgentOption createAgentOption);
+#if defined(J9VM_OPT_CRIU_SUPPORT)
+static void jvmtiHookVMPreparingForRestore(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData);
+static void jvmtiHookVMCRIURestore(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData);
+#endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
 
 #define INSTRUMENT_LIBRARY "instrument"
 
@@ -267,6 +271,51 @@ processAgentLibraryFromArgsList(J9JavaVM *vm, J9VMInitArgs *argsList, BOOLEAN lo
 }
 
 #if defined(J9VM_OPT_CRIU_SUPPORT)
+/**
+ * A hook method to invoke criuRestoreInitializeLib().
+ *
+ * @param[in] hook the VM hook interface, not used
+ * @param[in] eventNum the event number, not used
+ * @param[in] eventData the event data, not used
+ * @param[in] userData the registered user data
+ */
+static void
+jvmtiHookVMPreparingForRestore(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData)
+{
+	Trc_JVMTI_jvmtiHookVMPreparingForRestore_Entry();
+	criuRestoreInitializeLib(((J9RestoreEvent *)eventData)->currentThread->javaVM, (J9JVMTIEnv *)userData);
+	Trc_JVMTI_jvmtiHookVMPreparingForRestore_Exit();
+}
+
+/**
+ * A hook method to cleanup post-restore.
+ *
+ * @param[in] hook the VM hook interface, not used
+ * @param[in] eventNum the event number, not used
+ * @param[in] eventData the event data
+ * @param[in] userData the registered user data, not used
+ */
+static void
+jvmtiHookVMCRIURestore(J9HookInterface **hook, UDATA eventNum, void *eventData, void *userData)
+{
+	J9VMThread *currentThread = ((J9RestoreEvent *)eventData)->currentThread;
+	J9JavaVM *vm = currentThread->javaVM;
+	J9InternalVMFunctions const * const vmFuncs = vm->internalVMFunctions;
+
+	Trc_JVMTI_jvmtiHookVMCRIURestore_Entry();
+	vmFuncs->internalExitVMToJNI(currentThread);
+	if (J9_ARE_NO_BITS_SET(vm->checkpointState.flags, J9VM_CRIU_IS_JDWP_ENABLED)) {
+		/* Last part of cleanup if there was no JDWP agent specified.
+		 * This releases VM access hence can't be invoked within criuDisableHooks() from
+		 * J9HOOK_VM_PREPARING_FOR_RESTORE.
+		 */
+		jvmtiEnv *jvmti_env = vm->checkpointState.jvmtienv;
+		(*jvmti_env)->DisposeEnvironment(jvmti_env);
+	}
+	vmFuncs->internalEnterVMFromJNI(currentThread);
+	Trc_JVMTI_jvmtiHookVMCRIURestore_Exit();
+}
+
 /**
  * Add JVMTI capabilities before checkpoint.
  * This is required for debugger support when JIT is enabled.
@@ -301,9 +350,7 @@ criuAddCapabilities(J9JavaVM *vm, BOOLEAN jitEnabled) {
 		requiredCapabilities->can_get_line_numbers = 1;
 		requiredCapabilities->can_get_source_debug_extension = 1;
 		requiredCapabilities->can_maintain_original_method_order = 1;
-		requiredCapabilities->can_generate_single_step_events = 1;
 		requiredCapabilities->can_generate_exception_events = 1;
-		requiredCapabilities->can_generate_frame_pop_events = 1;
 		requiredCapabilities->can_generate_breakpoint_events = 1;
 		requiredCapabilities->can_generate_method_entry_events = 1;
 		requiredCapabilities->can_generate_method_exit_events = 1;
@@ -317,8 +364,6 @@ criuAddCapabilities(J9JavaVM *vm, BOOLEAN jitEnabled) {
 		if (JVMTI_ERROR_NONE != jvmtiRet) {
 			return JNI_ERR;
 		}
-		requiredCapabilities->can_generate_field_modification_events = potentialCapabilities.can_generate_field_modification_events;
-		requiredCapabilities->can_generate_field_access_events = potentialCapabilities.can_generate_field_access_events;
 		requiredCapabilities->can_pop_frame = potentialCapabilities.can_pop_frame;
 	}
 	jvmtiRet = (*jvmti_env)->AddCapabilities(jvmti_env, requiredCapabilities);
@@ -334,10 +379,9 @@ void
 criuRestoreInitializeLib(J9JavaVM *vm, J9JVMTIEnv *j9env)
 {
 	J9VMInitArgs *criuRestoreArgsList = vm->checkpointState.restoreArgsList;
-
-	processAgentLibraryFromArgsList(vm, criuRestoreArgsList, FALSE, OPTION_AGENTLIB);
-	processAgentLibraryFromArgsList(vm, criuRestoreArgsList, FALSE, OPTION_AGENTPATH);
-	processAgentLibraryFromArgsList(vm, criuRestoreArgsList, FALSE, OPTION_XRUNJDWP);
+	processAgentLibraryFromArgsList(vm, criuRestoreArgsList, TRUE, OPTION_AGENTLIB);
+	processAgentLibraryFromArgsList(vm, criuRestoreArgsList, TRUE, OPTION_AGENTPATH);
+	processAgentLibraryFromArgsList(vm, criuRestoreArgsList, TRUE, OPTION_XRUNJDWP);
 
 	if (J9_ARE_NO_BITS_SET(vm->checkpointState.flags, J9VM_CRIU_IS_JDWP_ENABLED)) {
 		J9JVMTIData * jvmtiData = vm->jvmtiData;
@@ -346,16 +390,6 @@ criuRestoreInitializeLib(J9JavaVM *vm, J9JVMTIEnv *j9env)
 		}
 	}
 }
-
-void
-criuRestoreStartAgent(J9JavaVM *vm)
-{
-	J9VMInitArgs *criuRestoreArgsList = vm->checkpointState.restoreArgsList;
-
-	processAgentLibraryFromArgsList(vm, criuRestoreArgsList, TRUE, OPTION_AGENTLIB);
-	processAgentLibraryFromArgsList(vm, criuRestoreArgsList, TRUE, OPTION_AGENTPATH);
-	processAgentLibraryFromArgsList(vm, criuRestoreArgsList, TRUE, OPTION_XRUNJDWP);
-}
 #endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
 
 IDATA J9VMDllMain(J9JavaVM *vm, IDATA stage, void *reserved)
@@ -418,14 +452,45 @@ IDATA J9VMDllMain(J9JavaVM *vm, IDATA stage, void *reserved)
 			hshelpUTRegister(vm);
 
 #if defined(J9VM_OPT_CRIU_SUPPORT)
-			/*
-			 * Adding capabilities is required before checkpoint if JIT is enabled.
-			 * Following code can be removed when JIT allows capabilities to be added after restore.
-			 */
-			if (vm->internalVMFunctions->isDebugOnRestoreEnabled(vm)) {
-				Trc_JVMTI_criuAddCapabilities_invoked();
-				/* ignore the failure, it won't cause a problem if JDWP is not enabled later */
-				criuAddCapabilities(vm, NULL != vm->jitConfig);
+			{
+				/* The isDebugEventOrFlagEnabled calculation matches a part of J9::Options::isFSDNeeded()
+				 * in compiler/control/J9Options.cpp.
+				 */
+				BOOLEAN isDebugEventOrFlagEnabled = J9_EVENT_IS_HOOKED_OR_RESERVED(vm->hookInterface, J9HOOK_VM_BREAKPOINT)
+						|| J9_EVENT_IS_HOOKED_OR_RESERVED(vm->hookInterface, J9HOOK_VM_FRAME_POP)
+						|| J9_EVENT_IS_HOOKED_OR_RESERVED(vm->hookInterface, J9HOOK_VM_FRAME_POPPED)
+						|| J9_EVENT_IS_HOOKED_OR_RESERVED(vm->hookInterface, J9HOOK_VM_GET_FIELD)
+						|| J9_EVENT_IS_HOOKED_OR_RESERVED(vm->hookInterface, J9HOOK_VM_PUT_FIELD)
+						|| J9_EVENT_IS_HOOKED_OR_RESERVED(vm->hookInterface, J9HOOK_VM_GET_STATIC_FIELD)
+						|| J9_EVENT_IS_HOOKED_OR_RESERVED(vm->hookInterface, J9HOOK_VM_PUT_STATIC_FIELD)
+#if defined (J9VM_INTERP_HOT_CODE_REPLACEMENT)
+						|| J9_EVENT_IS_HOOKED_OR_RESERVED(vm->hookInterface, J9HOOK_VM_POP_FRAMES_INTERRUPT)
+#endif
+#if defined(J9VM_JIT_FULL_SPEED_DEBUG)
+						|| (vm->requiredDebugAttributes & J9VM_DEBUG_ATTRIBUTE_CAN_ACCESS_LOCALS)
+#endif
+						|| J9_EVENT_IS_HOOKED_OR_RESERVED(vm->hookInterface, J9HOOK_VM_SINGLE_STEP);
+				J9InternalVMFunctions *vmFuncs = vm->internalVMFunctions;
+				BOOLEAN isDebugOnRestoreEnabled = !isDebugEventOrFlagEnabled
+						&& J9_ARE_ALL_BITS_SET(vm->checkpointState.flags, J9VM_CRIU_SUPPORT_DEBUG_ON_RESTORE)
+						&& vmFuncs->isCRaCorCRIUSupportEnabled(vm);
+
+				if (isDebugOnRestoreEnabled) {
+					J9HookInterface ** vmHook = vmFuncs->getVMHookInterface(vm);
+					if ((*vmHook)->J9HookRegisterWithCallSite(vmHook, J9HOOK_TAG_AGENT_ID | J9HOOK_VM_PREPARING_FOR_RESTORE, jvmtiHookVMPreparingForRestore, OMR_GET_CALLSITE(), jvmtiData, J9HOOK_AGENTID_FIRST)) {
+						goto _error;
+					}
+					if ((*vmHook)->J9HookRegisterWithCallSite(vmHook, J9HOOK_TAG_AGENT_ID | J9HOOK_VM_CRIU_RESTORE, jvmtiHookVMCRIURestore, OMR_GET_CALLSITE(), jvmtiData, J9HOOK_AGENTID_FIRST)) {
+						goto _error;
+					}
+					/* Adding capabilities is required before checkpoint if JIT is enabled.
+					 * Following code can be removed when JIT allows capabilities to be added after restore.
+					 */
+					Trc_JVMTI_criuAddCapabilities_invoked();
+					/* ignore the failure, it won't cause a problem if JDWP is not enabled later */
+					criuAddCapabilities(vm, NULL != vm->jitConfig);
+					vm->checkpointState.isDebugOnRestoreEnabled = isDebugOnRestoreEnabled;
+				}
 			}
 #endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
 

diff --git a/runtime/jvmti/jvmti_internal.h b/runtime/jvmti/jvmti_internal.h
@@ -2209,23 +2209,15 @@ jint JNICALL JVM_OnLoad(JavaVM *jvm, char* options, void *reserved);
 #if defined(J9VM_OPT_CRIU_SUPPORT)
 /**
 * Check if any agent library is specified in the CRIU restore option file,
-* if so, the checkpointState flag J9VM_CRIU_IS_JDWP_ENABLED is set, but the
-* actual library loading is deferred at criuRestoreStartAgent() invoked via
-* TRIGGER_J9HOOK_VM_CRIU_RESTORE().
+* if so, the checkpointState flag J9VM_CRIU_IS_JDWP_ENABLED is set, and
+* load/start the library loading.
+*
 * @param[in] vm the pointer to the J9JavaVM struct
 * @param[in] j9env - pointer to the current JVMTI environment
 * @return void
 */
 void
 criuRestoreInitializeLib(J9JavaVM *vm, J9JVMTIEnv *j9env);
-
-/**
- * Load and initialize libraries from CRIU restore option file.
- * @param[in] vm the pointer to the J9JavaVM struct
- * @return void
- */
-void
-criuRestoreStartAgent(J9JavaVM *vm);
 #endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
 
 /* ---------------- jvmtiSystemProperties.c ---------------- */

diff --git a/runtime/oti/j9nonbuilder.h b/runtime/oti/j9nonbuilder.h
@@ -4462,6 +4462,7 @@ typedef struct J9CRIUCheckpointState {
 	UDATA javaDebugThreadCount;
 	jvmtiEnv *jvmtienv;
 	jvmtiCapabilities requiredCapabilities;
+	BOOLEAN isDebugOnRestoreEnabled;
 } J9CRIUCheckpointState;
 #endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
 
@@ -5256,6 +5257,7 @@ typedef struct J9InternalVMFunctions {
 	BOOLEAN (*isNonPortableRestoreMode)(struct J9VMThread *currentThread);
 	BOOLEAN (*isJVMInPortableRestoreMode)(struct J9VMThread *currentThread);
 	BOOLEAN (*isDebugOnRestoreEnabled)(struct J9JavaVM *vm);
+	BOOLEAN (*isDebugAgentDisabled)(struct J9JavaVM *vm);
 	void (*setRequiredGhostFileLimit)(struct J9VMThread *currentThread, U_32 ghostFileLimit);
 	BOOLEAN (*runInternalJVMCheckpointHooks)(struct J9VMThread *currentThread, const char **nlsMsgFormat);
 	BOOLEAN (*runInternalJVMRestoreHooks)(struct J9VMThread *currentThread, const char **nlsMsgFormat);