orbis: fix dmem::release

kernel/orbis/src/dmem.cpp

@@ -26,6 +26,7 @@ struct DirectMemoryAllocation {
   struct Mapping {
     orbis::Process *process;
     rx::AddressRange vmRange;
+    std::uint64_t dmemOffset;
 
     void serialize(rx::Serializer &s) const {
       s.serialize(process->pid);
@@ -79,6 +80,58 @@ struct DirectMemoryAllocation {
     return result;
   }
 
+  void unmap(orbis::Process *process, rx::AddressRange range) {
+    for (std::size_t i = 0; i < mappings.size(); ++i) {
+      auto &map = mappings[i];
+
+      if (process != nullptr && map.process != process) {
+        continue;
+      }
+
+      auto blockRange = range.intersection(map.vmRange);
+
+      if (!blockRange.isValid()) {
+        continue;
+      }
+
+      if (map.vmRange == blockRange) {
+        if (i != mappings.size() - 1) {
+          std::swap(mappings[i], mappings.back());
+        }
+        mappings.pop_back();
+        --i;
+        continue;
+      }
+
+      if (map.vmRange.beginAddress() == blockRange.beginAddress()) {
+        map.vmRange = rx::AddressRange::fromBeginEnd(blockRange.endAddress(),
+                                                     map.vmRange.endAddress());
+        map.dmemOffset += blockRange.size();
+        continue;
+      }
+
+      if (map.vmRange.endAddress() == blockRange.endAddress()) {
+        map.vmRange = rx::AddressRange::fromBeginEnd(map.vmRange.beginAddress(),
+                                                     blockRange.beginAddress());
+        continue;
+      }
+
+      auto leftAllocation = rx::AddressRange::fromBeginEnd(
+          map.vmRange.beginAddress(), blockRange.beginAddress());
+
+      auto rightAllocation = rx::AddressRange::fromBeginEnd(
+          blockRange.endAddress(), map.vmRange.endAddress());
+
+      map.vmRange = leftAllocation;
+      mappings.push_back({
+          .process = process,
+          .vmRange = rightAllocation,
+          .dmemOffset = map.dmemOffset + (rightAllocation.beginAddress() -
+                                          leftAllocation.beginAddress()),
+      });
+    }
+  }
+
   bool operator==(const DirectMemoryAllocation &) const = default;
 };
 
@@ -326,6 +379,8 @@ orbis::ErrorCode orbis::dmem::release(unsigned dmemIndex,
     return ErrorCode::INVAL;
   }
 
+  orbis::kvector<DirectMemoryAllocation::Mapping> clearMappings;
+  {
     auto dmem = g_dmemPools[dmemIndex];
     std::lock_guard lock(*dmem);
 
@@ -336,22 +391,48 @@ orbis::ErrorCode orbis::dmem::release(unsigned dmemIndex,
       return ErrorCode::OPNOTSUPP;
     }
 
-  auto it = dmem->query(range.beginAddress());
+    auto beginIt = dmem->query(range.beginAddress());
 
-  if (it == dmem->end() || !it->isAllocated()) {
+    if (beginIt == dmem->end()) {
       return ErrorCode::NOENT;
     }
 
-  if (it->isPooled() && !pooled) {
+    auto endIt = beginIt;
+    while (endIt != dmem->end() && endIt.beginAddress() < range.endAddress()) {
+      if (!beginIt->isAllocated()) {
         return ErrorCode::NOENT;
       }
 
-  for (auto mapping : it->mappings) {
+      if (endIt->isPooled() && !pooled) {
-    mapping.process->invoke(
+        return ErrorCode::NOENT;
-        [=] { vmem::unmap(mapping.process, mapping.vmRange); });
       }
 
-  it->mappings.clear();
+      ++endIt;
+    }
+
+    for (auto it = beginIt; it != endIt; ++it) {
+      for (auto &mapping : it->mappings) {
+        auto mapRange = rx::AddressRange::fromBeginSize(mapping.dmemOffset,
+                                                        mapping.vmRange.size());
+        auto releaseRange = mapRange.intersection(range);
+
+        if (!releaseRange.isValid()) {
+          continue;
+        }
+
+        auto releaseVirtualRange = rx::AddressRange::fromBeginSize(
+            mapping.vmRange.beginAddress() +
+                (releaseRange.beginAddress() - mapRange.beginAddress()),
+            releaseRange.size());
+
+        clearMappings.push_back({
+            .process = mapping.process,
+            .vmRange = releaseVirtualRange,
+        });
+      }
+
+      it->unmap(nullptr, range);
+    }
 
     DirectMemoryAllocation allocation{};
     auto result = dmem->map(range.beginAddress(), range.size(), allocation,
@@ -363,6 +444,13 @@ orbis::ErrorCode orbis::dmem::release(unsigned dmemIndex,
 
     dmemDump(dmemIndex, rx::format("released {:x}-{:x}", range.beginAddress(),
                                    range.endAddress()));
+  }
+
+  for (auto mapping : clearMappings) {
+    mapping.process->invoke(
+        [=] { vmem::unmap(mapping.process, mapping.vmRange); });
+  }
+
   return {};
 }
 
@@ -647,41 +735,44 @@ orbis::ErrorCode orbis::dmem::map(orbis::Process *process, unsigned dmemIndex,
     return orbis::ErrorCode::ACCES;
   }
 
-  auto allocationInfoIt = dmem->query(offset);
+  auto beginIt = dmem->query(offset);
 
-  if (allocationInfoIt == dmem->end() || !allocationInfoIt->isAllocated()) {
+  if (beginIt == dmem->end()) {
-    if (allocationInfoIt != dmem->end()) {
-      dmemDump(
-          dmemIndex,
-          rx::format("map unallocated {:x}-{:x}, requested range {:x}-{:x}",
-                     allocationInfoIt.beginAddress(),
-                     allocationInfoIt.endAddress(), range.beginAddress(),
-                     range.endAddress()));
-    } else {
     dmemDump(dmemIndex, rx::format("map out of memory {:x}-{:x}",
                                    range.beginAddress(), range.endAddress()));
-    }
     return orbis::ErrorCode::ACCES;
   }
 
-  if (allocationInfoIt->isPooled()) {
+  auto endIt = beginIt;
+  while (endIt != dmem->end() && endIt.beginAddress() < offset + range.size()) {
+    if (!endIt->isAllocated() || endIt->isPooled()) {
       return orbis::ErrorCode::ACCES;
     }
 
-  if (!vmem::validateMemoryType(allocationInfoIt->getMemoryType(),
+    if (!vmem::validateMemoryType(endIt->getMemoryType(), protection)) {
-                                protection)) {
       return ErrorCode::ACCES;
     }
 
-  // if (!allocationInfoIt->mappings.empty() && !process->allowDmemAliasing) {
+    // if (!endIt->mappings.empty() && !process->allowDmemAliasing) {
     //   return ErrorCode::INVAL;
     // }
 
-  auto directRange = rx::AddressRange::fromBeginSize(offset, range.size())
+    ++endIt;
-                         .intersection(allocationInfoIt.range());
+  }
 
-  if (range.size() > directRange.size()) {
+  if (auto last = endIt; (--last).endAddress() < offset + range.size()) {
-    return orbis::ErrorCode::INVAL;
+    return orbis::ErrorCode::ACCES;
+  }
+
+  for (auto it = beginIt; it != endIt; ++it) {
+    auto itRange = it.range();
+    auto mappingRange = range.intersection(itRange);
+    it->mappings.push_back({
+        .process = process,
+        .vmRange = mappingRange,
+        .dmemOffset =
+            offset + (mappingRange.beginAddress() - itRange.beginAddress()),
+    });
   }
 
   auto physicalRange =
@@ -692,14 +783,8 @@ orbis::ErrorCode orbis::dmem::map(orbis::Process *process, unsigned dmemIndex,
                             vmem::toCpuProtection(protection));
   });
 
-  if (result == ErrorCode{}) {
+  rx::dieIf(result != ErrorCode{}, "failed to map physical memory");
-    allocationInfoIt->mappings.push_back({
+  return {};
-        .process = process,
-        .vmRange = range,
-    });
-  }
-
-  return result;
 }
 
 orbis::ErrorCode orbis::dmem::notifyUnmap(orbis::Process *process,
@@ -713,43 +798,11 @@ orbis::ErrorCode orbis::dmem::notifyUnmap(orbis::Process *process,
   auto dmem = g_dmemPools[dmemIndex];
   std::lock_guard lock(*dmem);
 
-  auto it = dmem->query(offset);
+  auto it = dmem->lowerBound(offset);
-  if (it == dmem->end()) {
-    return ErrorCode::INVAL;
-  }
 
-  for (auto mapIt = it->mappings.begin(); mapIt != it->mappings.end();) {
+  while (it != dmem->end() && it.beginAddress() < offset + range.size()) {
-    if (mapIt->process == process && mapIt->vmRange.intersects(range)) {
+    it->unmap(process, range);
-      auto blockRange = range.intersection(mapIt->vmRange);
+    ++it;
-      if (mapIt->vmRange == blockRange) {
-        mapIt = it->mappings.erase(mapIt);
-        break;
-      }
-
-      if (mapIt->vmRange.beginAddress() == blockRange.beginAddress()) {
-        mapIt->vmRange = rx::AddressRange::fromBeginEnd(
-            blockRange.endAddress(), mapIt->vmRange.endAddress());
-        break;
-      }
-
-      if (mapIt->vmRange.endAddress() == blockRange.endAddress()) {
-        mapIt->vmRange = rx::AddressRange::fromBeginEnd(
-            mapIt->vmRange.beginAddress(), blockRange.beginAddress());
-        break;
-      }
-
-      auto leftAllocation = rx::AddressRange::fromBeginEnd(
-          mapIt->vmRange.beginAddress(), blockRange.beginAddress());
-
-      auto rightAllocation = rx::AddressRange::fromBeginEnd(
-          blockRange.endAddress(), mapIt->vmRange.endAddress());
-
-      mapIt->vmRange = leftAllocation;
-      it->mappings.push_back({.process = process, .vmRange = rightAllocation});
-      break;
-    }
-
-    ++mapIt;
   }
 
   return {};