Skip to content

Commit a830a77

Browse files
E1ppE1pp
committed
YT-12720: Capture cancelation error in some new cases
commit_hash:db414aec089ec1f05bb59c309b763c4e403bae9e
1 parent 2cfac96 commit a830a77

12 files changed

+878
-42
lines changed

‎yt/yt/core/actions/cancelable_context.cpp

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,19 +25,29 @@ class TCancelableContext::TCancelableInvoker
2525
void Invoke(TClosure callback) override
2626
{
2727
YT_ASSERT(callback);
28+
auto guard = NDetail::MakeCancelableContextCurrentTokenGuard(Context_);
2829

2930
if (Context_->Canceled_) {
31+
callback.Reset();
3032
return;
3133
}
3234

33-
return UnderlyingInvoker_->Invoke(BIND_NO_PROPAGATE([this, this_ = MakeStrong(this), callback = std::move(callback)] {
34-
if (Context_->Canceled_) {
35-
return;
36-
}
37-
38-
TCurrentInvokerGuard guard(this);
39-
callback();
40-
}));
35+
return UnderlyingInvoker_->Invoke(BIND_NO_PROPAGATE(
36+
[
37+
this,
38+
this_ = MakeStrong(this),
39+
callback = std::move(callback)
40+
] () mutable {
41+
auto currentTokenGuard = NDetail::MakeCancelableContextCurrentTokenGuard(Context_);
42+
43+
if (Context_->Canceled_) {
44+
callback.Reset();
45+
return;
46+
}
47+
48+
TCurrentInvokerGuard guard(this);
49+
callback();
50+
}));
4151
}
4252

4353
private:
@@ -52,6 +62,11 @@ bool TCancelableContext::IsCanceled() const
5262
return Canceled_;
5363
}
5464

