src/shrpx_memcached_connection.cc - master - githubmirror/nghttp2

src/shrpx_memcached_connection.cc
788 строк20 KB
Tatsuhiro Tsujikawa
src: MSDMI part2
12 май 2026, 15:11
12 май 2026, 15:11f1aefd0
/*
 * nghttp2 - HTTP/2 C Library
 *
 * Copyright (c) 2015 Tatsuhiro Tsujikawa
 *
 * Permission is hereby granted, free of charge, to any person obtaining
 * a copy of this software and associated documentation files (the
 * "Software"), to deal in the Software without restriction, including
 * without limitation the rights to use, copy, modify, merge, publish,
 * distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so, subject to
 * the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
 * LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
 * OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
 * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */
#include "shrpx_memcached_connection.h"

#include <limits.h>
#include <sys/uio.h>

#include <cerrno>

#include "shrpx_memcached_request.h"
#include "shrpx_memcached_result.h"
#include "shrpx_config.h"
#include "shrpx_tls.h"
#include "shrpx_log.h"
#include "util.h"

namespace shrpx {

namespace {
void timeoutcb(struct ev_loop *loop, ev_timer *w, int revents) {
  auto conn = static_cast<Connection *>(w->data);
  auto mconn = static_cast<MemcachedConnection *>(conn->data);

  if (w == &conn->rt && !conn->expired_rt()) {
    return;
  }

  if (log_enabled(INFO)) {
    Log{INFO, mconn} << "Time out";
  }

  mconn->disconnect();
}
} // namespace

namespace {
void readcb(struct ev_loop *loop, ev_io *w, int revents) {
  auto conn = static_cast<Connection *>(w->data);
  auto mconn = static_cast<MemcachedConnection *>(conn->data);

  if (!mconn->on_read()) {
    mconn->reconnect_or_fail();
    return;
  }
}
} // namespace

namespace {
void writecb(struct ev_loop *loop, ev_io *w, int revents) {
  auto conn = static_cast<Connection *>(w->data);
  auto mconn = static_cast<MemcachedConnection *>(conn->data);

  if (!mconn->on_write()) {
    mconn->reconnect_or_fail();
    return;
  }
}
} // namespace

namespace {
void connectcb(struct ev_loop *loop, ev_io *w, int revents) {
  auto conn = static_cast<Connection *>(w->data);
  auto mconn = static_cast<MemcachedConnection *>(conn->data);

  if (!mconn->connected()) {
    mconn->disconnect();
    return;
  }

  writecb(loop, w, revents);
}
} // namespace

constexpr auto write_timeout = 10_s;
constexpr auto read_timeout = 10_s;

MemcachedConnection::MemcachedConnection(const Address *addr,
                                         struct ev_loop *loop, SSL_CTX *ssl_ctx,
                                         std::string_view sni_name,
                                         MemchunkPool *mcpool,
                                         std::mt19937 &gen)
  : conn_(loop, -1, nullptr, mcpool, write_timeout, read_timeout, {}, {},
          connectcb, readcb, timeoutcb, this, 0, 0., Proto::MEMCACHED),
    sni_name_(sni_name),
    connect_blocker_(
      gen, loop, [] {}, [] {}),
    addr_(addr),
    ssl_ctx_(ssl_ctx) {}

MemcachedConnection::~MemcachedConnection() { conn_.disconnect(); }

namespace {
void clear_request(std::deque<std::unique_ptr<MemcachedRequest>> &q) {
  for (auto &req : q) {
    if (req->cb) {
      req->cb(req.get(),
              MemcachedResult(MemcachedStatusCode::EXT_NETWORK_ERROR));
    }
  }
  q.clear();
}
} // namespace

void MemcachedConnection::disconnect() {
  clear_request(recvq_);
  clear_request(sendq_);

  sendbufv_.clear();
  sendsum_ = 0;

  parse_state_ = {};

  connected_ = false;

  conn_.disconnect();

  assert(recvbuf_.rleft() == 0);
  recvbuf_.reset();

  do_read_ = do_write_ = &MemcachedConnection::noop;
}

std::expected<void, Error> MemcachedConnection::initiate_connection() {
  assert(conn_.fd == -1);

  if (ssl_ctx_) {
    auto maybe_ssl = tls::create_ssl(ssl_ctx_);
    if (!maybe_ssl) {
      return std::unexpected{maybe_ssl.error()};
    }
    conn_.set_ssl(*maybe_ssl);
    conn_.tls.client_session_cache = &tls_session_cache_;
  }

  auto maybe_fd = util::create_nonblock_socket(addr_->family());
  if (!maybe_fd) {
    auto error = errno;
    Log{WARN, this} << "socket() failed; errno=" << error;

    return std::unexpected{maybe_fd.error()};
  }

  conn_.fd = *maybe_fd;

  int rv;
  rv = connect(conn_.fd, addr_->as_sockaddr(), addr_->size());
  if (rv != 0 && errno != EINPROGRESS) {
    auto error = errno;
    Log{WARN, this} << "connect() failed; errno=" << error;

    close(conn_.fd);
    conn_.fd = -1;

    return std::unexpected{Error::SYSCALL};
  }

  if (ssl_ctx_) {
    if (!util::numeric_host(sni_name_.data())) {
      SSL_set_tlsext_host_name(conn_.tls.ssl, sni_name_.data());
    }

    auto maybe_session = tls::reuse_tls_session(tls_session_cache_);
    if (maybe_session) {
      auto session = *maybe_session;
      SSL_set_session(conn_.tls.ssl, session);
      SSL_SESSION_free(session);
    }

    conn_.prepare_client_handshake();
  }

  if (log_enabled(INFO)) {
    Log{INFO, this} << "Connecting to memcached server";
  }

  ev_io_set(&conn_.wev, conn_.fd, EV_WRITE);
  ev_io_set(&conn_.rev, conn_.fd, EV_READ);

  ev_set_cb(&conn_.wev, connectcb);

  conn_.wlimit.startw();
  ev_timer_again(conn_.loop, &conn_.wt);

  return {};
}

std::expected<void, Error> MemcachedConnection::connected() {
  auto sock_error = util::get_socket_error(conn_.fd);
  if (sock_error != 0) {
    Log{WARN, this} << "memcached connect failed; addr="
                    << util::to_numeric_addr(addr_) << ": errno=" << sock_error;

    connect_blocker_.on_failure();

    conn_.wlimit.stopw();

    return std::unexpected{Error::CONNECT_FAIL};
  }

  if (log_enabled(INFO)) {
    Log{INFO, this} << "connected to memcached server";
  }

  conn_.rlimit.startw();

  ev_set_cb(&conn_.wev, writecb);

  if (conn_.tls.ssl) {
    conn_.again_rt();

    do_read_ = &MemcachedConnection::tls_handshake;
    do_write_ = &MemcachedConnection::tls_handshake;

    return {};
  }

  ev_timer_stop(conn_.loop, &conn_.wt);

  connected_ = true;

  connect_blocker_.on_success();

  do_read_ = &MemcachedConnection::read_clear;
  do_write_ = &MemcachedConnection::write_clear;

  return {};
}

std::expected<void, Error> MemcachedConnection::on_write() {
  return do_write_(*this);
}
std::expected<void, Error> MemcachedConnection::on_read() {
  return do_read_(*this);
}

std::expected<void, Error> MemcachedConnection::tls_handshake() {
  ERR_clear_error();

  conn_.last_read = std::chrono::steady_clock::now();

  if (auto rv = conn_.tls_handshake(); !rv) {
    if (rv.error() == Error::TLS_HANDSHAKE_INPROGRESS) {
      return {};
    }

    connect_blocker_.on_failure();
    return rv;
  }

  if (log_enabled(INFO)) {
    Log{INFO} << "SSL/TLS handshake completed";
  }

  auto &tlsconf = get_config()->tls;

  if (!tlsconf.insecure) {
    if (auto rv = tls::check_cert(conn_.tls.ssl, addr_, sni_name_); !rv) {
      connect_blocker_.on_failure();
      return rv;
    }
  }

  ev_timer_stop(conn_.loop, &conn_.rt);
  ev_timer_stop(conn_.loop, &conn_.wt);

  connected_ = true;

  connect_blocker_.on_success();

  do_read_ = &MemcachedConnection::read_tls;
  do_write_ = &MemcachedConnection::write_tls;

  return on_write();
}

std::expected<void, Error> MemcachedConnection::write_tls() {
  if (!connected_) {
    return {};
  }

  conn_.last_read = std::chrono::steady_clock::now();

  std::array<struct iovec, MAX_WR_IOVCNT> iovbuf;
  std::array<uint8_t, 16_k> buf;

  for (; !sendq_.empty();) {
    auto iov = fill_request_buffer(iovbuf);
    auto p = std::ranges::begin(buf);
    for (auto &v : iov) {
      auto n =
        std::min(static_cast<size_t>(std::ranges::end(buf) - p), v.iov_len);
      p =
        std::ranges::copy_n(static_cast<uint8_t *>(v.iov_base), as_signed(n), p)
          .out;
      if (p == std::ranges::end(buf)) {
        break;
      }
    }

    auto maybe_nwrite = conn_.write_tls({std::ranges::begin(buf), p});
    if (!maybe_nwrite) {
      return std::unexpected{maybe_nwrite.error()};
    }

    auto nwrite = *maybe_nwrite;
    if (nwrite == 0) {
      return {};
    }

    drain_send_queue(as_unsigned(nwrite));
  }

  conn_.wlimit.stopw();
  ev_timer_stop(conn_.loop, &conn_.wt);

  return {};
}

std::expected<void, Error> MemcachedConnection::read_tls() {
  if (!connected_) {
    return {};
  }

  conn_.last_read = std::chrono::steady_clock::now();

  for (;;) {
    auto maybe_data = conn_.read_tls(recvbuf_.wbuffer());
    if (!maybe_data) {
      return std::unexpected{maybe_data.error()};
    }

    auto data = *maybe_data;
    if (data.empty()) {
      return {};
    }

    recvbuf_.write(data.size());

    if (auto rv = parse_packet(); !rv) {
      return rv;
    }
  }

  return {};
}

std::expected<void, Error> MemcachedConnection::write_clear() {
  if (!connected_) {
    return {};
  }

  conn_.last_read = std::chrono::steady_clock::now();

  std::array<struct iovec, MAX_WR_IOVCNT> iovbuf;

  for (; !sendq_.empty();) {
    auto iov = fill_request_buffer(iovbuf);
    auto maybe_nwrite = conn_.writev_clear(iov);
    if (!maybe_nwrite) {
      return std::unexpected{maybe_nwrite.error()};
    }

    auto nwrite = *maybe_nwrite;
    if (nwrite == 0) {
      return {};
    }

    drain_send_queue(as_unsigned(nwrite));
  }

  conn_.wlimit.stopw();
  ev_timer_stop(conn_.loop, &conn_.wt);

  return {};
}

std::expected<void, Error> MemcachedConnection::read_clear() {
  if (!connected_) {
    return {};
  }

  conn_.last_read = std::chrono::steady_clock::now();

  for (;;) {
    auto maybe_data = conn_.read_clear(recvbuf_.wbuffer());
    if (!maybe_data) {
      return std::unexpected{maybe_data.error()};
    }

    auto data = *maybe_data;
    if (data.empty()) {
      return {};
    }

    recvbuf_.write(data.size());

    if (auto rv = parse_packet(); !rv) {
      return rv;
    }
  }

  return {};
}

std::expected<void, Error> MemcachedConnection::parse_packet() {
  auto in = recvbuf_.pos;

  for (;;) {
    auto busy = false;

    switch (parse_state_.state) {
    case MemcachedParseState::HEADER24: {
      if (recvbuf_.last - in < 24) {
        recvbuf_.drain_reset(as_unsigned(in - recvbuf_.pos));
        return {};
      }

      if (recvq_.empty()) {
        Log{WARN, this}
          << "Response received, but there is no in-flight request.";
        return std::unexpected{Error::MEMCACHED};
      }

      auto &req = recvq_.front();

      if (*in != MEMCACHED_RES_MAGIC) {
        Log{WARN, this} << "Response has bad magic: "
                        << static_cast<uint32_t>(*in);
        return std::unexpected{Error::MEMCACHED};
      }
      ++in;

      parse_state_.op = static_cast<MemcachedOp>(*in++);
      parse_state_.keylen = util::get_uint16(in);
      in += 2;
      parse_state_.extralen = *in++;
      // skip 1 byte reserved data type
      ++in;
      parse_state_.status_code =
        static_cast<MemcachedStatusCode>(util::get_uint16(in));
      in += 2;
      parse_state_.totalbody = util::get_uint32(in);
      in += 4;
      // skip 4 bytes opaque
      in += 4;
      parse_state_.cas = util::get_uint64(in);
      in += 8;

      if (req->op != parse_state_.op) {
        Log{WARN, this}
          << "opcode in response does not match to the request: want "
          << static_cast<uint32_t>(req->op) << ", got "
          << static_cast<uint32_t>(parse_state_.op);
        return std::unexpected{Error::MEMCACHED};
      }

      if (parse_state_.keylen != 0) {
        Log{WARN, this} << "zero length keylen expected: got "
                        << parse_state_.keylen;
        return std::unexpected{Error::MEMCACHED};
      }

      if (parse_state_.totalbody > 16_k) {
        Log{WARN, this} << "totalbody is too large: got "
                        << parse_state_.totalbody;
        return std::unexpected{Error::MEMCACHED};
      }

      if (parse_state_.op == MemcachedOp::GET &&
          parse_state_.status_code == MemcachedStatusCode::NO_ERROR &&
          parse_state_.extralen == 0) {
        Log{WARN, this} << "response for GET does not have extra";
        return std::unexpected{Error::MEMCACHED};
      }

      if (parse_state_.totalbody <
          parse_state_.keylen + parse_state_.extralen) {
        Log{WARN, this} << "totalbody is too short: totalbody "
                        << parse_state_.totalbody << ", want min "
                        << parse_state_.keylen + parse_state_.extralen;
        return std::unexpected{Error::MEMCACHED};
      }

      if (parse_state_.extralen) {
        parse_state_.state = MemcachedParseState::EXTRA;
        parse_state_.read_left = parse_state_.extralen;
      } else {
        parse_state_.state = MemcachedParseState::VALUE;
        parse_state_.read_left =
          parse_state_.totalbody - parse_state_.keylen - parse_state_.extralen;
      }
      busy = true;
      break;
    }
    case MemcachedParseState::EXTRA: {
      // We don't use extra for now. Just read and forget.
      auto n = std::min(static_cast<size_t>(recvbuf_.last - in),
                        parse_state_.read_left);

      parse_state_.read_left -= n;
      in += n;
      if (parse_state_.read_left) {
        recvbuf_.reset();
        return {};
      }
      parse_state_.state = MemcachedParseState::VALUE;
      // since we require keylen == 0, totalbody - extralen ==
      // valuelen
      parse_state_.read_left =
        parse_state_.totalbody - parse_state_.keylen - parse_state_.extralen;
      busy = true;
      break;
    }
    case MemcachedParseState::VALUE: {
      auto n = std::min(static_cast<size_t>(recvbuf_.last - in),
                        parse_state_.read_left);

      parse_state_.value.insert(std::ranges::end(parse_state_.value), in,
                                in + n);

      parse_state_.read_left -= n;
      in += n;
      if (parse_state_.read_left) {
        recvbuf_.reset();
        return {};
      }

      if (log_enabled(INFO)) {
        if (parse_state_.status_code != MemcachedStatusCode::NO_ERROR) {
          Log{INFO, this} << "response returned error status: "
                          << static_cast<uint16_t>(parse_state_.status_code);
        }
      }

      // We require at least one complete response to clear try count.
      try_count_ = 0;

      auto req = std::move(recvq_.front());
      recvq_.pop_front();

      if (sendq_.empty() && recvq_.empty()) {
        ev_timer_stop(conn_.loop, &conn_.rt);
      }

      if (!req->canceled && req->cb) {
        req->cb(req.get(), MemcachedResult(parse_state_.status_code,
                                           std::move(parse_state_.value)));
      }

      parse_state_ = {};
      break;
    }
    }

    if (!busy && in == recvbuf_.last) {
      break;
    }
  }

  assert(in == recvbuf_.last);
  recvbuf_.reset();

  return {};
}

#undef DEFAULT_WR_IOVCNT
#define DEFAULT_WR_IOVCNT 128

#if defined(IOV_MAX) && IOV_MAX < DEFAULT_WR_IOVCNT
#  define MAX_WR_IOVCNT IOV_MAX
#else // !defined(IOV_MAX) || IOV_MAX >= DEFAULT_WR_IOVCNT
#  define MAX_WR_IOVCNT DEFAULT_WR_IOVCNT
#endif // !defined(IOV_MAX) || IOV_MAX >= DEFAULT_WR_IOVCNT

std::span<struct iovec>
MemcachedConnection::fill_request_buffer(std::span<struct iovec> iov) {
  if (sendsum_ == 0) {
    for (auto &req : sendq_) {
      if (req->canceled) {
        continue;
      }
      if (serialized_size(req.get()) + sendsum_ > 1300) {
        break;
      }
      sendbufv_.emplace_back();
      sendbufv_.back().req = req.get();
      make_request(&sendbufv_.back(), req.get());
      sendsum_ += sendbufv_.back().left();
    }

    if (sendsum_ == 0) {
      sendq_.clear();
      return {};
    }
  }

  size_t iovcnt = 0;
  for (auto &buf : sendbufv_) {
    if (iovcnt + 2 > iov.size()) {
      break;
    }

    auto req = buf.req;
    if (buf.headbuf.rleft()) {
      iov[iovcnt++] = {buf.headbuf.pos, buf.headbuf.rleft()};
    }
    if (buf.send_value_left) {
      iov[iovcnt++] = {req->value.data() + req->value.size() -
                         buf.send_value_left,
                       buf.send_value_left};
    }
  }

  return iov.first(iovcnt);
}

void MemcachedConnection::drain_send_queue(size_t nwrite) {
  sendsum_ -= nwrite;

  while (nwrite > 0) {
    auto &buf = sendbufv_.front();
    auto &req = sendq_.front();
    if (req->canceled) {
      sendq_.pop_front();
      continue;
    }
    assert(buf.req == req.get());
    auto n = std::min(static_cast<size_t>(nwrite), buf.headbuf.rleft());
    buf.headbuf.drain(n);
    nwrite -= n;
    n = std::min(static_cast<size_t>(nwrite), buf.send_value_left);
    buf.send_value_left -= n;
    nwrite -= n;

    if (buf.headbuf.rleft() || buf.send_value_left) {
      break;
    }
    sendbufv_.pop_front();
    recvq_.push_back(std::move(sendq_.front()));
    sendq_.pop_front();
  }

  // start read timer only when we wait for responses.
  if (recvq_.empty()) {
    ev_timer_stop(conn_.loop, &conn_.rt);
  } else if (!ev_is_active(&conn_.rt)) {
    conn_.again_rt();
  }
}

size_t MemcachedConnection::serialized_size(MemcachedRequest *req) {
  switch (req->op) {
  case MemcachedOp::GET:
    return 24 + req->key.size();
  case MemcachedOp::ADD:
  default:
    return 24 + 8 + req->key.size() + req->value.size();
  }
}

void MemcachedConnection::make_request(MemcachedSendbuf *sendbuf,
                                       MemcachedRequest *req) {
  auto &headbuf = sendbuf->headbuf;

  std::ranges::fill(headbuf.buf, 0);

  headbuf[0] = MEMCACHED_REQ_MAGIC;
  headbuf[1] = static_cast<uint8_t>(req->op);
  switch (req->op) {
  case MemcachedOp::GET:
    util::put_uint16be(&headbuf[2], static_cast<uint16_t>(req->key.size()));
    util::put_uint32be(&headbuf[8], static_cast<uint32_t>(req->key.size()));
    headbuf.write(24);
    break;
  case MemcachedOp::ADD:
    util::put_uint16be(&headbuf[2], static_cast<uint16_t>(req->key.size()));
    headbuf[4] = 8;
    util::put_uint32be(&headbuf[8], static_cast<uint32_t>(8 + req->key.size() +
                                                          req->value.size()));
    util::put_uint32be(&headbuf[28], req->expiry);
    headbuf.write(32);
    break;
  }

  headbuf.write(as_uint8_span(std::span{req->key}));

  sendbuf->send_value_left = req->value.size();
}

std::expected<void, Error>
MemcachedConnection::add_request(std::unique_ptr<MemcachedRequest> req) {
  if (connect_blocker_.blocked()) {
    return std::unexpected{Error::INTERNAL};
  }

  sendq_.push_back(std::move(req));

  if (connected_) {
    signal_write();
    return {};
  }

  if (conn_.fd == -1) {
    if (auto rv = initiate_connection(); !rv) {
      connect_blocker_.on_failure();
      disconnect();
      return rv;
    }
  }

  return {};
}

// TODO should we start write timer too?
void MemcachedConnection::signal_write() { conn_.wlimit.startw(); }

void MemcachedConnection::reconnect_or_fail() {
  if (!connected_ || (recvq_.empty() && sendq_.empty())) {
    disconnect();
    return;
  }

  constexpr size_t MAX_TRY_COUNT = 3;

  if (++try_count_ >= MAX_TRY_COUNT) {
    if (log_enabled(INFO)) {
      Log{INFO, this} << "Tried " << MAX_TRY_COUNT
                      << " times, and all failed.  Aborting";
    }
    try_count_ = 0;
    disconnect();
    return;
  }

  std::vector<std::unique_ptr<MemcachedRequest>> q;
  q.reserve(recvq_.size() + sendq_.size());

  if (log_enabled(INFO)) {
    Log{INFO, this} << "Retry connection, enqueue "
                    << recvq_.size() + sendq_.size() << " request(s) again";
  }

  q.insert(std::ranges::end(q),
           std::make_move_iterator(std::ranges::begin(recvq_)),
           std::make_move_iterator(std::ranges::end(recvq_)));
  q.insert(std::ranges::end(q),
           std::make_move_iterator(std::ranges::begin(sendq_)),
           std::make_move_iterator(std::ranges::end(sendq_)));

  recvq_.clear();
  sendq_.clear();

  disconnect();

  sendq_.insert(std::ranges::end(sendq_),
                std::make_move_iterator(std::ranges::begin(q)),
                std::make_move_iterator(std::ranges::end(q)));

  if (!initiate_connection()) {
    connect_blocker_.on_failure();
    disconnect();
    return;
  }
}

} // namespace shrpx