Add external canceller for TEvPut requests (#30553)

serbel324 · web-flow · commit 3e835ba0ab8c · 2025-12-12T16:19:03.000+03:00
diff --git a/ydb/core/base/blobstorage.h b/ydb/core/base/blobstorage.h
@@ -39,6 +39,9 @@ static constexpr ui64 MaxCollectGarbageFlagsPerMessage = 10000;
 static constexpr TDuration VDiskCooldownTimeout = TDuration::Seconds(15);
 static constexpr TDuration VDiskCooldownTimeoutOnProxy = TDuration::Seconds(12);
 
+struct TMessageRelevanceTracker {};
+using TMessageRelevanceOwner = std::shared_ptr<TMessageRelevanceTracker>;
+using TMessageRelevanceWatcher = std::weak_ptr<TMessageRelevanceTracker>;
 
 struct TStorageStatusFlags {
     ui32 Raw = 0;
@@ -1071,6 +1074,7 @@ struct TEvBlobStorage {
         const bool IgnoreBlock = false;
         mutable NLWTrace::TOrbit Orbit;
         std::vector<std::pair<ui64, ui32>> ExtraBlockChecks; // (TabletId, Generation) pairs
+        std::optional<TMessageRelevanceWatcher> ExternalRelevanceWatcher;
 
         TEvPut(TCloneEventPolicy, const TEvPut& origin)
             : Id(origin.Id)
@@ -1081,18 +1085,21 @@ struct TEvBlobStorage {
             , IssueKeepFlag(origin.IssueKeepFlag)
             , IgnoreBlock(origin.IgnoreBlock)
             , ExtraBlockChecks(origin.ExtraBlockChecks)
+            , ExternalRelevanceWatcher(origin.ExternalRelevanceWatcher)
         {}
 
         TEvPut(const TLogoBlobID &id, TRope &&buffer, TInstant deadline,
                NKikimrBlobStorage::EPutHandleClass handleClass = NKikimrBlobStorage::TabletLog,
-               ETactic tactic = TacticDefault, bool issueKeepFlag = false, bool ignoreBlock = false)
+               ETactic tactic = TacticDefault, bool issueKeepFlag = false, bool ignoreBlock = false,
+               std::optional<TMessageRelevanceWatcher> externalRelevanceWatcher = std::nullopt)
             : Id(id)
             , Buffer(std::move(buffer))
             , Deadline(deadline)
             , HandleClass(handleClass)
             , Tactic(tactic)
             , IssueKeepFlag(issueKeepFlag)
             , IgnoreBlock(ignoreBlock)
+            , ExternalRelevanceWatcher(externalRelevanceWatcher)
         {
             Y_ABORT_UNLESS(Id, "EvPut invalid: LogoBlobId must have non-zero tablet field, id# %s", Id.ToString().c_str());
             Y_ABORT_UNLESS(Buffer.size() < (40 * 1024 * 1024),
@@ -1110,21 +1117,27 @@ struct TEvBlobStorage {
 
         TEvPut(const TLogoBlobID &id, TRcBuf &&buffer, TInstant deadline,
                NKikimrBlobStorage::EPutHandleClass handleClass = NKikimrBlobStorage::TabletLog,
-               ETactic tactic = TacticDefault, bool issueKeepFlag = false)
-            : TEvPut(id, TRope(std::move(buffer)), deadline, handleClass, tactic, issueKeepFlag)
+               ETactic tactic = TacticDefault, bool issueKeepFlag = false,
+               std::optional<TMessageRelevanceWatcher> externalRelevanceWatcher = std::nullopt)
+            : TEvPut(id, TRope(std::move(buffer)), deadline, handleClass, tactic, issueKeepFlag,
+                    /*ignoreBlock=*/false, std::move(externalRelevanceWatcher))
         {}
 
         TEvPut(const TLogoBlobID &id, const TString &buffer, TInstant deadline,
                NKikimrBlobStorage::EPutHandleClass handleClass = NKikimrBlobStorage::TabletLog,
-               ETactic tactic = TacticDefault, bool issueKeepFlag = false)
-            : TEvPut(id, TRope(buffer), deadline, handleClass, tactic, issueKeepFlag)
+               ETactic tactic = TacticDefault, bool issueKeepFlag = false,
+               std::optional<TMessageRelevanceWatcher> externalRelevanceWatcher = std::nullopt)
+            : TEvPut(id, TRope(buffer), deadline, handleClass, tactic, issueKeepFlag,
+                    /*ignoreBlock=*/false, std::move(externalRelevanceWatcher))
         {}
 
 
         TEvPut(const TLogoBlobID &id, const TSharedData &buffer, TInstant deadline,
                NKikimrBlobStorage::EPutHandleClass handleClass = NKikimrBlobStorage::TabletLog,
-               ETactic tactic = TacticDefault, bool issueKeepFlag = false)
-            : TEvPut(id, TRope(buffer), deadline, handleClass, tactic, issueKeepFlag)
+               ETactic tactic = TacticDefault, bool issueKeepFlag = false,
+               std::optional<TMessageRelevanceWatcher> externalRelevanceWatcher = std::nullopt)
+            : TEvPut(id, TRope(buffer), deadline, handleClass, tactic, issueKeepFlag,
+                    /*ignoreBlock=*/false, std::move(externalRelevanceWatcher))
         {}
 
         TString Print(bool isFull) const {
diff --git a/ydb/core/blobstorage/dsproxy/dsproxy.h b/ydb/core/blobstorage/dsproxy/dsproxy.h
@@ -38,6 +38,7 @@ const ui64 UnconfiguredBufferSizeLimit = 32 << 20;
 const TDuration ProxyEstablishSessionsTimeout = TDuration::Seconds(5);
 
 const TDuration DsMinimumDelayBetweenPutWakeups = TDuration::Seconds(1);
+const TDuration DsMaximumDelayBetweenPutWakeups = TDuration::Seconds(60);
 
 const ui64 BufferSizeThreshold = 1 << 20;
 
@@ -195,6 +196,8 @@ class TBlobStorageGroupRequestActor : public TActor<TBlobStorageGroupRequestActo
 
         std::optional<ui32> ForceGroupGeneration; // work only with this specific group generation and nothing else
         bool DoSendDeathNote = true; // unschedules DSProxy timeout on termination, be careful with disabling
+
+        std::optional<TMessageRelevanceWatcher> ExternalRelevanceWatcher = std::nullopt;
     };
 
     struct TTypeSpecificParameters {
@@ -237,6 +240,12 @@ class TBlobStorageGroupRequestActor : public TActor<TBlobStorageGroupRequestActo
         }
 
         Y_ABORT_UNLESS(CostModel);
+
+        if (params.Common.ExternalRelevanceWatcher) {
+            Relevance = std::move(*params.Common.ExternalRelevanceWatcher);
+        } else {
+            Relevance = std::make_shared<TMessageRelevanceTracker>();
+        }
     }
 
     virtual ~TBlobStorageGroupRequestActor() = default;
@@ -290,6 +299,8 @@ class TBlobStorageGroupRequestActor : public TActor<TBlobStorageGroupRequestActo
     static double GetTotalTimeMs(const NKikimrBlobStorage::TTimestamps& timestamps);
     static double GetVDiskTimeMs(const NKikimrBlobStorage::TTimestamps& timestamps);
 
+    bool CheckForExternalCancellation();
+
 private:
     void CheckPostponedQueue();
 
@@ -315,7 +326,7 @@ class TBlobStorageGroupRequestActor : public TActor<TBlobStorageGroupRequestActo
 private:
     const TActorId Source;
     const ui64 Cookie;
-    std::shared_ptr<TMessageRelevanceTracker> MessageRelevanceTracker = std::make_shared<TMessageRelevanceTracker>();
+    std::variant<TMessageRelevanceOwner, TMessageRelevanceWatcher> Relevance;
     ui32 RequestsInFlight = 0;
     std::unique_ptr<IEventBase> Response;
     const TMaybe<TGroupStat::EKind> LatencyQueueKind;
diff --git a/ydb/core/blobstorage/dsproxy/dsproxy_put.cpp b/ydb/core/blobstorage/dsproxy/dsproxy_put.cpp
@@ -311,6 +311,10 @@ class TBlobStorageGroupPutRequest : public TBlobStorageGroupRequestActor {
                 GetVDiskTimeMs(record.GetTimestamps()));
         }
 
+        if (CheckForExternalCancellation()) {
+            return;
+        }
+
         if (status == NKikimrProto::BLOCKED || status == NKikimrProto::DEADLINE) {
             TString error = TStringBuilder() << "Got VPutResult status# " << status << " from VDiskId# " << vdiskId;
             TPutImpl::TPutResultVec putResults;
@@ -387,6 +391,10 @@ class TBlobStorageGroupPutRequest : public TBlobStorageGroupRequestActor {
             }
         }
 
+        if (CheckForExternalCancellation()) {
+            return;
+        }
+
         // Handle put results
         bool isCauseRegistered = !RootCauseTrack.IsOn;
         TPutImpl::TPutResultVec putResults;
@@ -736,6 +744,10 @@ class TBlobStorageGroupPutRequest : public TBlobStorageGroupRequestActor {
             << " Not answered in "
             << (TActivationContext::Monotonic() - RequestStartTime) << " seconds");
 
+        if (CheckForExternalCancellation()) {
+            return;
+        }
+
         const TInstant now = TActivationContext::Now();
         while (!PutDeadlineMasks.empty()) {
             auto [deadline, mask] = *PutDeadlineMasks.begin();
@@ -841,7 +853,7 @@ class TBlobStorageGroupPutRequest : public TBlobStorageGroupRequestActor {
         }
 
         if (deadline != TInstant::Max()) {
-            Schedule(deadline, new TKikimrEvents::TEvWakeup);
+            Schedule(std::min(now + DsMaximumDelayBetweenPutWakeups, deadline), new TKikimrEvents::TEvWakeup);
         }
     }
 
diff --git a/ydb/core/blobstorage/dsproxy/dsproxy_request.cpp b/ydb/core/blobstorage/dsproxy/dsproxy_request.cpp
@@ -245,6 +245,7 @@ namespace NKikimr {
                         .ExecutionRelay = ev->Get()->ExecutionRelay,
                         .LatencyQueueKind = kind,
                         .ForceGroupGeneration = ev->Get()->ForceGroupGeneration,
+                        .ExternalRelevanceWatcher = ev->Get()->ExternalRelevanceWatcher,
                     },
                     .TimeStatsEnabled = Mon->TimeStats.IsEnabled(),
                     .Stats = PerDiskStats,
@@ -585,6 +586,7 @@ namespace NKikimr {
                                     .ExecutionRelay = ev->Get()->ExecutionRelay,
                                     .LatencyQueueKind = kind,
                                     .ForceGroupGeneration = forceGroupGeneration,
+                                    .ExternalRelevanceWatcher = ev->Get()->ExternalRelevanceWatcher,
                                 },
                                 .TimeStatsEnabled = Mon->TimeStats.IsEnabled(),
                                 .Stats = PerDiskStats,
@@ -1064,7 +1066,8 @@ namespace NKikimr {
 
             if constexpr (!std::is_same_v<T, TEvBlobStorage::TEvVStatus> &&
                     !std::is_same_v<T, TEvBlobStorage::TEvVAssimilate>) {
-                ev.MessageRelevanceTracker = MessageRelevanceTracker;
+                std::visit([&](const auto& relevance) { ev.MessageRelevanceTracker = relevance; },
+                        Relevance);
                 ui64 cost;
                 if constexpr (std::is_same_v<T, TEvBlobStorage::TEvVMultiPut>) {
                     bool internalQueue;
@@ -1128,6 +1131,20 @@ namespace NKikimr {
         return true;
     }
 
+    bool TBlobStorageGroupRequestActor::CheckForExternalCancellation() {
+        bool cancelled = false;
+        std::visit(
+            TOverloaded{
+                [&](const TMessageRelevanceOwner&) { cancelled = false; },
+                [&](const TMessageRelevanceWatcher& watcher) { cancelled = watcher.expired(); }
+            }, Relevance);
+
+        if (cancelled) {
+            ReplyAndDie(NKikimrProto::ERROR);
+        }
+        return cancelled;
+    }
+
     void TBlobStorageGroupProxy::Handle(TEvGetQueuesInfo::TPtr ev) {
         ui32 groupSize = Info->GetTotalVDisksNum();
         std::unique_ptr<TEvQueuesInfo> res = std::make_unique<TEvQueuesInfo>(groupSize);
diff --git a/ydb/core/blobstorage/ut_blobstorage/cancellation.cpp b/ydb/core/blobstorage/ut_blobstorage/cancellation.cpp
@@ -0,0 +1,51 @@
+#include <ydb/core/blobstorage/ut_blobstorage/lib/env.h>
+#include <ydb/core/blobstorage/ut_blobstorage/lib/common.h>
+
+#include <util/stream/null.h>
+
+#include <ydb/core/blobstorage/ut_blobstorage/lib/ut_helpers.h>
+
+#define Ctest Cnull
+
+Y_UNIT_TEST_SUITE(Cancellation) {
+    using TTestCtx = TTestCtxBase;
+
+    Y_UNIT_TEST(CancelPut) {
+        TBlobStorageGroupType erasure = TBlobStorageGroupType::Erasure4Plus2Block;
+        TTestCtx ctx({
+            .NodeCount = erasure.BlobSubgroupSize() + 1,
+            .Erasure = erasure,
+            .MaxPutTimeoutDSProxy = TDuration::Minutes(10),
+        });
+
+        ctx.Initialize();
+        ctx.AllocateEdgeActorOnSpecificNode(ctx.NodeCount);
+
+        TMessageRelevanceOwner owner = std::make_shared<TMessageRelevanceTracker>();
+
+        ctx.Env->Runtime->FilterFunction = [&](ui32, std::unique_ptr<IEventHandle>& ev) -> bool {
+            if (ev->GetTypeRewrite() == TEvBlobStorage::TEvVPut::EventType &&
+                    ev->Sender.NodeId() != ev->Recipient.NodeId()) {
+                UNIT_FAIL("Unexpected TEvVPut event# " << ev->Sender << "->" <<
+                        ev->Recipient << " " << ev->ToString());
+            }
+            return true;
+        };
+
+        {
+            TString data = MakeData(10);
+            TLogoBlobID blobId(1, 1, 1, 1, data.size(), 1);
+            TEvBlobStorage::TEvPut* ev = new TEvBlobStorage::TEvPut(blobId, data, TInstant::Max(),
+                    NKikimrBlobStorage::TabletLog, TEvBlobStorage::TEvPut::TacticDefault,
+                    false, owner);
+            owner.reset();
+            ctx.Env->Runtime->WrapInActorContext(ctx.Edge, [&] {
+                SendToBSProxy(ctx.Edge, ctx.GroupId, ev);
+            });
+        }
+
+        auto res = ctx.Env->WaitForEdgeActorEvent<TEvBlobStorage::TEvPutResult>(ctx.Edge, false,
+                TAppData::TimeProvider->Now() + TDuration::Seconds(65));
+        UNIT_ASSERT(res);
+    }
+}
diff --git a/ydb/core/blobstorage/ut_blobstorage/ya.make b/ydb/core/blobstorage/ut_blobstorage/ya.make
@@ -17,6 +17,7 @@ SRCS(
     backpressure.cpp
     block_race.cpp
     bsc_cache.cpp
+    cancellation.cpp
     counting_events.cpp
     corrupted_reads.cpp
     deadlines.cpp
diff --git a/ydb/core/blobstorage/vdisk/common/vdisk_events.h b/ydb/core/blobstorage/vdisk/common/vdisk_events.h
@@ -569,10 +569,8 @@ namespace NKikimr {
     // TEvVPut
     ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
-    struct TMessageRelevanceTracker {};
-
     struct TEventWithRelevanceTracker {
-        std::optional<std::weak_ptr<TMessageRelevanceTracker>> MessageRelevanceTracker;
+        std::optional<TMessageRelevanceWatcher> MessageRelevanceTracker;
     };
 
     struct TEvBlobStorage::TEvVPut

Original file line number	Diff line number	Diff line change
`@@ -311,6 +311,10 @@ class TBlobStorageGroupPutRequest : public TBlobStorageGroupRequestActor {`
`311`	`311`	`GetVDiskTimeMs(record.GetTimestamps()));`
`312`	`312`	`}`
`313`	`313`
	`314`	`+ if (CheckForExternalCancellation()) {`
	`315`	`+ return;`
	`316`	`+ }`
	`317`	`+`
`314`	`318`	`if (status == NKikimrProto::BLOCKED \|\| status == NKikimrProto::DEADLINE) {`
`315`	`319`	`TString error = TStringBuilder() << "Got VPutResult status# " << status << " from VDiskId# " << vdiskId;`
`316`	`320`	`TPutImpl::TPutResultVec putResults;`
`@@ -387,6 +391,10 @@ class TBlobStorageGroupPutRequest : public TBlobStorageGroupRequestActor {`
`387`	`391`	`}`
`388`	`392`	`}`
`389`	`393`
	`394`	`+ if (CheckForExternalCancellation()) {`
	`395`	`+ return;`
	`396`	`+ }`
	`397`	`+`
`390`	`398`	`// Handle put results`
`391`	`399`	`bool isCauseRegistered = !RootCauseTrack.IsOn;`
`392`	`400`	`TPutImpl::TPutResultVec putResults;`
`@@ -736,6 +744,10 @@ class TBlobStorageGroupPutRequest : public TBlobStorageGroupRequestActor {`
`736`	`744`	`<< " Not answered in "`
`737`	`745`	`<< (TActivationContext::Monotonic() - RequestStartTime) << " seconds");`
`738`	`746`
	`747`	`+ if (CheckForExternalCancellation()) {`
	`748`	`+ return;`
	`749`	`+ }`
	`750`	`+`
`739`	`751`	`const TInstant now = TActivationContext::Now();`
`740`	`752`	`while (!PutDeadlineMasks.empty()) {`
`741`	`753`	`auto [deadline, mask] = *PutDeadlineMasks.begin();`
`@@ -841,7 +853,7 @@ class TBlobStorageGroupPutRequest : public TBlobStorageGroupRequestActor {`
`841`	`853`	`}`
`842`	`854`
`843`	`855`	`if (deadline != TInstant::Max()) {`
`844`		`- Schedule(deadline, new TKikimrEvents::TEvWakeup);`
	`856`	`+ Schedule(std::min(now + DsMaximumDelayBetweenPutWakeups, deadline), new TKikimrEvents::TEvWakeup);`
`845`	`857`	`}`
`846`	`858`	`}`
`847`	`859`