[Orc] Memory Mapper fixes

[argentite]

Summary

[Orc] Actually save the callback

The callback function was captured by reference but it lived on the
stack and was in danger of being overwritten and could cause a crash.

[Orc] Only unmap shared memory in controller process destructor

By the time SharedMemoryMapper destructor is called, the RPC
connection is no longer available causing the release() call to
always fail. Instead at this point only shared memory regions
can be unmapped safely.

Deinitializers are called and mapped memory is released at the
executor side by ExecutorSharedMemoryMapperService::shutdown()
instead. Memory can also be released earlier by calling release()
earlier before RPC connection is closed.

[Orc] Provide correct Reservation address for slab allocations

When slab allocator is used, the MappingBase is not necessarily
the same as the original reservation base as the allocation could
be a part of the whole reservation.

In this case the original reservation address needs to be passed to
ExecutorSharedMemoryMapperService to associate the new allocation
with the original reservation.

[Orc] Improve deintialize and shutdown logic

When deinitializing, the allocation needs to be removed from the
allocation list of its associated reservation so that remaining
allocations can be deinitialized when releasing the reservation.

[lhames]

By the time SharedMemoryMapper destructor is called, the RPC
connection is no longer available causing the release() call to
always fail. Instead at this point only shared memory regions
can be unmapped safely.

I think that you can assert that there are no remaining allocations in the destructor -- these should all have been destroyed by a call to endSession(), which should happen before the allocator is destroyed.

Otherwise LGTM. Thanks Argentite!

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

35–37

Nice catch!

llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

305–308

So MappingBase is the lowest allocated address for this allocation?

llvm/lib/ExecutionEngine/Orc/TargetProcess/ExecutorSharedMemoryMapperService.cpp

284–286

No braces needed for single lines. You could sink the definition of ReservationAddrs down below the if test too:

if (Reservations.empty())
  return Error::success();

std::vector<ExecutorAddr> ReservationAddrs;
ReservationAddrs.reserve(Reservations.size());
for (const auto &R : Reservations)
  ReservationAddrs.push_back(ExecutorAddr::fromPtr(R.getFirst()));

288

Missing std::move.

[argentite]

Addressed the above issues

[argentite]

In D132313#3737642, @lhames wrote:

By the time SharedMemoryMapper destructor is called, the RPC
connection is no longer available causing the release() call to
always fail. Instead at this point only shared memory regions
can be unmapped safely.

I think that you can assert that there are no remaining allocations in the destructor -- these should all have been destroyed by a call to endSession(), which should happen before the allocator is destroyed.

We don't keep track of allocations in the controller process other than through the FinalizedAlloc class. That needs to be release()'d which is done only after deinitializing at MapperJITLinkMemoryManager::deallocate(). So I guess we are safe there.

llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

305–308

Yes it is actually the address that the slab allocator returns which may start at the middle of a reservation. We use this address as the identifier for this allocation.

[argentite]

Properly rebase everything

[argentite]

Properly base everything hopefully

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

 //=== MapperJITLinkMemoryManager.cpp - Memory management with MemoryMapper ===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/Support/Process.h"
 
