Enable JITServer post-restore only if explicitly specified

doc/compiler/control/OptionsPostRestore.md

@@ -155,3 +155,16 @@ then specified post-restore, the options processing code will
 invalidate all JIT code in the code cache, as well as prevent further
 AOT compilation (since the AOT method header would have already been
 validated).
+
+### JITServer
+
+The following table outlines when a client JVM will connect to a
+JITServer instance.
+
+||Non-Portable CRIU Pre-Checkpoint|Non-Portable CRIU Post-Restore|Portable CRIU Pre-Checkpoint|Portable CRIU Post-Restore|
+|--|--|--|--|--|
+|No Options Pre-Checkpoint; No Options Post-Restore|:x:|:x:|:x:|:x:|
+|No Options Pre-checkpoint; `-XX:+UseJITServer` Post-Restore|:x:|:white_check_mark:|:x:|:white_check_mark:|
+|`-XX:+UseJITServer` Pre-Checkpoint; No Options Post-Restore|:x:|:white_check_mark:|:white_check_mark:|:white_check_mark:|
+|`-XX:+UseJITServer` Pre-Checkpoint; `-XX:-UseJITServer` Post-Restore|:x:|:x:|:white_check_mark:|:x:|
+|`-XX:-UseJITServer` Pre-Checkpoint; `-XX:+UseJITServer` Post-Restore|:x:|:x:|:x:|:x:|

runtime/compiler/control/CompilationRuntime.hpp

@@ -733,8 +733,20 @@ class CompilationInfo
 
    bool resetStartAndElapsedTime()              { return _resetStartAndElapsedTime;  }
    void setResetStartAndElapsedTime(bool reset) { _resetStartAndElapsedTime = reset; }
+
+#if defined(J9VM_OPT_JITSERVER)
+   bool canPerformRemoteCompilationInCRIUMode()                   { return _canPerformRemoteCompilationInCRIUMode;       }
+   void setCanPerformRemoteCompilationInCRIUMode(bool remoteComp) { _canPerformRemoteCompilationInCRIUMode = remoteComp; }
+
+   bool remoteCompilationRequestedAtBootstrap()                   { return _remoteCompilationRequestedAtBootstrap;       }
+   void setRemoteCompilationRequestedAtBootstrap(bool remoteComp) { _remoteCompilationRequestedAtBootstrap = remoteComp; }
+
+   bool remoteCompilationExplicitlyDisabledAtBootstrap()                   { return _remoteCompilationExplicitlyDisabledAtBootstrap;       }
+   void setRemoteCompilationExplicitlyDisabledAtBootstrap(bool remoteComp) { _remoteCompilationExplicitlyDisabledAtBootstrap = remoteComp; }
 #endif
 
+#endif // #if defined(J9VM_OPT_CRIU_SUPPORT)
+
    TR_PersistentMemory *     persistentMemory() { return _persistentMemory; }
 
    TR::PersistentInfo    * getPersistentInfo() { return persistentMemory()->getPersistentInfo(); }
@@ -1380,7 +1392,12 @@ class CompilationInfo
    bool _vmMethodTraceEnabled;
    bool _vmExceptionEventsHooked;
    bool _resetStartAndElapsedTime;
+#if defined(J9VM_OPT_JITSERVER)
+   bool _canPerformRemoteCompilationInCRIUMode;
+   bool _remoteCompilationRequestedAtBootstrap;
+   bool _remoteCompilationExplicitlyDisabledAtBootstrap;
 #endif
+#endif // #if defined(J9VM_OPT_CRIU_SUPPORT)
 
    TR::Monitor *_vlogMonitor;
    TR::Monitor *_rtlogMonitor;

runtime/compiler/control/CompilationThread.cpp

@@ -1184,6 +1184,7 @@ TR::CompilationInfo::CompilationInfo(J9JITConfig *jitConfig) :
 #if defined(J9VM_JIT_DYNAMIC_LOOP_TRANSFER)
    _dltMonitor = TR::Monitor::create("JIT-DLTmonitor");
 #endif
