Client Request Cached Methods List from JIT Server

Added an option that makes the client request a list of cached
methods from the server on bootstrap, and set the count for
those methods as scount to improve rampup

Signed-off-by: Luke Li <luke.li@ibm.com>

runtime/compiler/control/HookedByTheJit.cpp

@@ -20,6 +20,9 @@
  * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0 OR GPL-2.0-only WITH OpenJDK-assembly-exception-1.0
  *******************************************************************************/
 
+#include <string>
+#include <unordered_set>
+
 #include <algorithm>
 #include <limits.h>
 #ifdef LINUX
@@ -425,6 +428,44 @@ static void jitHookInitializeSendTarget(J9HookInterface * * hook, UDATA eventNum
          }
       }
 
+#if defined(J9VM_OPT_JITSERVER)
+   bool methodCachedAtServer = false;
+   PersistentUnorderedSet<std::string> *serverAOTMethodSet =
+      (PersistentUnorderedSet<std::string> *) jitConfig->serverAOTMethodSet;
+
+   if (serverAOTMethodSet != NULL)
+      {
+      // Construct a signature
+      J9UTF8 *className;
+      J9UTF8 *name;
+      J9UTF8 *signature;
+      getClassNameSignatureFromMethod(method, className, name, signature);
+      // SigLen calculation as used in TRJ9VMBase::printTruncatedSignature
+      int32_t sigLen = J9UTF8_LENGTH(className) + J9UTF8_LENGTH(name) +
+         J9UTF8_LENGTH(signature) + 2;
+
+      if (sigLen < 1024)
+         {
+         char sigC[1024];
+         sigLen = sprintf(sigC, "%.*s.%.*s%.*s",
+                           J9UTF8_LENGTH(className), utf8Data(className),
+                           J9UTF8_LENGTH(name), utf8Data(name),
+                           J9UTF8_LENGTH(signature), utf8Data(signature));
+
+         // contains
+         methodCachedAtServer =
+            (serverAOTMethodSet->find(std::string(sigC)) != serverAOTMethodSet->end());
+
+         if (TR::Options::getVerboseOption(TR_VerboseJITServer)
+            && TR::Options::getVerboseOption(TR_VerboseCounts))
+            if (methodCachedAtServer)
+               TR_VerboseLog::writeLineLocked(TR_Vlog_JITServer,
+                                             "Method %s was cached at the server",
+                                             sigC);
+         }
+      }
+#endif // J9VM_OPT_JITSERVER
+
    if (TR::Options::getAOTCmdLineOptions()->anOptionSetContainsACountValue())
       {
       TR::OptionSet * optionSet = findOptionSet(method, true);
@@ -628,6 +669,24 @@ static void jitHookInitializeSendTarget(J9HookInterface * * hook, UDATA eventNum
             }
 #endif // defined(J9VM_INTERP_AOT_COMPILE_SUPPORT) && defined(J9VM_OPT_SHARED_CLASSES) && (defined(TR_HOST_X86) || defined(TR_HOST_POWER) || defined(TR_HOST_S390) || defined(TR_HOST_ARM) || defined(TR_HOST_ARM64))
          } // if (TR::Options::sharedClassCache())