-#include <limits>
-
 using namespace llvm::jitlink;
 
 namespace llvm {
 namespace orc {
 
 class MapperJITLinkMemoryManager::InFlightAlloc
     : public JITLinkMemoryManager::InFlightAlloc {
 public:
   InFlightAlloc(MapperJITLinkMemoryManager &Parent, LinkGraph &G,
                 ExecutorAddr AllocAddr,
                 std::vector<MemoryMapper::AllocInfo::SegInfo> Segs)
       : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
 
   void finalize(OnFinalizedFunction OnFinalize) override {
     MemoryMapper::AllocInfo AI;
     AI.MappingBase = AllocAddr;
 
     std::swap(AI.Segments, Segs);
     std::swap(AI.Actions, G.allocActions());
 
-    Parent.Mapper->initialize(AI, [&](Expected<ExecutorAddr> Result) {
+    Parent.Mapper->initialize(AI, [OnFinalize = std::move(OnFinalize)](
+                                      Expected<ExecutorAddr> Result) mutable {
       if (!Result) {
         OnFinalize(Result.takeError());
         return;
       }
 
       OnFinalize(FinalizedAlloc(*Result));
     });
   }
 

llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

 //===- MemoryMapper.cpp - Cross-process memory mapper ------------*- C++ -*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/MemoryMapper.h"
 
 #include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
 #include "llvm/Support/WindowsError.h"
 
+#include <algorithm>
+
 #if defined(LLVM_ON_UNIX) && !defined(__ANDROID__)
 #include <fcntl.h>
 #include <sys/mman.h>
 #include <unistd.h>
 #elif defined(_WIN32)
 #include <windows.h>
 #endif
 
char *SharedMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
 
   ExecutorAddrDiff Offset = Addr - R->first;
 
   return static_cast<char *>(R->second.LocalAddr) + Offset;
 }
 
 void SharedMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
                                     OnInitializedFunction OnInitialized) {
-  auto Reservation = Reservations.find(AI.MappingBase);
-  assert(Reservation != Reservations.end() &&
-         "Attempt to initialize unreserved range");
+  auto Reservation = std::lower_bound(
+      Reservations.rbegin(), Reservations.rend(), AI.MappingBase,
+      [](const auto &A, const auto &B) { return A.first > B; });
+  assert(Reservation != Reservations.rend() && "Attempt to initialize unreserved range");
 
   tpctypes::SharedMemoryFinalizeRequest FR;
 
   AI.Actions.swap(FR.Actions);
 
   FR.Segments.reserve(AI.Segments.size());
 
   for (auto Segment : AI.Segments) {
EPC.callSPSWrapperAsync<
           Error SerializationErr, Expected<ExecutorAddr> Result) mutable {
         if (SerializationErr) {
           cantFail(Result.takeError());
           return OnInitialized(std::move(SerializationErr));
         }
 
         OnInitialized(std::move(Result));
       },
-      SAs.Instance, AI.MappingBase, std::move(FR));
+      SAs.Instance, Reservation->first, std::move(FR));
 }
 
 void SharedMemoryMapper::deinitialize(
     ArrayRef<ExecutorAddr> Allocations,
     MemoryMapper::OnDeinitializedFunction OnDeinitialized) {
   EPC.callSPSWrapperAsync<
       rt::SPSExecutorSharedMemoryMapperServiceDeinitializeSignature>(
       SAs.Deinitialize,

 #else
   OnReleased(make_error<StringError>(
       "SharedMemoryMapper is not supported on this platform yet",
       inconvertibleErrorCode()));
 #endif
 }
 
 SharedMemoryMapper::~SharedMemoryMapper() {
-  std::vector<ExecutorAddr> ReservationAddrs;
-  if (!Reservations.empty()) {
-    std::lock_guard<std::mutex> Lock(Mutex);
-    {
-      ReservationAddrs.reserve(Reservations.size());
-      for (const auto &R : Reservations) {
-        ReservationAddrs.push_back(R.first);
-      }
-    }
-  }
+  for (const auto R : Reservations) {
 
-  std::promise<MSVCPError> P;
-  auto F = P.get_future();
-  release(ReservationAddrs, [&](Error Err) { P.set_value(std::move(Err)); });
-  // FIXME: Release can actually fail. The error should be propagated.
-  // Meanwhile, a better option is to explicitly call release().
-  cantFail(F.get());
+#if defined(LLVM_ON_UNIX) && !defined(__ANDROID__)
+
+    munmap(R.second.LocalAddr, R.second.Size);
+
+#elif defined(_WIN32)
+
+  UnmapViewOfFile(R.second.LocalAddr);
+
+#endif
+
+  }
 }
 
 } // namespace orc
 
 } // namespace llvm

llvm/lib/ExecutionEngine/Orc/TargetProcess/ExecutorSharedMemoryMapperService.cpp

Error AllErr = Error::success();
     std::lock_guard<std::mutex> Lock(Mutex);
 
     for (auto Base : Bases) {
       if (Error Err = shared::runDeallocActions(
               Allocations[Base].DeinitializationActions)) {
         AllErr = joinErrors(std::move(AllErr), std::move(Err));
       }
 
+      // Remove the allocation from the allocation list of its reservation
+      for (auto &Reservation : Reservations) {
+        auto AllocationIt =
+            std::find(Reservation.second.Allocations.begin(),
+                      Reservation.second.Allocations.end(), Base);
+        if (AllocationIt != Reservation.second.Allocations.end()) {
+          Reservation.second.Allocations.erase(AllocationIt);
+          break;
+        }
+      }
+
       Allocations.erase(Base);
     }
   }
 
   return AllErr;
 }
 
 Error ExecutorSharedMemoryMapperService::release(
#else
   return make_error<StringError>(
       "SharedMemoryMapper is not supported on this platform yet",
       inconvertibleErrorCode());
 #endif
 }
 
 Error ExecutorSharedMemoryMapperService::shutdown() {
   std::vector<ExecutorAddr> ReservationAddrs;
-  {
-    std::lock_guard<std::mutex> Lock(Mutex);
 
-    if (Reservations.empty())
-      return Error::success();
 
-    ReservationAddrs.reserve(Reservations.size());
-    for (const auto &R : Reservations) {
-      ReservationAddrs.push_back(ExecutorAddr::fromPtr(R.getFirst()));
-    }
 
-    Reservations.clear();
-  }
-
   return release(ReservationAddrs);
 }
 
 void ExecutorSharedMemoryMapperService::addBootstrapSymbols(
     StringMap<ExecutorAddr> &M) {
   M[rt::ExecutorSharedMemoryMapperServiceInstanceName] =
       ExecutorAddr::fromPtr(this);
   M[rt::ExecutorSharedMemoryMapperServiceReserveWrapperName] =
       ExecutorAddr::fromPtr(&reserveWrapper);