1 // Copyright 2022 gRPC authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 #ifndef GRPC_SRC_CORE_LIB_PROMISE_CANCEL_CALLBACK_H
16 #define GRPC_SRC_CORE_LIB_PROMISE_CANCEL_CALLBACK_H
17
18 #include <grpc/support/port_platform.h>
19
20 #include "src/core/lib/promise/context.h"
21 #include "src/core/lib/promise/detail/promise_like.h"
22 #include "src/core/lib/resource_quota/arena.h"
23
24 namespace grpc_core {
25
26 namespace cancel_callback_detail {
27
28 template <typename Fn>
29 class Handler {
30 public:
Handler(Fn fn)31 GPR_ATTRIBUTE_ALWAYS_INLINE_FUNCTION explicit Handler(Fn fn)
32 : fn_(std::move(fn)) {}
33 Handler(const Handler&) = delete;
34 Handler& operator=(const Handler&) = delete;
~Handler()35 GPR_ATTRIBUTE_ALWAYS_INLINE_FUNCTION ~Handler() {
36 if (!done_) {
37 promise_detail::Context<Arena> ctx(arena_.get());
38 fn_();
39 }
40 }
Handler(Handler && other)41 GPR_ATTRIBUTE_ALWAYS_INLINE_FUNCTION Handler(Handler&& other) noexcept
42 : fn_(std::move(other.fn_)), done_(other.done_) {
43 other.done_ = true;
44 }
45 GPR_ATTRIBUTE_ALWAYS_INLINE_FUNCTION Handler& operator=(
46 Handler&& other) noexcept {
47 fn_ = std::move(other.fn_);
48 done_ = other.done_;
49 other.done_ = true;
50 }
51
Done()52 GPR_ATTRIBUTE_ALWAYS_INLINE_FUNCTION void Done() { done_ = true; }
53
54 private:
55 Fn fn_;
56 // Since cancellation happens at destruction time we need to either capture
57 // context here (via the arena), or make sure that no promise is destructed
58 // without an Arena context on the stack. The latter is an eternal game of
59 // whackamole, so we're choosing the former for now.
60 // TODO(ctiller): re-evaluate at some point in the future.
61 RefCountedPtr<Arena> arena_ =
62 HasContext<Arena>() ? GetContext<Arena>()->Ref() : nullptr;
63 bool done_ = false;
64 };
65
66 } // namespace cancel_callback_detail
67
68 // Wrap main_fn so that it calls cancel_fn if the promise is destroyed prior to
69 // completion.
70 // Returns a promise with the same result type as main_fn.
71 template <typename MainFn, typename CancelFn>
OnCancel(MainFn main_fn,CancelFn cancel_fn)72 GPR_ATTRIBUTE_ALWAYS_INLINE_FUNCTION inline auto OnCancel(MainFn main_fn,
73 CancelFn cancel_fn) {
74 return [on_cancel =
75 cancel_callback_detail::Handler<CancelFn>(std::move(cancel_fn)),
76 main_fn = promise_detail::PromiseLike<MainFn>(
77 std::move(main_fn))]() mutable {
78 auto r = main_fn();
79 if (r.ready()) {
80 on_cancel.Done();
81 }
82 return r;
83 };
84 }
85
86 // Similar to OnCancel, but returns a factory that uses main_fn to construct the
87 // resulting promise. If the factory is dropped without being called, cancel_fn
88 // is called.
89 template <typename MainFn, typename CancelFn>
OnCancelFactory(MainFn main_fn,CancelFn cancel_fn)90 GPR_ATTRIBUTE_ALWAYS_INLINE_FUNCTION inline auto OnCancelFactory(
91 MainFn main_fn, CancelFn cancel_fn) {
92 return [on_cancel =
93 cancel_callback_detail::Handler<CancelFn>(std::move(cancel_fn)),
94 main_fn = std::move(main_fn)]() mutable {
95 auto r = main_fn();
96 on_cancel.Done();
97 return r;
98 };
99 };
100
101 } // namespace grpc_core
102
103 #endif // GRPC_SRC_CORE_LIB_PROMISE_CANCEL_CALLBACK_H
104