[promises] Cancellation callback (grpc#31863)

* [promises] Cancellation callback

* Automated change: Fix sanity tests

* Automated change: Fix sanity tests

Co-authored-by: ctiller <ctiller@users.noreply.github.com>

src/core/BUILD

@@ -475,6 +475,19 @@ grpc_cc_library(
     ],
 )
 
+grpc_cc_library(
+    name = "cancel_callback",
+    language = "c++",
+    public_hdrs = [
+        "lib/promise/cancel_callback.h",
+    ],
+    deps = [
+        "poll",
+        "promise_like",
+        "//:gpr_platform",
+    ],
+)
+
 grpc_cc_library(
     name = "promise_factory",
     external_deps = ["absl/meta:type_traits"],

src/core/lib/promise/cancel_callback.h

@@ -0,0 +1,78 @@
+// Copyright 2022 gRPC authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef GRPC_CORE_LIB_PROMISE_CANCEL_CALLBACK_H
+#define GRPC_CORE_LIB_PROMISE_CANCEL_CALLBACK_H
+
+#include <grpc/support/port_platform.h>
+
+#include <utility>
+
+#include "src/core/lib/promise/detail/promise_like.h"
+#include "src/core/lib/promise/poll.h"
+
+namespace grpc_core {
+
+namespace cancel_callback_detail {
+
+template <typename Fn>
+class Handler {
+ public:
+  explicit Handler(Fn fn) : fn_(std::move(fn)) {}
+  Handler(const Handler&) = delete;
+  Handler& operator=(const Handler&) = delete;
+  ~Handler() {
+    if (!done_) {
+      fn_();
+    }
+  }
+  Handler(Handler&& other) noexcept
+      : fn_(std::move(other.fn_)), done_(other.done_) {
+    other.done_ = true;
+  }
+  Handler& operator=(Handler&& other) noexcept {
+    fn_ = std::move(other.fn_);
+    done_ = other.done_;
+    other.done_ = true;
+  }
+
+  void Done() { done_ = true; }
+
+ private:
+  Fn fn_;
+  bool done_ = false;
+};
+
+}  // namespace cancel_callback_detail
+
+// Wrap main_fn so that it calls cancel_fn if the promise is destroyed prior to
+// completion.
+// Returns a promise with the same result type as main_fn.
+template <typename MainFn, typename CancelFn>
+auto OnCancel(MainFn main_fn, CancelFn cancel_fn) {
+  return [on_cancel =
+              cancel_callback_detail::Handler<CancelFn>(std::move(cancel_fn)),
+          main_fn = promise_detail::PromiseLike<MainFn>(
+              std::move(main_fn))]() mutable {
+    auto r = main_fn();
+    if (!absl::holds_alternative<Pending>(r)) {
+      on_cancel.Done();
+    }
+    return r;
+  };
+}
+
+}  // namespace grpc_core
+
+#endif  // GRPC_CORE_LIB_PROMISE_CANCEL_CALLBACK_H

test/core/promise/BUILD

@@ -55,6 +55,16 @@ grpc_cc_test(
     deps = ["//src/core:context"],
 )
 
+grpc_cc_test(
+    name = "cancel_callback_test",
+    srcs = ["cancel_callback_test.cc"],
+    external_deps = ["gtest"],
+    language = "c++",
+    uses_event_engine = False,
+    uses_polling = False,
+    deps = ["//src/core:cancel_callback"],
+)
+
 grpc_cc_test(
     name = "promise_test",
     srcs = ["promise_test.cc"],

test/core/promise/cancel_callback_test.cc

@@ -0,0 +1,41 @@
+// Copyright 2022 gRPC authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/core/lib/promise/cancel_callback.h"
+
+#include "gtest/gtest.h"
+
+namespace grpc_core {
+
+TEST(CancelCallback, DoesntCallCancelIfCompleted) {
+  auto x = OnCancel([]() { return 42; },
+                    []() { FAIL() << "Should never reach here"; });
+  EXPECT_EQ(x(), Poll<int>(42));
+}
+
+TEST(CancelCallback, CallsCancelIfNotCompleted) {
+  bool called = false;
+  {
+    auto x = OnCancel([]() { return 42; }, [&called]() { called = true; });
+    EXPECT_EQ(called, false);
+  }
+  EXPECT_EQ(called, true);
+}
+
+}  // namespace grpc_core
+
+int main(int argc, char** argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}