+#if defined(J9VM_OPT_JITSERVER)
+      else if (methodCachedAtServer) // Not cached in SCC but the server
+         {
+         int32_t scount = TR_INITIAL_SCOUNT;
+         if (optionsAOT)
+            {
+            scount = optionsAOT->getInitialSCount();
+            if ((scount == TR_QUICKSTART_INITIAL_SCOUNT) || (scount == TR_INITIAL_SCOUNT))
+               {
+               // If scount is not user specified (coarse way due to info being lost
+               // from options parsing)
+               scount= std::min(getCount(romMethod, optionsJIT, optionsAOT),
+                                       optionsAOT->getInitialSCount());
+               }
+            }
+         count = scount;
+         }
+#endif // J9VM_OPT_JITSERVER
       if (count == -1) // count didn't change yet
          {
          if (!TR::Options::getCountsAreProvidedByUser() &&
@@ -4181,6 +4240,14 @@ void JitShutdown(J9JITConfig * jitConfig)
 
    TR::CompilationInfo * compInfo = TR::CompilationInfo::get(jitConfig);
 
+   PersistentUnorderedSet<std::string> *serverAOTMethodSet =
+      (PersistentUnorderedSet<std::string> *) jitConfig->serverAOTMethodSet;
+   if (serverAOTMethodSet)
+      {
+      serverAOTMethodSet->~unordered_set();
+      TR_Memory::jitPersistentFree((void *)serverAOTMethodSet);
+      }
+
 #if defined(J9VM_OPT_CRIU_SUPPORT)
    if (jitConfig->javaVM->internalVMFunctions->isCRaCorCRIUSupportEnabled(vmThread))
       {

runtime/compiler/control/JITServerCompilationThread.cpp

@@ -20,6 +20,8 @@
  * SPDX-License-Identifier: EPL-2.0 OR Apache-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0 OR GPL-2.0-only WITH OpenJDK-assembly-exception-1.0
  *******************************************************************************/
 
+#include <string.h>
+
 #include "control/JITServerCompilationThread.hpp"
 
 #include "codegen/CodeGenerator.hpp"
@@ -202,9 +204,14 @@ outOfProcessCompilationEnd(TR_MethodToBeCompiled *entry, TR::Compilation *comp)
       CachedAOTMethod *freshMethodRecord = NULL;
       if (!methodRecord)
          {
-         freshMethodRecord = CachedAOTMethod::create(compInfoPT->getDefiningClassChainRecord(), compInfoPT->getMethodIndex(),
-                                                     entry->_optimizationPlan->getOptLevel(), clientData->getAOTHeaderRecord(),
-                                                     comp->getSerializationRecords(), codeCacheHeader, codeSize, dataCacheHeader, dataSize);
+         freshMethodRecord = CachedAOTMethod::create(compInfoPT->getDefiningClassChainRecord(),
+                                                     compInfoPT->getMethodIndex(),
+                                                     entry->_optimizationPlan->getOptLevel(),
+                                                     clientData->getAOTHeaderRecord(),
+                                                     comp->getSerializationRecords(),
+                                                     codeCacheHeader, codeSize,
+                                                     dataCacheHeader, dataSize,
+                                                     comp->signature());
          methodRecord = freshMethodRecord;
          }
 
@@ -1101,6 +1108,69 @@ TR::CompilationInfoPerThreadRemote::processEntry(TR_MethodToBeCompiled &entry, J
       stream->writeError(compilationLowPhysicalMemory, (uint64_t) computeServerMemoryState(getCompilationInfo()));
       abortCompilation = true;
       }
+   catch (const JITServer::StreamAotCacheMapRequest &e)
+      {
+      const std::string& aotCacheName = e.getCacheName();
+
+      if (TR::Options::getVerboseOption(TR_VerboseJITServer))
+         {
+         TR_VerboseLog::writeLineLocked(TR_Vlog_JITServer,
+            "compThreadID=%d handling request for AOT cache %s method list",
+            getCompThreadId(), aotCacheName.c_str());
+         }
+
+      auto aotCacheMap = compInfo->getJITServerAOTCacheMap();
+      bool pending = false;
+      auto aotCache = aotCacheMap->get(aotCacheName, 0, pending);
+
+      std::vector<std::string> methodSignaturesV;
+      if (aotCache)
+         {
+         auto cachedAOTMethod = aotCache->getCachedMethodHead();
+         auto cachedMethodMonitor = aotCache->getCachedMethodMonitor();
+
+         methodSignaturesV.reserve(aotCache->getCachedMethodMap().size());
+
+         try
+            {
+               {
+               OMR::CriticalSection cs(cachedMethodMonitor);
+
+               for (;cachedAOTMethod != NULL;
+                  cachedAOTMethod = cachedAOTMethod->getNextRecord())
+                  {
+                  const SerializedAOTMethod &serializedAOTMethod = cachedAOTMethod->data();
+                  methodSignaturesV.push_back(std::string(serializedAOTMethod.signature()));
+                  }
+               }
+            }
+         catch (const std::bad_alloc &e)
+            {
+            if (TR::Options::isAnyVerboseOptionSet(TR_VerboseJITServer))
+               TR_VerboseLog::writeLineLocked(TR_Vlog_FAILURE,
+                  "std::bad_alloc: %s",
+                  e.what());
+            }
+
+         if (TR::Options::getVerboseOption(TR_VerboseJITServer))
+            {
+            TR_VerboseLog::writeLineLocked(TR_Vlog_JITServer,
+                                          "Sending the list of AOT methods size %d",
+                                          methodSignaturesV.size());
+            }
+         }
+      else // Failed getting aotCache, treat pending as a failure
+         {
+         if (TR::Options::getVerboseOption(TR_VerboseJITServer))
+            {
+            TR_VerboseLog::writeLineLocked(TR_Vlog_JITServer, "Failed getting aotCache");
+            }
+         }
+      stream->write(JITServer::MessageType::AOTCacheMap_reply, methodSignaturesV);
+
+      abortCompilation = true;
+      deleteStream = true;
+      }
 
    // Acquire VM access
    //

runtime/compiler/control/rossa.cpp

@@ -24,6 +24,9 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include <vector>
+#include <string>
+
 #ifdef WINDOWS
 // Undefine the winsockapi because winsock2 defines it. Removes warnings.
 #if defined(_WINSOCKAPI_) && !defined(_WINSOCK2API_)
@@ -1917,6 +1920,83 @@ onLoadInternal(
    return 0;
    }
 
+#if defined(J9VM_OPT_JITSERVER)
+static int32_t J9THREAD_PROC fetchServerCachedAOTMethods(void * entryarg)
+   {
+   J9JITConfig *jitConfig = (J9JITConfig *) entryarg;
+   J9JavaVM *vm = jitConfig->javaVM;
+   TR::CompilationInfo * compInfo = TR::CompilationInfo::get(jitConfig);
+   TR::PersistentInfo *persistentInfo = compInfo->getPersistentInfo();
+
+   j9thread_t osThread = (j9thread_t) jitConfig->serverAOTQueryThread;
+   J9VMThread *vmThread = NULL;
+
+   int rc = vm->internalVMFunctions->internalAttachCurrentThread(vm, &vmThread, NULL,
+                              J9_PRIVATE_FLAGS_DAEMON_THREAD | J9_PRIVATE_FLAGS_NO_OBJECT |
+                              J9_PRIVATE_FLAGS_SYSTEM_THREAD | J9_PRIVATE_FLAGS_ATTACHED_THREAD,
+                              osThread);
+
+   if (rc != JNI_OK)
+   {
+      return rc;
+   }
+
+   try
+      {
+      JITServer::ClientStream *client = new (PERSISTENT_NEW) JITServer::ClientStream(persistentInfo);
+      client->write(JITServer::MessageType::AOTCacheMap_request,
+                     persistentInfo->getJITServerAOTCacheName());
+
+      client->read();
+      auto result = client->getRecvData<std::vector<std::string>>();
+
+      std::vector<std::string> cachedMethods = std::get<0>(result);
+
+      if (TR::Options::getVerboseOption(TR_VerboseJITServer))
+         TR_VerboseLog::writeLineLocked(TR_Vlog_JITServer, "Received %d methods",
+                                       cachedMethods.size());
+
+      PersistentUnorderedSet<std::string> *serverAOTMethodSet =
+         new (PERSISTENT_NEW) PersistentUnorderedSet<std::string>(
+            PersistentUnorderedSet<std::string>::allocator_type
+               (TR::Compiler->persistentAllocator()));
+
+      for (const auto &methodSig : cachedMethods)
+         {
+         serverAOTMethodSet->insert(methodSig);
+         }
+
+      client->~ClientStream();
+      TR_Memory::jitPersistentFree(client);
+
+      FLUSH_MEMORY(TR::Compiler->target.isSMP());
+      jitConfig->serverAOTMethodSet = (void *) serverAOTMethodSet;
+      }
+   catch (const JITServer::StreamFailure &e)
+      {
+      if (TR::Options::isAnyVerboseOptionSet(TR_VerboseJITServer, TR_VerboseCompilationDispatch))
+         TR_VerboseLog::writeLineLocked(TR_Vlog_FAILURE,
+            "JITServer::StreamFailure: %s",
+            e.what());
+
+      JITServerHelpers::postStreamFailure(
+         OMRPORT_FROM_J9PORT(vm->portLibrary),
+         compInfo, e.retryConnectionImmediately(), true);
+      }
+   catch (const std::bad_alloc &e)
+      {
+      if (TR::Options::isAnyVerboseOptionSet(TR_VerboseJITServer, TR_VerboseCompilationDispatch))
+         TR_VerboseLog::writeLineLocked(TR_Vlog_FAILURE,
+            "std::bad_alloc: %s",
+            e.what());
+      }
+
+   vm->internalVMFunctions->DetachCurrentThread((JavaVM *) vm);
+   j9thread_exit(NULL);
+
+   return 0;
+   }
+#endif // J9VM_OPT_JITSERVER
 
 extern "C" int32_t
 aboutToBootstrap(J9JavaVM * javaVM, J9JITConfig * jitConfig)
@@ -2184,6 +2264,43 @@ aboutToBootstrap(J9JavaVM * javaVM, J9JITConfig * jitConfig)
    UT_MODULE_LOADED(J9_UTINTERFACE_FROM_VM(javaVM));
    Trc_JIT_VMInitStages_Event1(curThread);
    Trc_JIT_portableSharedCache_enabled_or_disabled(curThread, J9_ARE_ANY_BITS_SET(javaVM->extendedRuntimeFlags2, J9_EXTENDED_RUNTIME2_ENABLE_PORTABLE_SHARED_CACHE) ? 1 : 0);
+
+#if defined(J9VM_OPT_JITSERVER)
+   jitConfig->serverAOTMethodSet = NULL;
+   if (TR::Options::getCmdLineOptions()->getOption(TR_RequestJITServerCachedMethods))
+      {
+      // Ask the server for its cached methods
+      if (compInfo->getPersistentInfo()->getRemoteCompilationMode() == JITServer::CLIENT)
+         {
+         if (JITServerHelpers::isServerAvailable())
+            {
+            if (TR::Options::getVerboseOption(TR_VerboseJITServer))
+               TR_VerboseLog::writeLineLocked(TR_Vlog_JITServer,
+                  "Creating a thread to ask the server for its cached methods");
+
+            IDATA result = javaVM->internalVMFunctions->createThreadWithCategory(
+               (omrthread_t *) &(jitConfig->serverAOTQueryThread),
+               javaVM->defaultOSStackSize,
+               J9THREAD_PRIORITY_NORMAL,
+               0,
+               &fetchServerCachedAOTMethods,
+               (void *) jitConfig,
+               J9THREAD_CATEGORY_SYSTEM_JIT_THREAD
+            );
+
+            if (result != J9THREAD_SUCCESS)
+               {
+               jitConfig->serverAOTMethodSet = NULL;
+               if (TR::Options::getVerboseOption(TR_VerboseJITServer))
+                  TR_VerboseLog::writeLineLocked(TR_Vlog_JITServer,
+                                                "Query thread not created");
+               }
+            }
+         }
+      }
+#endif // J9VM_OPT_JITSERVER
+
+
    return 0;
    }
 

runtime/compiler/net/CommunicationStream.hpp

@@ -128,7 +128,7 @@ class CommunicationStream
    // likely to lose an increment when merging/rebasing/etc.
    //
    static const uint8_t MAJOR_NUMBER = 1;
-   static const uint16_t MINOR_NUMBER = 65; // ID: YxVkiLqD7B1LhYMv58y8
+   static const uint16_t MINOR_NUMBER = 66; // ID: KM2yGpoPgxfVKNeftLTB
    static const uint8_t PATCH_NUMBER = 0;
    static uint32_t CONFIGURATION_FLAGS;
 

runtime/compiler/net/MessageTypes.cpp

@@ -257,6 +257,8 @@ const char *messageNames[] =
    "KnownObjectTable_getReferenceField",
    "KnownObjectTable_getKnownObjectTableDumpInfo",
    "AOTCache_getROMClassBatch",
+   "AOTCacheMap_request",
+   "AOTCacheMap_reply"
    };
 
    static_assert(sizeof(messageNames) / sizeof(messageNames[0]) == MessageType_MAXTYPE,

runtime/compiler/net/MessageTypes.hpp

@@ -284,6 +284,9 @@ enum MessageType : uint16_t
 
    AOTCache_getROMClassBatch,
 
+   AOTCacheMap_request,
+   AOTCacheMap_reply,
+
    MessageType_MAXTYPE
    };
    extern const char *messageNames[];

runtime/compiler/net/ServerStream.hpp

@@ -179,6 +179,11 @@ class ServerStream : public CommunicationStream
             {
             return getArgsRaw<T...>(_cMsg);
             }
+         case MessageType::AOTCacheMap_request:
+            {
+            std::string cacheName = std::get<0>(getArgsRaw<std::string>(_cMsg));
+            throw StreamAotCacheMapRequest(cacheName);
+            }
          default:
             {
             throw StreamMessageTypeMismatch(MessageType::compilationRequest, _cMsg.type());

runtime/compiler/net/StreamExceptions.hpp

@@ -81,6 +81,25 @@ class StreamClientSessionTerminate: public virtual std::exception
    uint64_t _clientId;
    };
 
+
+class StreamAotCacheMapRequest: public virtual std::exception
+   {
+public:
+   StreamAotCacheMapRequest(std::string cacheName) : _cacheName(cacheName)
+      {
+      }
+   virtual const char* what() const throw()
+      {
+      return "Requesting AOT cache content";
+      }
+   const std::string& getCacheName() const
+      {
+      return _cacheName;
+      }
+private:
+   const std::string _cacheName;
+   };
+
 class StreamOOO : public virtual std::exception
    {
    public: