Merge pull request #19808 from cjjdespres/disclaim-free-message-buffer

Use custom allocator for JITServer client MessageBuffers

Author: mpirvu

runtime/compiler/control/HookedByTheJit.cpp

@@ -5320,6 +5320,11 @@ static void jitStateLogic(J9JITConfig * jitConfig, TR::CompilationInfo * compInf
          static char *disableIdleRATCleanup = feGetEnv("TR_disableIdleRATCleanup");
          if (disableIdleRATCleanup == NULL)
             persistentInfo->getRuntimeAssumptionTable()->reclaimMarkedAssumptionsFromRAT(-1);
+
+#if defined(J9VM_OPT_JITSERVER)
+         if (compInfo->getMethodQueueSize() == 0)
+            JITServer::MessageBuffer::tryFreePersistentAllocator();
+#endif /* defined(J9VM_OPT_JITSERVER) */
          }
 
       // Logic related to IdleCPU exploitation

runtime/compiler/control/rossa.cpp

@@ -1792,6 +1792,9 @@ onLoadInternal(
          }
       }
 
+   if (compInfo->getPersistentInfo()->getRemoteCompilationMode() != JITServer::NONE)
+      JITServer::MessageBuffer::initTotalBuffersMonitor();
+
    if (compInfo->getPersistentInfo()->getRemoteCompilationMode() == JITServer::SERVER)
       {
       JITServer::CommunicationStream::initConfigurationFlags();

runtime/compiler/env/PersistentAllocator.cpp

@@ -757,6 +757,21 @@ PersistentAllocator::disclaimAllSegments()
 #endif // LINUX
    return numSegDisclaimed;
    }
+
+void
+TR::PersistentAllocator::adviseDontNeedSegments()
+   {
+#ifdef LINUX
+   j9thread_monitor_enter(_segmentMonitor);
+   for (auto segmentIterator = _segments.begin(); segmentIterator != _segments.end(); ++segmentIterator)
+      {
+      J9MemorySegment &segment = *segmentIterator;
+      size_t segLength = segment.heapTop - segment.heapBase;
+      madvise(segment.heapBase, segLength, MADV_DONTNEED);
+      }
+   j9thread_monitor_exit(_segmentMonitor);
+#endif
+   }
 } // namespace J9
 
 void *

runtime/compiler/env/PersistentAllocator.hpp

@@ -77,6 +77,11 @@ class PersistentAllocator
    int disclaimAllSegments();
    int getNumSegments() const { return _numSegments; }
 
+   // Issue MADV_DONTNEED on the segments allocated by this allocator. This has the effect
+   // of zeroing out that memory if it has no persistent backing (see madvise(2)), so you must
+   // be sure that these segments are not actually in use in that case.
+   void adviseDontNeedSegments();
+
 private:
 
    // Persistent block header

runtime/compiler/net/MessageBuffer.cpp

@@ -21,18 +21,53 @@
  *******************************************************************************/
 
 #include "net/MessageBuffer.hpp"
+#include "infra/CriticalSection.hpp"
+#include "env/VerboseLog.hpp"
+#include "control/Options.hpp"
 #include <cstring>
 
 namespace JITServer
 {
+
+TR::Monitor *MessageBuffer::_totalBuffersMonitor = NULL;
+int MessageBuffer::_totalBuffers = 0;
+TR::PersistentAllocator *MessageBuffer::_allocator = NULL;
+
 MessageBuffer::MessageBuffer() :
-   _capacity(INITIAL_BUFFER_SIZE),
-   _allocator(TR::Compiler->persistentGlobalAllocator())
+   _capacity(INITIAL_BUFFER_SIZE)
    {
+   OMR::CriticalSection cs(getTotalBuffersMonitor());
+
+   if (!_allocator)
+      {
+      if (J9::PersistentInfo::_remoteCompilationMode == JITServer::CLIENT)
+         {
+         uint32_t memoryType = MEMORY_TYPE_VIRTUAL; // Force the usage of mmap for allocation
+         TR::PersistentAllocatorKit kit(1 << 20/*1 MB*/, *TR::Compiler->javaVM, memoryType);
+         _allocator = new (TR::Compiler->rawAllocator) TR::PersistentAllocator(kit);
+         }
+      else
+         {
+         _allocator = &TR::Compiler->persistentGlobalAllocator();
+         }
+      }
+
    _storage = allocateMemory(_capacity);
    if (!_storage)
       throw std::bad_alloc();
    _curPtr = _storage;
+   _totalBuffers++;
+   }
+
+MessageBuffer::~MessageBuffer()
+   {
+   OMR::CriticalSection cs(getTotalBuffersMonitor());
+
+   freeMemory(_storage);
+   _totalBuffers--;
+
+   if ((0 == _totalBuffers) && (J9::PersistentInfo::_remoteCompilationMode == JITServer::CLIENT))
+      _allocator->adviseDontNeedSegments();
    }
 
 void
@@ -79,7 +114,7 @@ MessageBuffer::writeData(const void *dataStart, uint32_t dataSize, uint8_t paddi
    _curPtr += dataSize + paddingSize;
    return offset(data);
    }
- 
+
 uint8_t
 MessageBuffer::alignCurrentPositionOn64Bit()
    {
@@ -108,5 +143,22 @@ MessageBuffer::computeRequiredCapacity(uint32_t requiredSize)
       extendedCapacity *= 2;
    return extendedCapacity;
    }
-};
 
