[promises] Cancellation callback (#31863)

* [promises] Cancellation callback

* Automated change: Fix sanity tests

* Automated change: Fix sanity tests

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

CMakeLists.txt

@@ -850,6 +850,7 @@ if(gRPC_BUILD_TESTS)
   add_dependencies(buildtests_cxx c_slice_buffer_test)
   add_dependencies(buildtests_cxx call_finalization_test)
   add_dependencies(buildtests_cxx cancel_ares_query_test)
+  add_dependencies(buildtests_cxx cancel_callback_test)
   add_dependencies(buildtests_cxx cel_authorization_engine_test)
   add_dependencies(buildtests_cxx certificate_provider_registry_test)
   add_dependencies(buildtests_cxx certificate_provider_store_test)
@@ -7281,6 +7282,44 @@ target_link_libraries(cancel_ares_query_test
 )
 
 
+endif()
+if(gRPC_BUILD_TESTS)
+
+add_executable(cancel_callback_test
+  test/core/promise/cancel_callback_test.cc
+  third_party/googletest/googletest/src/gtest-all.cc
+  third_party/googletest/googlemock/src/gmock-all.cc
+)
+
+target_include_directories(cancel_callback_test
+  PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}
+    ${CMAKE_CURRENT_SOURCE_DIR}/include
+    ${_gRPC_ADDRESS_SORTING_INCLUDE_DIR}
+    ${_gRPC_RE2_INCLUDE_DIR}
+    ${_gRPC_SSL_INCLUDE_DIR}
+    ${_gRPC_UPB_GENERATED_DIR}
+    ${_gRPC_UPB_GRPC_GENERATED_DIR}
+    ${_gRPC_UPB_INCLUDE_DIR}
+    ${_gRPC_XXHASH_INCLUDE_DIR}
+    ${_gRPC_ZLIB_INCLUDE_DIR}
+    third_party/googletest/googletest/include
+    third_party/googletest/googletest
+    third_party/googletest/googlemock/include
+    third_party/googletest/googlemock
+    ${_gRPC_PROTO_GENS_DIR}
+)
+
+target_link_libraries(cancel_callback_test
+  ${_gRPC_BASELIB_LIBRARIES}
+  ${_gRPC_PROTOBUF_LIBRARIES}
+  ${_gRPC_ZLIB_LIBRARIES}
+  ${_gRPC_ALLTARGETS_LIBRARIES}
+  absl::type_traits
+  absl::variant
+)
+
+
 endif()
 if(gRPC_BUILD_TESTS)
 

build_autogenerated.yaml

@@ -5128,6 +5128,20 @@ targets:
   deps:
   - grpc++_test_config
   - grpc++_test_util
+- name: cancel_callback_test
+  gtest: true
+  build: test
+  language: c++
+  headers:
+  - src/core/lib/promise/cancel_callback.h
+  - src/core/lib/promise/detail/promise_like.h
+  - src/core/lib/promise/poll.h
+  src:
+  - test/core/promise/cancel_callback_test.cc
+  deps:
+  - absl/meta:type_traits
+  - absl/types:variant
+  uses_polling: false
 - name: cel_authorization_engine_test
   gtest: true
   build: test

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();
+}

tools/run_tests/generated/tests.json

@@ -1371,6 +1371,30 @@
     ],
     "uses_polling": true
   },
+  {
+    "args": [],
+    "benchmark": false,
+    "ci_platforms": [
+      "linux",
+      "mac",
+      "posix",
+      "windows"
+    ],
+    "cpu_cost": 1.0,
+    "exclude_configs": [],
+    "exclude_iomgrs": [],
+    "flaky": false,
+    "gtest": true,
+    "language": "c++",
+    "name": "cancel_callback_test",
+    "platforms": [
+      "linux",
+      "mac",
+      "posix",
+      "windows"
+    ],
+    "uses_polling": false
+  },
   {
     "args": [],
     "benchmark": false,