+
 #if defined(J9VM_OPT_CRIU_SUPPORT)
    _crMonitor = TR::Monitor::create("JIT-CheckpointRestoreMonitor");
    _checkpointStatus = TR_CheckpointStatus::NO_CHECKPOINT_IN_PROGRESS;
@@ -1205,7 +1206,15 @@ TR::CompilationInfo::CompilationInfo(J9JITConfig *jitConfig) :
    _vmExceptionEventsHooked = exceptionCatchEventHooked || exceptionThrowEventHooked;
 
    _resetStartAndElapsedTime = false;
+
+#if defined(J9VM_OPT_JITSERVER)
+   _canPerformRemoteCompilationInCRIUMode = false;
+   _remoteCompilationRequestedAtBootstrap = false;
+   _remoteCompilationExplicitlyDisabledAtBootstrap = false;
 #endif
+
+#endif // #if defined(J9VM_OPT_CRIU_SUPPORT)
+
    _iprofilerBufferArrivalMonitor = TR::Monitor::create("JIT-IProfilerBufferArrivalMonitor");
    _classUnloadMonitor = TR::MonitorTable::get()->getClassUnloadMonitor(); // by this time this variable is initialized
                                              // TODO: hang these monitors to something persistent
@@ -7413,14 +7422,7 @@ TR::CompilationInfoPerThreadBase::cannotPerformRemoteComp(
    {
    return
 #if defined(J9VM_OPT_CRIU_SUPPORT)
-          // In non-portable restore (i.e. single checkpoint) mode prevent remote compilations until
-          // after the process is restored. In portable restore mode checkpoints are always possible
-          // so there's no point to delaying remote compilations.
-          (_jitConfig->javaVM->internalVMFunctions->isCheckpointAllowed(vmThread) && _jitConfig->javaVM->internalVMFunctions->isNonPortableRestoreMode(vmThread)) ||
-
-          // If -XX:-UseJITServer is specified post-restore, then the Remote Compilation Mode will
-          // be set to JITServer::NONE
-          _compInfo.getPersistentInfo()->getRemoteCompilationMode() == JITServer::NONE ||
+          (_jitConfig->javaVM->internalVMFunctions->isCheckpointAllowed(vmThread) && !_compInfo.canPerformRemoteCompilationInCRIUMode()) ||
 #endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
           !JITServer::ClientStream::isServerCompatible(OMRPORT_FROM_J9PORT(_jitConfig->javaVM->portLibrary)) ||
           (!JITServerHelpers::isServerAvailable() && !JITServerHelpers::shouldRetryConnection(OMRPORT_FROM_J9PORT(_jitConfig->javaVM->portLibrary))) ||

runtime/compiler/control/J9Options.cpp

@@ -2300,26 +2300,35 @@ bool J9::Options::preProcessJitServer(J9JavaVM *vm, J9JITConfig *jitConfig)
          int32_t xxUseJITServerArgIndex = FIND_ARG_IN_VMARGS(EXACT_MATCH, xxUseJITServerOption, 0);
          int32_t xxDisableUseJITServerArgIndex = FIND_ARG_IN_VMARGS(EXACT_MATCH, xxDisableUseJITServerOption, 0);
 
-         bool explicitClientMode = xxUseJITServerArgIndex > xxDisableUseJITServerArgIndex;
+         bool useJitServerExplicitlySpecified = xxUseJITServerArgIndex > xxDisableUseJITServerArgIndex;
+
 #if defined(J9VM_OPT_CRIU_SUPPORT)
+         bool useJitServerExplicitlyDisabled = xxDisableUseJITServerArgIndex > xxUseJITServerArgIndex;
+
          struct J9InternalVMFunctions *ifuncs = vm->internalVMFunctions;
          J9VMThread *currentThread = ifuncs->currentVMThread(vm);
+
          // Enable JITServer client mode if
          // 1) CRIU support is enabled
-         // 2) non-portable restore mode is enabled
-         // 3) client mode is not explicitly disabled
-         // In portable restore mode let the user explicitly decide whether to enable JITServer.
-         bool implicitClientMode = ifuncs->isCRIUSupportEnabled(currentThread) &&
-                                   ifuncs->isNonPortableRestoreMode(currentThread) &&
-                                   (xxUseJITServerArgIndex >= xxDisableUseJITServerArgIndex);
+         // 2) client mode is not explicitly disabled
+         bool implicitClientMode = ifuncs->isCRIUSupportEnabled(currentThread) && !useJitServerExplicitlyDisabled;
 #endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
 
-         if (explicitClientMode
+         if (useJitServerExplicitlySpecified
 #if defined(J9VM_OPT_CRIU_SUPPORT)
              || implicitClientMode
 #endif /* defined(J9VM_OPT_CRIU_SUPPORT) */
          )
             {
+#if defined(J9VM_OPT_CRIU_SUPPORT)
+            if (implicitClientMode && useJitServerExplicitlySpecified)
+               {
+               compInfo->setRemoteCompilationRequestedAtBootstrap(true);
+               if (!ifuncs->isNonPortableRestoreMode(currentThread))
+                   compInfo->setCanPerformRemoteCompilationInCRIUMode(true);
+               }
+#endif
+
             J9::PersistentInfo::_remoteCompilationMode = JITServer::CLIENT;
             compInfo->getPersistentInfo()->setSocketTimeout(DEFAULT_JITCLIENT_TIMEOUT);
 
@@ -2355,7 +2364,14 @@ bool J9::Options::preProcessJitServer(J9JavaVM *vm, J9JITConfig *jitConfig)
                compInfo->getPersistentInfo()->setJITServerAOTCacheName(name);
                }
             }
+#if defined(J9VM_OPT_CRIU_SUPPORT)
+         else if (useJitServerExplicitlyDisabled)
+            {
+            compInfo->setRemoteCompilationExplicitlyDisabledAtBootstrap(true);
+            }
+#endif // #if defined(J9VM_OPT_CRIU_SUPPORT)
          }
+
       if (!JITServerParseCommonOptions(vm->vmArgsArray, vm, compInfo))
          {
          // Could not parse JITServer options successfully

runtime/compiler/control/OptionsPostRestore.cpp

@@ -280,13 +280,17 @@ void
 J9::OptionsPostRestore::processJitServerOptions()
    {
 #if defined(J9VM_OPT_JITSERVER)
-   if (_argIndexUseJITServer >= _argIndexDisableUseJITServer)
+   bool jitserverEnabled
+      = ((_argIndexUseJITServer > _argIndexDisableUseJITServer)
+          && !_compInfo->remoteCompilationExplicitlyDisabledAtBootstrap())
+        || ((_argIndexUseJITServer == _argIndexDisableUseJITServer)
+             && _compInfo->remoteCompilationRequestedAtBootstrap());
+
+   if (jitserverEnabled)
       {
       // Needed for GET_OPTION_VALUE_RESTORE_ARGS
       J9JavaVM *vm = _jitConfig->javaVM;
 
-      // TODO: Determine what other JITServer init is required
-
       // Parse common options
       if (!TR::Options::JITServerParseCommonOptions(vm->checkpointState.restoreArgsList, vm, _compInfo))
          {
@@ -314,11 +318,14 @@ J9::OptionsPostRestore::processJitServerOptions()
          }
       // Re-compute client UID post restore
       uint64_t clientUID = JITServerHelpers::generateUID();
-      _compInfo->getPersistentInfo()->setClientUID(clientUID);
       _jitConfig->clientUID = clientUID;
+      _compInfo->getPersistentInfo()->setClientUID(clientUID);
+      _compInfo->getPersistentInfo()->setServerUID(0);
+      _compInfo->setCanPerformRemoteCompilationInCRIUMode(true);
       }
    else
       {
+      _compInfo->setCanPerformRemoteCompilationInCRIUMode(false);
       _compInfo->getPersistentInfo()->setClientUID(0);
       _compInfo->getPersistentInfo()->setServerUID(0);
       _jitConfig->clientUID = 0;

test/functional/cmdLineTests/criu/build.xml

@@ -75,6 +75,7 @@
 			<fileset dir="${src}/../" includes="*.xml" />
 			<fileset dir="${src}/../" includes="*.mk" />
 			<fileset dir="${src}/../" includes="*.sh" />
+			<fileset dir="${src}/../../jitserver" includes="jitserverconfig.sh" />
 		</copy>
 	</target>
 

test/functional/cmdLineTests/criu/criuJitServerScript.sh

@@ -0,0 +1,70 @@
+#!/bin/sh
+
+#
+# Copyright IBM Corp. and others 2023
+#
+# This program and the accompanying materials are made available under
+# the terms of the Eclipse Public License 2.0 which accompanies this
+# distribution and is available at https://www.eclipse.org/legal/epl-2.0/
+# or the Apache License, Version 2.0 which accompanies this distribution and
+# is available at https://www.apache.org/licenses/LICENSE-2.0.
+#
+# This Source Code may also be made available under the following
+# Secondary Licenses when the conditions for such availability set
+# forth in the Eclipse Public License, v. 2.0 are satisfied: GNU
+# General Public License, version 2 with the GNU Classpath
+# Exception [1] and GNU General Public License, version 2 with the
+# OpenJDK Assembly Exception [2].
+#
+# [1] https://www.gnu.org/software/classpath/license.html
+# [2] https://openjdk.org/legal/assembly-exception.html
+#
+# SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 OR GPL-2.0 WITH Classpath-exception-2.0 OR LicenseRef-GPL-2.0 WITH Assembly-exception
+#
+
+echo "start running script";
+# the expected arguments are:
+# $1 is the TEST_ROOT
+# $2 is the TEST_JDK_BIN
+# $3 is the JVM_OPTIONS
+# $4 is the MAINCLASS
+# $5 is the APP ARGS
+# $6 is the NUM_CHECKPOINT
+# $7 is the KEEP_CHECKPOINT
+
+source $1/jitserverconfig.sh
+
+echo "export GLIBC_TUNABLES=glibc.cpu.hwcaps=-XSAVEC,-XSAVE,-AVX2,-ERMS,-AVX,-AVX_Fast_Unaligned_Load";
+export GLIBC_TUNABLES=glibc.pthread.rseq=0:glibc.cpu.hwcaps=-XSAVEC,-XSAVE,-AVX2,-ERMS,-AVX,-AVX_Fast_Unaligned_Load
+echo "export LD_BIND_NOT=on";
+export LD_BIND_NOT=on
+
+JITSERVER_PORT=$(random_port)
+JITSERVER_OPTIONS="-XX:JITServerPort=$JITSERVER_PORT"
+
+echo "Starting $2/jitserver $JITSERVER_OPTIONS"
+$2/jitserver $JITSERVER_OPTIONS &
+sleep 2
+
+$2/java -XX:+EnableCRIUSupport -XX:JITServerPort=$JITSERVER_PORT $3 -cp "$1/criu.jar" $4 $5 -XX:JITServerPort=$JITSERVER_PORT $6 >testOutput 2>&1;
+
+if [ "$7" != true ]; then
+    NUM_CHECKPOINT=$6
+    for ((i=0; i<$NUM_CHECKPOINT; i++)); do
+        sleep 2;
+        criu restore -D ./cpData --shell-job;
+    done
+fi
+
+cat testOutput;
+
+if  [ "$7" != true ]; then
+    rm -rf testOutput
+    echo "Removed testOutput file"
+fi
+
+echo "Terminating $2/jitserver $JITSERVER_OPTIONS"
+pkill -9 -xf "$2/jitserver $JITSERVER_OPTIONS"
+sleep 2
+
+echo "finished script";