Merge pull request #19166 from cjjdespres/0.44-handle-cache-unavailability

mpirvu · web-flow · commit 747f86c91abe · 2024-03-18T12:32:37.000-04:00
(0.44) Enable per-connection JITServer AOT cache disabling
diff --git a/runtime/compiler/control/CompilationThread.cpp b/runtime/compiler/control/CompilationThread.cpp
@@ -7940,6 +7940,7 @@ TR::CompilationInfoPerThreadBase::preCompilationTasks(J9VMThread * vmThread,
             && persistentInfo->getJITServerUseAOTCache() // Ideally we would check if the options allow us to use the AOT cache for this method
             && persistentInfo->getJITServerAOTCacheIgnoreLocalSCC() // Need to be using the new implementation
             && !entry->_doNotLoadFromJITServerAOTCache
+            && !persistentInfo->doNotRequestJITServerAOTCacheStore()
             && !_compInfo.getLowCompDensityMode() // AOT req is sent to JITServer and we don't want to send JITServer any messages in this mode
             && !cannotDoRemoteCompilation; // Make sure remote compilations are possible (no strong guarantees)
 #endif /* defined(J9VM_OPT_JITSERVER) */
diff --git a/runtime/compiler/control/JITClientCompilationThread.cpp b/runtime/compiler/control/JITClientCompilationThread.cpp
@@ -190,7 +190,8 @@ handleServerMessage(JITServer::ClientStream *client, TR_J9VM *fe, JITServer::Mes
       if ((response != MessageType::compilationCode) &&
           (response != MessageType::compilationFailure) &&
           (response != MessageType::AOTCache_serializedAOTMethod) &&
-          (response != MessageType::AOTCache_storedAOTMethod))
+          (response != MessageType::AOTCache_storedAOTMethod) &&
+          (response != MessageType::AOTCache_failure))
          client->writeError(JITServer::MessageType::compilationInterrupted, 0 /* placeholder */);
 
       if (TR::Options::isAnyVerboseOptionSet(TR_VerboseJITServer, TR_VerboseCompilationDispatch))
@@ -210,6 +211,7 @@ handleServerMessage(JITServer::ClientStream *client, TR_J9VM *fe, JITServer::Mes
       case MessageType::compilationFailure:
       case MessageType::AOTCache_storedAOTMethod:
       case MessageType::AOTCache_serializedAOTMethod:
+      case MessageType::AOTCache_failure:
       case MessageType::compilationThreadCrashed:
          done = true;
          break;
@@ -224,7 +226,9 @@ handleServerMessage(JITServer::ClientStream *client, TR_J9VM *fe, JITServer::Mes
          {
          uint64_t serverUID = std::get<0>(client->getRecvData<uint64_t>());
          uint64_t previousUID = compInfo->getPersistentInfo()->getServerUID();
+         bool hasEverConnectedToServer = compInfo->getPersistentInfo()->hasEverConnectedToServer();
          compInfo->getPersistentInfo()->setServerUID(serverUID);
+         compInfo->getPersistentInfo()->setHasEverConnectedToServer();
 
          auto unloadedClasses = comp->getPersistentInfo()->getUnloadedClassAddresses();
          std::vector<TR_AddressRange> ranges;
@@ -242,11 +246,18 @@ handleServerMessage(JITServer::ClientStream *client, TR_J9VM *fe, JITServer::Mes
             compInfo->getclassesCachedAtServer().clear();
             }
 
-         auto deserializer = compInfo->getJITServerAOTDeserializer();
-         // Reset AOT deserializer if connected to a new server (cached serialization records are now invalid),
-         // except if this is the first server this client has connected to.
-         if (deserializer && (0 != previousUID) && (previousUID != serverUID))
-            deserializer->reset(compInfoPT);
+         // This is a connection to a new server after a previous disconnection
+         if (hasEverConnectedToServer && (previousUID != serverUID))
+            {
+            // Reset AOT deserializer (cached serialization records are now invalid)
+            auto deserializer = compInfo->getJITServerAOTDeserializer();
+            if (deserializer)
+               deserializer->reset(compInfoPT);
+
+            // Do not forbid AOT cache stores or loads (this server might be able to fulfill them)
+            compInfo->getPersistentInfo()->setDoNotRequestJITServerAOTCacheLoad(false);
+            compInfo->getPersistentInfo()->setDoNotRequestJITServerAOTCacheStore(false);
+            }
 
          client->write(response, ranges, unloadedClasses->getMaxRanges(), serializedCHTable);
 
@@ -3123,6 +3134,7 @@ remoteCompile(J9VMThread *vmThread, TR::Compilation *compiler, TR_ResolvedMethod
       aotCacheStore = entry->_useAOTCacheCompilation &&
                       compInfo->methodCanBeJITServerAOTCacheStored(compiler->signature(), compilee->convertToMethod()->methodType());
       aotCacheLoad = persistentInfo->getJITServerUseAOTCache() &&
+                     !persistentInfo->doNotRequestJITServerAOTCacheLoad() &&
                      !TR::CompilationInfo::isCompiled(method) &&
                      !entry->_doNotLoadFromJITServerAOTCache &&
                      compInfo->methodCanBeJITServerAOTCacheLoaded(compiler->signature(), compilee->convertToMethod()->methodType());
@@ -3502,6 +3514,36 @@ remoteCompile(J9VMThread *vmThread, TR::Compilation *compiler, TR_ResolvedMethod
          entry->_compErrCode = compilationFailure;
          compiler->failCompilation<JITServer::ServerCompilationFailure>("JITServer compilation thread has crashed.");
          }
+      else if (JITServer::MessageType::AOTCache_failure == response)
+         {
+         auto recv = client->getRecvData<bool, bool>();
+         bool aotCacheUnavailable = std::get<0>(recv);
+         bool aotCacheStoreUnavailable = std::get<1>(recv);
+
+         auto persistentInfo = compInfo->getPersistentInfo();
+         if (aotCacheUnavailable)
+            {
+            persistentInfo->setDoNotRequestJITServerAOTCacheLoad(true);
+            persistentInfo->setDoNotRequestJITServerAOTCacheStore(true);
+            if (TR::Options::getVerboseOption(TR_VerboseJITServer))
+               TR_VerboseLog::writeLineLocked(TR_Vlog_JITServer,
+                  "Server AOT cache unavailable (serverUID=%llu). Disabling future AOT cache requests for this server.",
+                  (unsigned long long)persistentInfo->getServerUID());
+            }
+         else if (aotCacheStoreUnavailable)
+            {
+            persistentInfo->setDoNotRequestJITServerAOTCacheStore(true);
+            if (TR::Options::getVerboseOption(TR_VerboseJITServer))
+               TR_VerboseLog::writeLineLocked(TR_Vlog_JITServer,
+                  "Cannot store methods in server AOT cache (serverUID=%llu). Disabling future AOT cache store requests for this server.",
+                  (unsigned long long)persistentInfo->getServerUID());
+            }
+         entry->_compErrCode = compilationFailure;
+         entry->_doNotLoadFromJITServerAOTCache = true;
+         if (entry->_compilationAttemptsLeft > 0)
+            entry->_tryCompilingAgain = true;
+         compiler->failCompilation<JITServer::ServerCompilationFailure>("JITServer compilation failed.");
+         }
       else
          {
          TR_ASSERT(JITServer::MessageType::compilationFailure == response,
diff --git a/runtime/compiler/control/JITServerCompilationThread.cpp b/runtime/compiler/control/JITServerCompilationThread.cpp
@@ -640,8 +640,8 @@ TR::CompilationInfoPerThreadRemote::processEntry(TR_MethodToBeCompiled &entry, J
       auto &chtableMods             = std::get<14>(req);
       useAotCompilation             = std::get<15>(req);
       bool isInStartupPhase         = std::get<16>(req);
-      bool aotCacheStore            = std::get<17>(req);
-      bool aotCacheLoad             = std::get<18>(req);
+      bool requestedAOTCacheStore   = std::get<17>(req);
+      bool requestedAOTCacheLoad    = std::get<18>(req);
       _methodIndex                  = std::get<19>(req);
       uintptr_t classChainOffset    = std::get<20>(req);
       auto &ramClassChain           = std::get<21>(req);
@@ -950,8 +950,8 @@ TR::CompilationInfoPerThreadRemote::processEntry(TR_MethodToBeCompiled &entry, J
       entry._stream = stream; // Add the stream to the entry
 
       auto aotCache = clientSession->getOrCreateAOTCache(stream);
-      _aotCacheStore = aotCacheStore && aotCache && JITServerAOTCacheMap::cacheHasSpace();
-      aotCacheLoad = aotCacheLoad && aotCache;
+      _aotCacheStore = requestedAOTCacheStore && aotCache && JITServerAOTCacheMap::cacheHasSpace();
+      bool aotCacheLoad = requestedAOTCacheLoad && aotCache;
       if (aotCache && !aotCacheLoad)
          aotCache->incNumCacheBypasses();
 
@@ -981,6 +981,23 @@ TR::CompilationInfoPerThreadRemote::processEntry(TR_MethodToBeCompiled &entry, J
 
       if (aotCacheLoad)
          aotCacheHit = serveCachedAOTMethod(entry, ramMethod, clazz, &clientOptPlan, clientSession, scratchSegmentProvider);
+
+      // If the client requests that the server use AOT cache offsets during AOT cache compilations, then
+      // the client will be ignoring its local SCC (if it even exists) for the duration of the compilation.
+      // This means that if the client requests an AOT cache store in that scenario and the server
+      // cannot fulfill that request, the server must abort the compilation.
+      // If the client isn't requesting server offsets, then the compilation can still continue, albeit as
+      // a regular AOT compilation.
+      if (clientSession->useServerOffsets(stream) &&
+          requestedAOTCacheStore &&
+          !aotCacheHit &&
+          !_aotCacheStore)
+         {
+         abortCompilation = true;
+         bool aotCacheUnavailable = !aotCache;
+         bool aotCacheStoreUnavailable = aotCacheUnavailable || !JITServerAOTCacheMap::cacheHasSpace();
+         stream->write(JITServer::MessageType::AOTCache_failure, aotCacheUnavailable, aotCacheStoreUnavailable);
+         }
       }
    catch (const JITServer::StreamFailure &e)
       {
diff --git a/runtime/compiler/env/J9PersistentInfo.hpp b/runtime/compiler/env/J9PersistentInfo.hpp
@@ -165,6 +165,7 @@ class PersistentInfo : public OMR::PersistentInfoConnector
          _socketTimeoutMs(0),
          _clientUID(0),
          _serverUID(0),
+         _hasEverConnectedToServer(false),
          _JITServerMetricsPort(38500),
          _JITServerHealthPort(38600),
          _requireJITServer(false),
@@ -175,6 +176,8 @@ class PersistentInfo : public OMR::PersistentInfoConnector
          _JITServerAOTCacheDir(),
          _JITServerAOTCacheDelayMethodRelocation(false),
          _JITServerAOTCacheIgnoreLocalSCC(false),
+         _doNotRequestJITServerAOTCacheLoad(false),
+         _doNotRequestJITServerAOTCacheStore(false),
 #endif /* defined(J9VM_OPT_JITSERVER) */
       OMR::PersistentInfoConnector(pm)
       {}
@@ -353,6 +356,8 @@ class PersistentInfo : public OMR::PersistentInfoConnector
    void setClientUID(uint64_t val) { _clientUID = val; }
    uint64_t getServerUID() const { return _serverUID; }
    void setServerUID(uint64_t val) { _serverUID = val; }
+   bool hasEverConnectedToServer() const { return _hasEverConnectedToServer; }
+   void setHasEverConnectedToServer() { _hasEverConnectedToServer = true; }
    uint32_t getJITServerMetricsPort() const { return _JITServerMetricsPort; }
    void setJITServerMetricsPort(uint32_t port) { _JITServerMetricsPort = port; }
    uint32_t getJITServerHealthPort() const { return _JITServerHealthPort; }
@@ -373,6 +378,10 @@ class PersistentInfo : public OMR::PersistentInfoConnector
    void setJITServerAOTCacheDelayMethodRelocation(bool b) { _JITServerAOTCacheDelayMethodRelocation = b; }
    bool getJITServerAOTCacheIgnoreLocalSCC() const { return _JITServerAOTCacheIgnoreLocalSCC; }
    void setJITServerAOTCacheIgnoreLocalSCC(bool b) { _JITServerAOTCacheIgnoreLocalSCC = b; }
+   bool doNotRequestJITServerAOTCacheLoad() const { return _doNotRequestJITServerAOTCacheLoad; }
+   void setDoNotRequestJITServerAOTCacheLoad(bool b) { _doNotRequestJITServerAOTCacheLoad = b; }
+   bool doNotRequestJITServerAOTCacheStore() const { return _doNotRequestJITServerAOTCacheStore; }
+   void setDoNotRequestJITServerAOTCacheStore(bool b) { _doNotRequestJITServerAOTCacheStore = b; }
 #endif /* defined(J9VM_OPT_JITSERVER) */
 
    private:
@@ -463,6 +472,7 @@ class PersistentInfo : public OMR::PersistentInfoConnector
    uint32_t    _socketTimeoutMs; // timeout for communication sockets used in out-of-process JIT compilation
    uint64_t    _clientUID;
    uint64_t    _serverUID; // At the client, this represents the UID of the server the client is connected to
+   bool        _hasEverConnectedToServer; // At the client, true if the client has connected to a server at some point in the past
    uint32_t    _JITServerMetricsPort; // Port for receiving http metrics requests from Prometheus; only used at server
    uint32_t    _JITServerHealthPort; // Port for receiving readiness/liveness probes from Kubernetes; only used at server
    bool        _requireJITServer;
@@ -472,8 +482,12 @@ class PersistentInfo : public OMR::PersistentInfoConnector
    bool        _JITServerUseAOTCachePersistence; // Whether to persist the JITServer AOT caches at the server
    std::string _JITServerAOTCacheDir;  // Directory where the JITServer persistent AOT caches are located
    bool        _JITServerAOTCacheDelayMethodRelocation; // At the client, whether to delay deserialized method relocation or not
-    // At the client, whether or not to use the new AOT cache implementation (with serialization record IDs as SCC offsets)
+   // At the client, whether or not to use the new AOT cache implementation (with serialization record IDs as SCC offsets)
    bool        _JITServerAOTCacheIgnoreLocalSCC;
+   // True if the client should not request AOT cache loads during this server connection
+   bool        _doNotRequestJITServerAOTCacheLoad;
+   // True if the client should not request AOT cache stores during this server connection
+   bool        _doNotRequestJITServerAOTCacheStore;
 #endif /* defined(J9VM_OPT_JITSERVER) */
    };
 
diff --git a/runtime/compiler/net/CommunicationStream.hpp b/runtime/compiler/net/CommunicationStream.hpp
@@ -118,7 +118,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 = 58; // ID: jdaegRYmBxGBJ/FG0fcY
+   static const uint16_t MINOR_NUMBER = 59; // ID: h4Ihmr7j+jt8ie/rpCIX
    static const uint8_t PATCH_NUMBER = 0;
    static uint32_t CONFIGURATION_FLAGS;
 
diff --git a/runtime/compiler/net/MessageTypes.cpp b/runtime/compiler/net/MessageTypes.cpp
@@ -30,6 +30,7 @@ const char *messageNames[] =
    "compilationFailure",
    "AOTCache_storedAOTMethod",
    "AOTCache_serializedAOTMethod",
+   "AOTCache_failure",
    "mirrorResolvedJ9Method",
    "get_params_to_construct_TR_j9method",
    "getUnloadedClassRangesAndCHTable",
diff --git a/runtime/compiler/net/MessageTypes.hpp b/runtime/compiler/net/MessageTypes.hpp
@@ -34,6 +34,7 @@ enum MessageType : uint16_t
    compilationFailure,
    AOTCache_storedAOTMethod, // Final response to a compilation request that was an AOT cache store that ignored the client SCC
    AOTCache_serializedAOTMethod,// Final response to a compilation request that was an AOT cache hit
+   AOTCache_failure,// Final response to a compilation request that could not be completed due to AOT cache failure
    mirrorResolvedJ9Method,
    get_params_to_construct_TR_j9method,
    getUnloadedClassRangesAndCHTable,
