Reset JNI Addresses if FSD on restore

If debug is specified on restore, ensure that any proactively compiled
JNI Methods are reset to their original address.

JNI methods are not compiled under FSD, but they are compiled
proactively in the checkpoint hook. However, if debug is specified on
restore, all proactively compiled methods, except for JNI thunks, get
invalidated. Thus, the extra field of the J9Method of a JNI method
needs to get reset to whatever the VM initialized them to originally.

This is necessary to ensure an environment that is consistent with FSD
mode. The invalidation of proactively compiled non-JNI methods (which
are non-FSD bodies) ensure that non-FSD code does not execute
post-restore. At the checkpoint hook, only FSD bodies are on the stacks
of the threads; although the compiler currently does not reuse these
FSD bodies, they will continue to execute on restore until an OSR
transition, but this is OK because they are already set up for
involuntary OSR. The missing piece was the JNI methods, which this
commit addresses.

Signed-off-by: Irwin D'Souza <[email protected]>

runtime/compiler/control/CompileBeforeCheckpoint.cpp

@@ -53,6 +53,12 @@ TR::CompileBeforeCheckpoint::CompileBeforeCheckpoint(TR::Region &region, J9VMThr
 void
 TR::CompileBeforeCheckpoint::queueMethodForCompilationBeforeCheckpoint(J9Method *j9method, bool recomp)
    {
+   /* Do this before releasing the CRRuntime Monitor */
+   if (_compInfo->isJNINative(j9method))
+      {
+      _compInfo->getCRRuntime()->pushJNIAddr(j9method, j9method->extra);
+      }
+
    /* Release CR Runtime Monitor since compileMethod below will acquire the
     * Comp Monitor.
     */

runtime/compiler/control/OptionsPostRestore.cpp

@@ -530,7 +530,10 @@ J9::OptionsPostRestore::invalidateCompiledMethodsIfNeeded(bool invalidateAll)
       javaVM->internalVMFunctions->allClassesEndDo(&classWalkState);
 
       if (invalidateAll)
+         {
          _compInfo->getCRRuntime()->purgeMemoizedCompilations();
+         _compInfo->getCRRuntime()->resetJNIAddr();
+         }
 
       j9nls_printf(PORTLIB, (UDATA) J9NLS_WARNING, J9NLS_JIT_CHECKPOINT_RESTORE_CODE_INVALIDATED);
       }

runtime/compiler/runtime/CRRuntime.cpp

@@ -90,6 +90,7 @@ TR::CRRuntime::CRRuntime(J9JITConfig *jitConfig, TR::CompilationInfo *compInfo)
    _failedComps(),
    _forcedRecomps(),
    _impMethodForCR(),
+   _jniMethodAddr(),
    _proactiveCompEnv(),
    _restoreTime(0)
    {
@@ -164,22 +165,25 @@ TR::CRRuntime::waitOnCRRuntimeMonitor()
    _crRuntimeMonitor->wait();
    }
 
+template<typename T>
 void
-TR::CRRuntime::pushMemoizedCompilation(TR_MemoizedCompilations& list, J9Method *method)
+TR::CRRuntime::pushMemoizedCompilation(TR_MemoizedCompilations& list, J9Method *method, void *data)
    {
-   auto newEntry = new (_compInfo->persistentMemory()) TR_MemoizedComp(method);
+   auto newEntry = new (_compInfo->persistentMemory()) T(method, data);
    if (newEntry)
       list.add(newEntry);
    }
 
 J9Method *
-TR::CRRuntime::popMemoizedCompilation(TR_MemoizedCompilations& list)
+TR::CRRuntime::popMemoizedCompilation(TR_MemoizedCompilations& list, void **data)
    {
    J9Method *method = NULL;
    auto memComp = list.pop();
    if (memComp)
       {
       method = memComp->getMethod();
+      if (data)
+         *data = memComp->getData();
       jitPersistentFree(memComp);
       }
    return method;
@@ -265,6 +269,46 @@ TR::CRRuntime::purgeMemoizedCompilations()
    purgeMemoizedCompilation(_impMethodForCR);
    }
 
+void
+TR::CRRuntime::pushFailedCompilation(J9Method *method)
+   {
+   pushMemoizedCompilation<TR_MemoizedComp>(_failedComps, method);
+   }
+
+void
+TR::CRRuntime::pushForcedRecompilation(J9Method *method)
+   {
+   pushMemoizedCompilation<TR_MemoizedComp>(_forcedRecomps, method);
+   }
+
+void
+TR::CRRuntime::pushImportantMethodForCR(J9Method *method)
+   {
+   pushMemoizedCompilation<TR_MemoizedComp>(_impMethodForCR, method);
+   }
+
+void
+TR::CRRuntime::pushJNIAddr(J9Method *method, void *addr)
+   {
+   pushMemoizedCompilation<TR_JNIMethodAddr>(_jniMethodAddr, method, addr);
+   }
+
+void
+TR::CRRuntime::resetJNIAddr()
+   {
+   OMR::CriticalSection resetJNI(getCRRuntimeMonitor());
+   while (!_jniMethodAddr.isEmpty())
+      {
+      J9Method *method;
+      void *addr;
+      while ((method = popJNIAddr(&addr)))
+         {
+         TR_ASSERT_FATAL(addr, "JNI Address to be reset cannot be NULL!");
+         _compInfo->setJ9MethodExtra(method, reinterpret_cast<intptr_t>(addr));
+         }
+      }
+   }
+
 void
 TR::CRRuntime::triggerRecompilationForPreCheckpointGeneratedFSDBodies(J9VMThread *vmThread)
    {

runtime/compiler/runtime/CRRuntime.hpp

@@ -131,12 +131,14 @@ class CRRuntime
    /* The following methods should be only be invoked with the CR Runtime
     * Monitor in hand.
     */
-   void pushFailedCompilation(J9Method *method)    { pushMemoizedCompilation(_failedComps, method);    }
-   J9Method * popFailedCompilation()               { return popMemoizedCompilation(_failedComps);      }
-   void pushForcedRecompilation(J9Method *method)  { pushMemoizedCompilation(_forcedRecomps, method);  }
-   J9Method * popForcedRecompilation()             { return popMemoizedCompilation(_forcedRecomps);    }
-   void pushImportantMethodForCR(J9Method *method) { pushMemoizedCompilation(_impMethodForCR, method); }
-   J9Method * popImportantMethodForCR()            { return popMemoizedCompilation(_impMethodForCR);   }
+   void pushFailedCompilation(J9Method *method);
+   void pushForcedRecompilation(J9Method *method);
+   void pushImportantMethodForCR(J9Method *method);
+   void pushJNIAddr(J9Method *method, void *addr);
+   J9Method * popFailedCompilation()    { return popMemoizedCompilation(_failedComps);         }
+   J9Method * popForcedRecompilation()  { return popMemoizedCompilation(_forcedRecomps);       }
+   J9Method * popImportantMethodForCR() { return popMemoizedCompilation(_impMethodForCR);      }
+   J9Method * popJNIAddr(void **addr)   { return popMemoizedCompilation(_jniMethodAddr, addr); }
 
    /**
     * @brief Remove appropriate methods from all of the memoized lists. This
@@ -154,6 +156,12 @@ class CRRuntime
     */
    void purgeMemoizedCompilations();
 
+   /**
+    * @brief Reset JNI Addresses to what was set by the VM when they were
+    *        first initialized.
+    */
+   void resetJNIAddr();
+
    /**
     * @brief Processing method that the CR Runtime Thread executes.
     */
@@ -196,14 +204,35 @@ class CRRuntime
       public:
       TR_PERSISTENT_ALLOC(TR_MemoryBase::CompilationInfo);
 
-      TR_MemoizedComp(J9Method *method) { _method = method; }
+      TR_MemoizedComp(J9Method *method, void *data = NULL)
+         : _method(method)
+         {}
+
       J9Method *getMethod() { return _method; }
 
+      virtual void *getData() { TR_ASSERT_FATAL(false, "Should not be called!\n"); return NULL; }
+
       private:
       J9Method *_method;
       };
    typedef TR_LinkHead0<TR_MemoizedComp> TR_MemoizedCompilations;
 
+   class TR_JNIMethodAddr : public TR_MemoizedComp
+      {
+      public:
+      TR_PERSISTENT_ALLOC(TR_MemoryBase::CompilationInfo);
+
+      TR_JNIMethodAddr(J9Method *method, void *data)
+         : TR_MemoizedComp(method),
+           _oldJNIAddr(data)
+         {}
+
+      virtual void *getData() { return _oldJNIAddr; }
+
+      private:
+      void *_oldJNIAddr;
+      };
+
    struct ProactiveCompEnv
       {
 #if defined(J9VM_OPT_JITSERVER)
@@ -315,8 +344,10 @@ class CRRuntime
     *
     * @param list A reference to the TR_MemoizedCompilations list
     * @param method The method whose compilation should be memoized
+    * @param data The additional data to add for the J9Method
     */
-   void pushMemoizedCompilation(TR_MemoizedCompilations& list, J9Method *method);
+   template<typename T>
+   void pushMemoizedCompilation(TR_MemoizedCompilations& list, J9Method *method, void *data = NULL);
 
    /**
     * @brief Helper method to pop a J9Method from a list that is used to memoize
@@ -325,10 +356,11 @@ class CRRuntime
     *        Comp Monitor in hand.
     *
     * @param list A reference to the TR_MemoizedCompilations list
+    * @param data A pointer to memory to return the additional data (if it exists)
     *
     * @return the J9Method at the front of the list; NULL if the list is empty.
     */
-   J9Method * popMemoizedCompilation(TR_MemoizedCompilations& list);
+   J9Method * popMemoizedCompilation(TR_MemoizedCompilations& list, void **data = NULL);
 
    /**
     * @brief Helper method to remove a J9Method from a list that is used to
@@ -391,6 +423,7 @@ class CRRuntime
    TR_MemoizedCompilations _failedComps;
    TR_MemoizedCompilations _forcedRecomps;
    TR_MemoizedCompilations _impMethodForCR;
+   TR_MemoizedCompilations _jniMethodAddr;
 
    ProactiveCompEnv _proactiveCompEnv;