feat: support cancellation tokens in Curl REST client and payloads

Author: ajayky-os

google/cloud/internal/curl_http_payload.cc

@@ -35,6 +35,10 @@ std::multimap<std::string, std::string> CurlHttpPayload::DebugHeaders() const {
   return impl_->headers();
 }
 
+void CurlHttpPayload::Cancel() {
+  if (impl_) impl_->Cancel();
+}
+
 StatusOr<std::string> ReadAll(std::unique_ptr<HttpPayload> payload,
                               std::size_t read_size) {
   std::string output_buffer;

google/cloud/internal/curl_http_payload.h

@@ -45,6 +45,8 @@ class CurlHttpPayload : public HttpPayload {
 
   std::multimap<std::string, std::string> DebugHeaders() const override;
 
+  void Cancel() override;
+
  private:
   friend class CurlRestResponse;
   CurlHttpPayload(std::unique_ptr<CurlImpl> impl, Options options);

google/cloud/internal/curl_impl.cc

@@ -125,6 +125,12 @@ static int SeekFunction(  // NOLINT(misc-use-anonymous-namespace)
              : CURL_SEEKFUNC_FAIL;
 }
 
+static int TransferInfoFunction(  // NOLINT(misc-use-anonymous-namespace)
+    void* userdata, curl_off_t, curl_off_t, curl_off_t, curl_off_t) {
+  auto* const request = reinterpret_cast<CurlImpl*>(userdata);
+  return request->TransferInfoCallback();
+}
+
 }  // extern "C"
 
 std::size_t SpillBuffer::CopyFrom(absl::Span<char const> src) {
@@ -244,7 +250,12 @@ CurlImpl::~CurlImpl() {
   CleanupHandles();
 
   CurlHandle::ReturnToPool(*factory_, std::move(handle_));
-  factory_->CleanupMultiHandle(std::move(multi_), HandleDisposition::kKeep);
+  factory_->CleanupMultiHandle(ReleaseMulti(), HandleDisposition::kKeep);
+}
+
+CurlMulti CurlImpl::ReleaseMulti() {
+  std::lock_guard<std::mutex> lk(multi_mu_);
+  return std::move(multi_);
 }
 
 void CurlImpl::SetHeader(HttpHeader header) {
@@ -460,8 +471,6 @@ Status CurlImpl::MakeRequest(HttpMethod method, RestContext& context,
   if (!status.ok()) return OnTransferError(context, std::move(status));
 
   if (method == HttpMethod::kGet) {
-    status = handle_.SetOption(CURLOPT_NOPROGRESS, 1L);
-    if (!status.ok()) return OnTransferError(context, std::move(status));
     if (download_stall_timeout_ != std::chrono::seconds::zero()) {
       // NOLINTNEXTLINE(google-runtime-int) - libcurl *requires* long
       auto const timeout = static_cast<long>(download_stall_timeout_.count());
@@ -553,9 +562,43 @@ StatusOr<std::size_t> CurlImpl::Read(absl::Span<char> output) {
   // This context is discarded.  Any interesting information was already
   // captured when the request was started.
   RestContext context;
+  if (cancellation_token_) {
+    context.set_cancellation_token(cancellation_token_);
+  }
   return ReadImpl(context, std::move(output));
 }
 
+void CurlImpl::SetCancellationToken(
+    std::shared_ptr<std::atomic<bool>> token) {
+  if (token) {
+    if (cancellation_token_->load(std::memory_order_relaxed)) {
+      token->store(true, std::memory_order_relaxed);
+    }
+    cancellation_token_ = std::move(token);
+    cancellable_ = true;
+  }
+}
+
+void CurlImpl::Cancel() {
+  cancellation_token_->store(true, std::memory_order_relaxed);
+#if CURL_AT_LEAST_VERSION(7, 68, 0)
+  // The lock keeps `multi_` alive and owned by this request while the wakeup
+  // is delivered: without it the transfer thread could concurrently return
+  // the handle to the pool, where another request may already be using it.
+  std::lock_guard<std::mutex> lk(multi_mu_);
+  if (multi_) {
+    (void)curl_multi_wakeup(multi_.get());
+  }
+#endif
+}
+
+int CurlImpl::TransferInfoCallback() {
+  if (cancellation_token_->load(std::memory_order_relaxed)) {
+    return 1;
+  }
+  return 0;
+}
+
 std::size_t CurlImpl::WriteCallback(absl::Span<char> response) {
   handle_.FlushDebug(__func__);
   TRACE_STATE() << ", begin"
@@ -624,6 +667,17 @@ std::size_t CurlImpl::HeaderCallback(absl::Span<char> response) {
 Status CurlImpl::MakeRequestImpl(RestContext& context) {
   TRACE_STATE() << ", url_=" << url_;
 
+  if (context.cancellation_token() &&
+      context.cancellation_token() != cancellation_token_) {
+    SetCancellationToken(context.cancellation_token());
+  }
+
+  if (cancellation_token_->load(std::memory_order_relaxed)) {
+    return OnTransferError(
+        context,
+        internal::CancelledError("Request cancelled", GCP_ERROR_INFO()));
+  }
+
   Status status;
   status = handle_.SetOption(CURLOPT_URL, url_.c_str());
   if (!status.ok()) return OnTransferError(context, std::move(status));
@@ -645,6 +699,13 @@ Status CurlImpl::MakeRequestImpl(RestContext& context) {
   handle_.SetOptionUnchecked(CURLOPT_HTTP_VERSION,
                              VersionToCurlCode(http_version_));
 
+  status = handle_.SetOption(CURLOPT_NOPROGRESS, 0L);
+  if (!status.ok()) return OnTransferError(context, std::move(status));
+  status = handle_.SetOption(CURLOPT_XFERINFOFUNCTION, &TransferInfoFunction);
+  if (!status.ok()) return OnTransferError(context, std::move(status));
+  status = handle_.SetOption(CURLOPT_XFERINFODATA, this);
+  if (!status.ok()) return OnTransferError(context, std::move(status));
+
   auto error = curl_multi_add_handle(multi_.get(), handle_.handle_.get());
 
   // This indicates that we are using the API incorrectly. The application
@@ -669,6 +730,17 @@ StatusOr<std::size_t> CurlImpl::ReadImpl(RestContext& context,
   avail_ = output;
   TRACE_STATE() << ", begin";
 
+  if (context.cancellation_token() &&
+      context.cancellation_token() != cancellation_token_) {
+    SetCancellationToken(context.cancellation_token());
+  }
+
+  if (cancellation_token_->load(std::memory_order_relaxed)) {
+    return OnTransferError(
+        context,
+        internal::CancelledError("Request cancelled", GCP_ERROR_INFO()));
+  }
+
   // Before calling WaitForHandles(), move any data from the spill buffer
   // into the output buffer. It is possible that WaitForHandles() will
   // never call WriteCallback() (e.g., because PerformWork() closed the
@@ -807,6 +879,15 @@ StatusOr<int> CurlImpl::PerformWork() {
       // (see above) tells libcurl that it cannot receive more data.
       if (closing_) continue;
       if (multi_info_read_result != CURLE_OK) {
+        // CURLE_ABORTED_BY_CALLBACK maps to `kAborted`, but when the abort
+        // came from TransferInfoCallback() observing the cancellation token
+        // the caller asked for the cancellation: report `kCancelled` so it is
+        // not mistaken for a permanent failure.
+        if (multi_info_read_result == CURLE_ABORTED_BY_CALLBACK &&
+            cancellation_token_->load(std::memory_order_relaxed)) {
+          return internal::CancelledError("Request cancelled",
+                                          GCP_ERROR_INFO());
+        }
         return CurlHandle::AsStatus(multi_info_read_result, __func__);
       }
       if (multi_remove_result != CURLM_OK) {
@@ -822,6 +903,9 @@ Status CurlImpl::PerformWorkUntil(absl::FunctionRef<bool()> predicate) {
   TRACE_STATE() << ", begin";
   int repeats = 0;
   while (!predicate()) {
+    if (cancellation_token_->load(std::memory_order_relaxed)) {
+      return internal::CancelledError("Request cancelled", GCP_ERROR_INFO());
+    }
     handle_.FlushDebug(__func__);
     TRACE_STATE() << ", repeats=" << repeats;
     auto running_handles = PerformWork();
@@ -839,7 +923,14 @@ Status CurlImpl::PerformWorkUntil(absl::FunctionRef<bool()> predicate) {
 }
 
 Status CurlImpl::WaitForHandles(int& repeats) {
+#if !CURL_AT_LEAST_VERSION(7, 68, 0)
+  // Without curl_multi_wakeup() a Cancel() from another thread cannot
+  // interrupt the wait, so poll frequently -- but only when some caller can
+  // actually cancel this transfer; other transfers keep the long timeout.
+  int const timeout_ms = cancellable_ ? 50 : 1000;
+#else
   int const timeout_ms = 1000;
+#endif
   int numfds = 0;
   CURLMcode result;
 #if CURL_AT_LEAST_VERSION(7, 66, 0)
@@ -890,7 +981,7 @@ Status CurlImpl::OnTransferError(RestContext& context, Status status) {
   // While the handle is suspect, there is probably nothing wrong with the
   // CURLM* handle. That just represents a local resource, such as data
   // structures for epoll(7) or select(2).
-  factory_->CleanupMultiHandle(std::move(multi_), HandleDisposition::kKeep);
+  factory_->CleanupMultiHandle(ReleaseMulti(), HandleDisposition::kKeep);
 
   return status;
 }
@@ -903,7 +994,7 @@ void CurlImpl::OnTransferDone() {
   // in PerformWork(). Release the handles back to the factory as soon as
   // possible, so they can be reused for any other requests.
   CurlHandle::ReturnToPool(*factory_, std::move(handle_));
-  factory_->CleanupMultiHandle(std::move(multi_), HandleDisposition::kKeep);
+  factory_->CleanupMultiHandle(ReleaseMulti(), HandleDisposition::kKeep);
 }
 
 std::optional<std::string> CurlOptProxy(Options const& options) {

google/cloud/internal/curl_impl.h

@@ -29,10 +29,12 @@
 #include "google/cloud/version.h"
 #include "absl/types/span.h"
 #include <array>
+#include <atomic>
 #include <chrono>
 #include <cstdint>
 #include <map>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <string>
 #include <vector>
@@ -80,9 +82,11 @@ class CurlImpl {
   ~CurlImpl();
 
   CurlImpl(CurlImpl const&) = delete;
-  CurlImpl(CurlImpl&&) = default;
+  // Not movable: `multi_mu_` synchronizes Cancel(), which may run on a
+  // different thread, with the handoff of `multi_` back to the pool.
+  CurlImpl(CurlImpl&&) = delete;
   CurlImpl& operator=(CurlImpl const&) = delete;
-  CurlImpl& operator=(CurlImpl&&) = default;
+  CurlImpl& operator=(CurlImpl&&) = delete;
 
   void SetHeader(HttpHeader header);
   void SetHeaders(HttpHeaders const& headers);
@@ -105,6 +109,10 @@ class CurlImpl {
   bool HasUnreadData() const;
   StatusOr<std::size_t> Read(absl::Span<char> output);
 
+  void SetCancellationToken(std::shared_ptr<std::atomic<bool>> token);
+  void Cancel();
+  int TransferInfoCallback();
+
   // Called from libcurl callbacks for received data.
   std::size_t WriteCallback(absl::Span<char> response);
   std::size_t HeaderCallback(absl::Span<char> response);
@@ -132,6 +140,10 @@ class CurlImpl {
 
   // Cleanup the CURL handles, leaving them ready for reuse.
   void CleanupHandles();
+  // Take ownership of `multi_` away from this request, synchronizing with any
+  // concurrent Cancel(). Call this instead of `std::move(multi_)` before
+  // returning the handle to the pool.
+  CurlMulti ReleaseMulti();
   // Perform at least part of the request.
   StatusOr<int> PerformWork();
   // Loop on PerformWork until a condition is met.
@@ -146,7 +158,16 @@ class CurlImpl {
   std::vector<HttpHeader> pending_request_headers_;
   CurlHeaders request_headers_;
   CurlHandle handle_;
+  // Guards the handoff of `multi_` back to the pool against a concurrent
+  // Cancel(), which calls `curl_multi_wakeup()` on it from another thread.
+  std::mutex multi_mu_;
   CurlMulti multi_;
+  std::shared_ptr<std::atomic<bool>> cancellation_token_ =
+      std::make_shared<std::atomic<bool>>(false);
+  // True once an external cancellation token was installed, i.e. some other
+  // thread may call Cancel() on this transfer. Only read and written on the
+  // transfer thread.
+  bool cancellable_ = false;
 
   bool logging_enabled_;
   bool follow_location_;

google/cloud/internal/curl_rest_client.cc

@@ -133,6 +133,9 @@ StatusOr<std::unique_ptr<CurlImpl>> CurlRestClient::CreateCurlImpl(
   auto handle = CurlHandle::MakeFromPool(*handle_factory_);
   auto impl = std::make_unique<CurlImpl>(std::move(handle), handle_factory_,
                                          options, pqc_ec_curves_);
+  if (context.cancellation_token()) {
+    impl->SetCancellationToken(context.cancellation_token());
+  }
   if (credentials_) {
     auto auth_headers = credentials_->AuthenticationHeaders(
         std::chrono::system_clock::now(), endpoint_address_);

google/cloud/internal/http_payload.h

@@ -46,6 +46,9 @@ class HttpPayload {
   virtual std::multimap<std::string, std::string> DebugHeaders() const {
     return {};
   }
+
+  /// Cancels an in-progress or subsequent read.
+  virtual void Cancel() {}
 };
 
 // This function makes one or more HttpPayload::Read calls and writes all the

google/cloud/internal/rest_context.h

@@ -18,7 +18,9 @@
 #include "google/cloud/internal/http_header.h"
 #include "google/cloud/options.h"
 #include "google/cloud/version.h"
+#include <atomic>
 #include <chrono>
+#include <memory>
 #include <optional>
 #include <string>
 #include <utility>
@@ -106,10 +108,18 @@ class RestContext {
     appconnect_time_ = us;
   }
 
+  std::shared_ptr<std::atomic<bool>> cancellation_token() const {
+    return cancellation_token_;
+  }
+  void set_cancellation_token(std::shared_ptr<std::atomic<bool>> token) {
+    cancellation_token_ = std::move(token);
+  }
+
  private:
   friend bool operator==(RestContext const& lhs, RestContext const& rhs);
   Options options_;
   HttpHeaders headers_;
+  std::shared_ptr<std::atomic<bool>> cancellation_token_;
   std::optional<std::string> local_ip_address_;
   std::optional<std::int32_t> local_port_;
   std::optional<std::string> primary_ip_address_;

google/cloud/internal/rest_context_test.cc

@@ -78,6 +78,15 @@ TEST_F(RestContextTest, Equality) {
   EXPECT_THAT(lhs, Eq(rhs));
 }
 
+TEST_F(RestContextTest, CancellationToken) {
+  RestContext context;
+  EXPECT_EQ(context.cancellation_token(), nullptr);
+
+  auto token = std::make_shared<std::atomic<bool>>(false);
+  context.set_cancellation_token(token);
+  EXPECT_EQ(context.cancellation_token(), token);
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace rest_internal

google/cloud/internal/tracing_http_payload.h

@@ -35,6 +35,9 @@ class TracingHttpPayload : public HttpPayload {
   bool HasUnreadData() const override;
   StatusOr<std::size_t> Read(absl::Span<char> buffer) override;
   std::multimap<std::string, std::string> DebugHeaders() const override;
+  void Cancel() override {
+    if (impl_) impl_->Cancel();
+  }
 
  private:
   std::unique_ptr<HttpPayload> impl_;

google/cloud/internal/tracing_http_payload_test.cc

@@ -124,6 +124,17 @@ TEST(TracingHttpPayload, Failure) {
           SpanHasEvents(MakeReadMatcher(16, 16), MakeReadMatcher(16)))));
 }
 
+TEST(TracingHttpPayload, Cancel) {
+  auto impl = std::make_unique<MockHttpPayload>();
+  EXPECT_CALL(*impl, Cancel).Times(1);
+
+  RestRequest request("https://example.com/ignored");
+  auto span = MakeSpanHttp(request, "GET");
+
+  TracingHttpPayload payload(std::move(impl), std::move(span));
+  payload.Cancel();
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace rest_internal