65+
const TError& TCancelableContext::GetCancelationError() const
66+
{
67+
return CancelationError_;
68+
}
69+
5570
void TCancelableContext::Cancel(const TError& error)
5671
{
5772
THashSet<TWeakPtr<TCancelableContext>> propagateToContexts;
@@ -70,8 +85,7 @@ void TCancelableContext::Cancel(const TError& error)
7085
Handlers_.FireAndClear(error);
7186

7287
for (const auto& weakContext : propagateToContexts) {
73-
auto context = weakContext.Lock();
74-
if (context) {
88+
if (auto context = weakContext.Lock()) {
7589
context->Cancel(error);
7690
}
7791
}

‎yt/yt/core/actions/cancelable_context.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@ class TCancelableContext
2323
//! Returns |true| iff the context is canceled.
2424
bool IsCanceled() const;
2525

26+
//! Only safe to use after IsCanceled returned |true|.
27+
const TError& GetCancelationError() const;
28+
2629
//! Marks the context as canceled raising the handlers
2730
//! and propagates cancelation.
2831
void Cancel(const TError& error);

‎yt/yt/core/actions/cancelation_token-inl.h

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
#ifndef CANCELATION_TOKEN_INL_H_
2+
#error "Direct inclusion of this file is not allowed, include cancelation_token.h"
3+
// For the sake of sane code completion.
4+
#include "cancelation_token.h"
5+
#endif
6+
7+
namespace NYT::NDetail {
8+
9+
////////////////////////////////////////////////////////////////////////////////
10+
11+
template <class TDecayedConcrete>
12+
constexpr TDecayedConcrete& TAnyCancelationToken::TStorage::AsConcrete() & noexcept
13+
{
14+
using TPtr = TDecayedConcrete*;
15+
16+
if constexpr (SmallToken<TDecayedConcrete>) {
17+
return *std::launder(reinterpret_cast<TPtr>(&Storage_));
18+
} else {
19+
return *static_cast<TPtr>(*std::launder(reinterpret_cast<void**>(&Storage_)));
20+
}
21+
}
22+
23+
template <class TDecayedConcrete>
24+
constexpr const TDecayedConcrete& TAnyCancelationToken::TStorage::AsConcrete() const & noexcept
25+
{
26+
using TPtr = const TDecayedConcrete*;
27+
28+
if constexpr (SmallToken<TDecayedConcrete>) {
29+
return *std::launder(reinterpret_cast<TPtr>(&Storage_));
30+
} else {
31+
return *static_cast<TPtr>(*std::launder(reinterpret_cast<void const* const*>(&Storage_)));
32+
}
33+
}
34+
35+
template <class TDecayedConcrete>
36+
constexpr TDecayedConcrete&& TAnyCancelationToken::TStorage::AsConcrete() && noexcept
37+
{
38+
using TPtr = TDecayedConcrete*;
39+
40+
if constexpr (SmallToken<TDecayedConcrete>) {
41+
return std::move(*std::launder(reinterpret_cast<TPtr>(&Storage_)));
42+
} else {
43+
return std::move(*static_cast<TPtr>(*std::launder(reinterpret_cast<void**>(&Storage_))));
44+
}
45+
}
46+
47+
////////////////////////////////////////////////////////////////////////////////
48+
49+
template <CCancelationToken TDecayedConcrete>
50+
TAnyCancelationToken::TVTable TAnyCancelationToken::TVTable::Create() noexcept
51+
{
52+
TVTable table = {};
53+
table.Dtor_ = +[] (TStorage& what) {
54+
TAlloc allocator = {};
55+
56+
auto* ptr = &what.template AsConcrete<TDecayedConcrete>();
57+
std::destroy_at(ptr);
58+
59+
if constexpr (!SmallToken<TDecayedConcrete>) {
60+
TTraits::deallocate(allocator, reinterpret_cast<std::byte*>(ptr), sizeof(TDecayedConcrete));
61+
}
62+
};
63+
64+
table.CopyCtor_ = +[] (TStorage& where, const TStorage& what) -> void {
65+
TAlloc allocator = {};
66+
67+
if constexpr (SmallToken<TDecayedConcrete>) {
68+
where.Set();
69+
} else {
70+
where.Set(TTraits::allocate(allocator, sizeof(TDecayedConcrete)));
71+
}
72+
73+
TTraits::template construct<TDecayedConcrete>(
74+
allocator,
75+
&where.template AsConcrete<TDecayedConcrete>(),
76+
what.template AsConcrete<TDecayedConcrete>());
77+
};
78+
79+
table.MoveCtor_ = +[] (TStorage& where, TStorage&& what) -> void {
80+
if constexpr (SmallToken<TDecayedConcrete>) {
81+
TAlloc allocator = {};
82+
83+
where.Set();
84+
85+
TTraits::template construct<TDecayedConcrete>(
86+
allocator,
87+
&where.template AsConcrete<TDecayedConcrete>(),
88+
std::move(what).template AsConcrete<TDecayedConcrete>());
89+
TTraits::template destroy<TDecayedConcrete>(
90+
allocator,
91+
&what.template AsConcrete<TDecayedConcrete>());
92+
} else {
93+
where.Set(static_cast<void*>(&what));
94+
}
95+
};
96+
97+
table.IsCancelationRequested_ = +[] (const TStorage& what) -> bool {
98+
return what.template AsConcrete<TDecayedConcrete>().IsCancelationRequested();
99+
};
100+
101+
table.CancellationError_ = +[] (const TStorage& what) -> const TError& {
102+
return what.template AsConcrete<TDecayedConcrete>().GetCancelationError();
103+
};
104+
105+
return table;
106+
}
107+
108+
////////////////////////////////////////////////////////////////////////////////
109+
110+
template <class TToken>
111+
requires (!std::same_as<TAnyCancelationToken, std::remove_cvref_t<TToken>> &&
112+
CCancelationToken<std::remove_cvref_t<TToken>>)
113+
TAnyCancelationToken::TAnyCancelationToken(TToken&& token)
114+
{
115+
Set<TToken>(std::forward<TToken>(token));
116+
}
117+
118+
template <class TToken>
119+
void TAnyCancelationToken::Set(TToken&& token)
120+
{
121+
using TDecayed = std::remove_cvref_t<TToken>;
122+
TAlloc allocator = {};
123+
124+
Reset();
125+
126+
VTable_ = &StaticTable<TDecayed>;
127+
128+
if constexpr (SmallToken<TDecayed>) {
129+
Storage_.Set();
130+
} else {
131+
Storage_.Set(TTraits::allocate(allocator, sizeof(TDecayed)));
132+
}
133+
134+
TTraits::template construct<TDecayed>(
135+
allocator,
136+
&Storage_.template AsConcrete<TDecayed>(),
137+
std::forward<TToken>(token));
138+
}
139+
140+
////////////////////////////////////////////////////////////////////////////////
141+
142+
} // namespace NYT::NDetail

0 commit comments

Comments
 (0)