+void
+MessageBuffer::tryFreePersistentAllocator()
+   {
+   if (J9::PersistentInfo::_remoteCompilationMode != JITServer::CLIENT)
+      return;
+
+   OMR::CriticalSection cs(getTotalBuffersMonitor());
+
+   if ((_totalBuffers != 0) || (_allocator == NULL))
+      return;
+
+   _allocator->~PersistentAllocator();
+   TR::Compiler->rawAllocator.deallocate(_allocator);
+   _allocator = NULL;
+   if (TR::Options::getVerboseOption(TR_VerbosePerformance))
+      TR_VerboseLog::writeLineLocked(TR_Vlog_PERF, "Freed message buffer storage allocator");
+   }
+};

runtime/compiler/net/MessageBuffer.hpp

@@ -27,6 +27,7 @@
 #include "env/TRMemory.hpp"
 #include "OMR/Bytes.hpp" // for alignNoCheck
 #include "env/CompilerEnv.hpp"
+#include "infra/Monitor.hpp"
 
 namespace JITServer
 {
@@ -45,16 +46,17 @@ namespace JITServer
 
    Variable _curPtr defines the boundary of the current data. Reading/writing to/from buffer
    will always advance the pointer.
+
+   A shared, reference-counted persistent allocator is used for all message buffers at a JITServer client.
+   The freed persistent memory tracked by the allocator is disclaimed when the last MessageBuffer is destroyed,
+   and the allocator itself can be destroyed with tryFreePersistentAllocator() if there are no active connections
+   to a server. Servers use the persistent global allocator.
  */
 class MessageBuffer
    {
 public:
    MessageBuffer();
-
-   ~MessageBuffer()
-      {
-      freeMemory(_storage);
-      }
+   ~MessageBuffer();
 
 
    /**
@@ -204,18 +206,28 @@ class MessageBuffer
 
    uint32_t getCapacity() const { return _capacity; }
 
+   // Must be called before any client-server communication takes place
+   static void initTotalBuffersMonitor() { _totalBuffersMonitor = TR::Monitor::create("JIT-JITServerTotalBuffersMonitor"); }
+
+   // Try to free the custom persistent allocator for message buffers. This method does nothing
+   // if the JVM is not in client mode, or if there is at least one active message buffer.
+   static void tryFreePersistentAllocator();
 
 private:
    static const size_t INITIAL_BUFFER_SIZE = 32768; // Initial buffer size is 32K
    uint32_t offset(char *addr) const { return addr - _storage; }
-   char *allocateMemory(uint32_t capacity) { return static_cast<char *>(_allocator.allocate(capacity)); }
-   void freeMemory(char *storage) { _allocator.deallocate(storage); }
+   char *allocateMemory(uint32_t capacity) { return static_cast<char *>(_allocator->allocate(capacity)); }
+   void freeMemory(char *storage) { _allocator->deallocate(storage); }
    uint32_t computeRequiredCapacity(uint32_t requiredSize);
+   static TR::Monitor *getTotalBuffersMonitor() { return _totalBuffersMonitor; }
 
    uint32_t _capacity;
    char *_storage;
    char *_curPtr;
-   TR::PersistentAllocator &_allocator;
+
+   static TR::Monitor *_totalBuffersMonitor;
+   static int _totalBuffers;
+   static TR::PersistentAllocator *_allocator;
    };
 };
 #endif