Use custom allocator for JITServer client MessageBuffers

A global, reference-counted allocator is now created at JITServer
clients that is used to allocate the memory used by the internal
MessageBuffer storage buffer during client/server communication.
This allocator's memory segments are disclaimed if the number of
connections to the server drops to zero, and the compiler will also
periodically attempt to destroy the allocator entirely.

Signed-off-by: Christian Despres <[email protected]>

Author: cjjdespres

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