[Filestore] Validate iovec address ranges in use

Author: antonmyagkov

cloud/filestore/libs/server/server.cpp

@@ -505,7 +505,7 @@ NProto::TError TryAdjustIovecOffsets(
         iovec.SetBase(offset + reinterpret_cast<ui64>(metadata.Address));
     }
 
-    return {};
+    return state->LockAddressRanges(regionId, iovecs);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -763,6 +763,18 @@ class TRequestHandler final
             [=, this] (const auto& response) {
                 Y_UNUSED(response);
 
+                if constexpr (
+                    (std::is_same_v<TRequest, NProto::TWriteDataRequest> ||
+                     std::is_same_v<TRequest, NProto::TReadDataRequest>) &&
+                    std::is_same_v<TAppContext, TFileStoreContext>)
+                {
+                    if (AppCtx.State && Request->IovecsSize() > 0) {
+                        AppCtx.State->ReleaseAddressRanges(
+                            Request->GetRegionId(),
+                            Request->GetIovecs());
+                    }
+                }
+
                 if (AtomicCas(&RequestState, ExecutionCompleted, ExecutingRequest)) {
                     // will be processed on executor thread
                     EnqueueCompletion(ExecCtx.CompletionQueue.get(), tag);

cloud/filestore/libs/server/server_memory_state.cpp

@@ -145,6 +145,59 @@ TResultOrError<TMmapRegionMetadata> TServerState::GetMmapRegion(ui64 mmapId)
     return it->second.ToMetadata();
 }
 
+NProto::TError TServerState::LockAddressRanges(
+    ui64 mmapId,
+    const google::protobuf::RepeatedPtrField<NProto::TIovec>& iovecs)
+{
+    TLightWriteGuard guard(StateLock);
+    auto it = MmapRegions.find(mmapId);
+    if (it == MmapRegions.end()) {
+        return MakeError(
+            E_TRANSPORT_ERROR,
+            Sprintf("Mmap region not found: %lu", mmapId));
+    }
+
+    NKikimr::TIntervalSet<ui64> addressRanges;
+
+    for (const auto& iovec: iovecs) {
+        addressRanges.Add(iovec.GetBase(), iovec.GetBase() + iovec.GetLength());
+    }
+
+    if (addressRanges & it->second.GetLockedAddressRanges()) {
+        return MakeError(
+            E_TRANSPORT_ERROR,
+            Sprintf(
+                "Address ranges from the mmap region %lu are already in use",
+                mmapId));
+    }
+
+    it->second.GetLockedAddressRanges().Add(addressRanges);
+    return {};
+}
+
+NProto::TError TServerState::ReleaseAddressRanges(
+    ui64 mmapId,
+    const google::protobuf::RepeatedPtrField<NProto::TIovec>& iovecs)
+{
+    TLightWriteGuard guard(StateLock);
+    auto it = MmapRegions.find(mmapId);
+    if (it == MmapRegions.end()) {
+        return MakeError(
+            E_TRANSPORT_ERROR,
+            Sprintf("Mmap region not found: %lu", mmapId));
+    }
+
+    NKikimr::TIntervalSet<ui64> addressRanges;
+
+    for (const auto& iovec: iovecs) {
+        it->second.GetLockedAddressRanges().Subtract(
+            iovec.GetBase(),
+            iovec.GetBase() + iovec.GetLength());
+    }
+
+    return {};
+}
+
 NProto::TError TServerState::PingMmapRegion(ui64 mmapId)
 {
     TLightWriteGuard guard(StateLock);

cloud/filestore/libs/server/server_memory_state.h

@@ -11,6 +11,10 @@
 #include <util/system/file.h>
 #include <util/system/spinlock.h>
 
+#include <cloud/filestore/public/api/protos/data.pb.h>
+
+#include <contrib/ydb/core/util/interval_set.h>
+
 #include <tuple>
 #include <unordered_map>
 #include <utility>
@@ -63,8 +67,13 @@ struct TMmapRegion
         Metadata.LatestActivityTimestamp = TInstant::Now();
     }
 
+    NKikimr::TIntervalSet<ui64>& GetLockedAddressRanges() {
+        return LockedAddressRanges;
+    }
+
 private:
     TMmapRegionMetadata Metadata;
+    NKikimr::TIntervalSet<ui64> LockedAddressRanges;
 };
 
 class TServerState
@@ -86,6 +95,14 @@ class TServerState
 
     NProto::TError PingMmapRegion(ui64 mmapId);
 
+    NProto::TError LockAddressRanges(
+        ui64 mmapId,
+        const google::protobuf::RepeatedPtrField<NProto::TIovec>& iovecs);
+
+    NProto::TError ReleaseAddressRanges(
+        ui64 mmapId,
+        const google::protobuf::RepeatedPtrField<NProto::TIovec>& iovecs);
+
 private:
     TLightRWLock StateLock;
     std::unordered_map<ui64, TMmapRegion> MmapRegions;

cloud/filestore/libs/server/server_memory_state_ut.cpp

@@ -278,6 +278,67 @@ Y_UNIT_TEST_SUITE(TServerStateTest)
             getResult.GetResult().LatestActivityTimestamp,
             oldTimestamp);
     }
+
+    Y_UNIT_TEST(ShouldRejectLockingIntersectingRegions)
+    {
+        TTestEnv env;
+        TServerState state(env.GetBasePath());
+
+        const size_t fileSize = 4096;
+        TString relativePath = env.CreateTestFile("test_file.dat", fileSize);
+
+        auto result = state.CreateMmapRegion(relativePath, fileSize);
+        UNIT_ASSERT_C(!HasError(result), FormatError(result.GetError()));
+
+        const auto mmapInfo = result.ExtractResult();
+
+        google::protobuf::RepeatedPtrField<NProto::TIovec> iovecs;
+        auto iovec = iovecs.Add();
+        iovec->SetBase(0);
+        iovec->SetLength(100);
+
+        iovec = iovecs.Add();
+        iovec->SetBase(4000);
+        iovec->SetLength(100);
+
+        result = state.LockAddressRanges(mmapInfo.Id, iovecs);
+        UNIT_ASSERT_C(!HasError(result), FormatError(result.GetError()));
+
+        result = state.LockAddressRanges(mmapInfo.Id, iovecs);
+        UNIT_ASSERT(HasError(result));
+
+        {
+            google::protobuf::RepeatedPtrField<NProto::TIovec> iovecs;
+            auto iovec = iovecs.Add();
+            iovec->SetBase(50);
+            iovec->SetLength(1000);
+
+            result = state.LockAddressRanges(mmapInfo.Id, iovecs);
+            UNIT_ASSERT(HasError(result));
+        }
+
+        {
+            google::protobuf::RepeatedPtrField<NProto::TIovec> iovecs;
+            auto iovec = iovecs.Add();
+            iovec->SetBase(50);
+            iovec->SetLength(10);
+
+            result = state.LockAddressRanges(mmapInfo.Id, iovecs);
+            UNIT_ASSERT(HasError(result));
+        }
+
+        state.ReleaseAddressRanges(mmapInfo.Id, iovecs);
+
+        {
+            google::protobuf::RepeatedPtrField<NProto::TIovec> iovecs;
+            auto iovec = iovecs.Add();
+            iovec->SetBase(50);
+            iovec->SetLength(10);
+
+            result = state.LockAddressRanges(mmapInfo.Id, iovecs);
+            UNIT_ASSERT_C(!HasError(result), FormatError(result.GetError()));
+        }
+    }
 }
 
 }   // namespace NCloud::NFileStore::NServer