Cancellable async_resolve

scratch.cpp

//
// Copyright (c) 2021 Richard Hodges (hodges.r@gmail.com)
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//

#include <asio.hpp>
#include <iostream>
#include <asio/experimental/append.hpp>

template<class Handler>
struct resolve_state
: std::enable_shared_from_this<resolve_state<Handler>>
{
  resolve_state(Handler handler)
  : handler(std::move(handler))
  {}

  using executor_type = asio::associated_executor_t<Handler>;

  executor_type get_executor()
  {
    return asio::get_associated_executor(handler);
  }

  void run(std::string_view host, std::string_view service)
  {
    resolver.async_resolve(
      host,
      service,
      asio::bind_executor(
        get_executor(),
        [self = this->shared_from_this()](std::error_code ec, asio::ip::tcp::resolver::results_type results) mutable
        {
          if (!std::exchange(self->handled, true))
          {
            auto h = std::move(self->handler);
            self.reset();
            h(ec, results);
          }
        }));

    auto slot = asio::get_associated_cancellation_slot(handler);
    slot.assign([weak = this->weak_from_this()](asio::cancellation_type type){
      if (auto self = weak.lock())
      {
        if (!std::exchange(self->handled, true))
        {
          auto h = std::move(self->handler);
          self.reset();
          asio::post(
            asio::experimental::append(
              std::move(h),
              std::error_code(asio::error::operation_aborted),
              asio::ip::tcp::resolver::results_type {}));
        }
      }
    });
  }

  Handler handler;
  asio::ip::tcp::resolver resolver { get_executor() };
  std::error_code ec = {};
  asio::ip::tcp::resolver::results_type results = {};
  bool handled = false;
};

template<class CompletionHandler>
auto
async_cancellable_resolve(std::string host, std::string service,
                          CompletionHandler &&token)
{
  return asio::async_initiate<CompletionHandler, void(std::error_code, asio::ip::tcp::resolver::results_type)>(
  [host, service](auto &&handler) mutable {
    using Handler = std::decay_t<decltype(handler)>;
    using State = resolve_state<Handler>;
    auto pstate = std::make_shared<State>(std::forward<Handler>(handler));
    pstate->run(host, service);
  }, token);
}


int
main()
{
  using std::error_code;
  using namespace std::literals;

  auto then = std::chrono::system_clock::now();
  asio::io_context ioc;
  asio::cancellation_signal sig;
  async_cancellable_resolve("www.facebook.com", "http",
                  asio::bind_cancellation_slot(sig.slot(),
                                               [then](error_code ec,
                                                      asio::ip::tcp::resolver::results_type results) {
                                                 auto now = std::chrono::system_clock::now();
                                                 using milli_d = std::chrono::duration<double, std::milli>;
                                                 auto diff = std::chrono::duration_cast<milli_d>(now - then).count();
                                                 std::cout << "resolved in " << diff << "ms : error = " << ec
                                                           << ", results: ";
                                                 for (auto &r: results)
                                                 {
                                                   std::cout << r.endpoint() << '\n';
                                                 }
                                                 std::cout << '\n';
                                               }));

  auto t = asio::steady_timer(ioc, 1us);
  t.async_wait([&sig](error_code) {
    sig.emit(asio::cancellation_type::all);
  });

  ioc.run();
}