[Orc] Properly deallocate mapped memory in MapperJITLinkMemoryManager

[argentite]

Summary

When memory is deallocated from MapperJITLinkMemoryManager deinitialize
actions are run through mapper and in case of InProcessMapper, memory
protections of the region are reset to read/write as they were previously
changed and can be reused in future.

[argentite]

We are resetting permissions to read/write in InProcessMapper::deinitialize(). It may make sense to set them to read/write in InProcessMapper::prepare() instead. But in that case extra syscalls are made during startup instead of shutdown.

[lhames]

LGTM. Thanks Argentite!

llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h

public:
 
   void release(ArrayRef<ExecutorAddr> Reservations,
                OnReleasedFunction OnRelease) override;
 
   ~InProcessMemoryMapper() override;
 
 private:
   struct Allocation {
+    size_t Size;
     std::vector<shared::WrapperFunctionCall> DeinitializationActions;
   };
   using AllocationMap = DenseMap<ExecutorAddr, Allocation>;
 
   struct Reservation {
     size_t Size;
     std::vector<ExecutorAddr> Allocations;
   };

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

if (SelectedRange.empty()) { // no already reserved range was found
     Mapper->reserve(TotalAllocation, std::move(CompleteAllocation));
   } else {
     CompleteAllocation(SelectedRange);
   }
 }
 
 void MapperJITLinkMemoryManager::deallocate(
     std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
+  std::vector<ExecutorAddr> Bases;
+  Bases.reserve(Allocs.size());
+  for (auto &FA : Allocs) {
+    ExecutorAddr Addr = FA.getAddress();
+    Bases.push_back(Addr);
+  }
+
+  Mapper->deinitialize(Bases, [this, Allocs = std::move(Allocs),
+                               OnDeallocated = std::move(OnDeallocated)](
+                                  llvm::Error Err) mutable {
+    if (Err)
+      OnDeallocated(std::move(Err));
+
+    {
       std::lock_guard<std::mutex> Lock(Mutex);
 
       for (auto &FA : Allocs) {
         ExecutorAddr Addr = FA.getAddress();
         ExecutorAddrDiff Size = UsedMemory[Addr];
-    UsedMemory.erase(Addr);
 
+        UsedMemory.erase(Addr);
         AvailableMemory.push_back({Addr, Addr + Size});
-    FA.release();
       }
+    }
+
     OnDeallocated(Error::success());
+  });
 }
 
 } // end namespace orc
 } // end namespace llvm

llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

 
 char *InProcessMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
   return Addr.toPtr<char *>();
 }
 
 void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
                                        OnInitializedFunction OnInitialized) {
   ExecutorAddr MinAddr(~0ULL);
+  ExecutorAddr MaxAddr(0);
 
   for (auto &Segment : AI.Segments) {
     auto Base = AI.MappingBase + Segment.Offset;
     auto Size = Segment.ContentSize + Segment.ZeroFillSize;
 
     if (Base < MinAddr)
       MinAddr = Base;
 
+    if (Base + Size > MaxAddr)
+      MaxAddr = Base + Size;
+
     std::memset((Base + Segment.ContentSize).toPtr<void *>(), 0,
                 Segment.ZeroFillSize);
 
     if (auto EC = sys::Memory::protectMappedMemory({Base.toPtr<void *>(), Size},
                                                    Segment.Prot)) {
       return OnInitialized(errorCodeToError(EC));
     }
     if (Segment.Prot & sys::Memory::MF_EXEC)
       sys::Memory::InvalidateInstructionCache(Base.toPtr<void *>(), Size);
   }
 
   auto DeinitializeActions = shared::runFinalizeActions(AI.Actions);
   if (!DeinitializeActions)
     return OnInitialized(DeinitializeActions.takeError());
 
   {
     std::lock_guard<std::mutex> Lock(Mutex);
+
+    // This is the maximum range whose permission have been possibly modified
+    Allocations[MinAddr].Size = MaxAddr - MinAddr;
     Allocations[MinAddr].DeinitializationActions =
         std::move(*DeinitializeActions);
     Reservations[AI.MappingBase.toPtr<void *>()].Allocations.push_back(MinAddr);
   }
 
   OnInitialized(MinAddr);
 }
 
 void InProcessMemoryMapper::deinitialize(
     ArrayRef<ExecutorAddr> Bases,
     MemoryMapper::OnDeinitializedFunction OnDeinitialized) {
   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));
       }
 
+      // Reset protections to read/write so the area can be reused
+      if (auto EC = sys::Memory::protectMappedMemory(
+              {Base.toPtr<void *>(), Allocations[Base].Size},
+              sys::Memory::ProtectionFlags::MF_READ |
+                  sys::Memory::ProtectionFlags::MF_WRITE)) {
+        AllErr = joinErrors(std::move(AllErr), errorCodeToError(EC));
+      }
+
       Allocations.erase(Base);
     }
   }
 
   OnDeinitialized(std::move(AllErr));
 }
 
 void InProcessMemoryMapper::release(ArrayRef<ExecutorAddr> Bases,

llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp

TEST(MapperJITLinkMemoryManagerTest, InProcess) {
   const char *TargetMem1 = TargetAddr1.toPtr<const char *>();
   StringRef TargetHello1(TargetMem1, Hello.size());
   EXPECT_EQ(Hello, TargetHello1);
 
   const char *TargetMem2 = TargetAddr2.toPtr<const char *>();
   StringRef TargetHello2(TargetMem2, Hello.size());
   EXPECT_EQ(Hello, TargetHello2);
 
+  EXPECT_EQ(Counter->DeinitCount, 0);
+
   auto Err2 = MemMgr->deallocate(std::move(*FA1));
   EXPECT_THAT_ERROR(std::move(Err2), Succeeded());
 
+  EXPECT_EQ(Counter->DeinitCount, 1);
+
   auto Err3 = MemMgr->deallocate(std::move(*FA2));
   EXPECT_THAT_ERROR(std::move(Err3), Succeeded());
+
+  EXPECT_EQ(Counter->DeinitCount, 2);
 }
 
 } // namespace