chore(facade): Handle Socket::Recv() returning 0 explicitly (#6066)

Currently helio translates 0 in Socket:Recv to ECONNABORTED.
This works logically well unless we want to differrentiate between
properly closed sockets (result=0) and those that were closed due to an error.
This PR starts handling result=0 states for sockets even though helio does not yet return them.

Also clean up leftovers from the provided buffers.

Author: romange

src/facade/dragonfly_connection.cc

@@ -1037,6 +1037,11 @@ io::Result<bool> Connection::CheckForHttpProto() {
     if (!recv_sz) {
       return make_unexpected(recv_sz.error());
     }
+    if (recv_sz == 0) {
+      // Peer closed connection.
+      return false;
+    }
+
     io_buf_.CommitWrite(*recv_sz);
     string_view ib = ToSV(io_buf_.InputBuffer());
     if (ib.size() >= 2 && ib[0] == 22 && ib[1] == 3) {
@@ -1139,13 +1144,12 @@ void Connection::ConnectionFlow() {
     // to reproduce: nc localhost 6379  and then run invalid sequence: *1 <enter> *1 <enter>
     error_code ec2 = socket_->Shutdown(SHUT_WR);
     LOG_IF(WARNING, ec2) << "Could not shutdown socket " << ec2;
-    if (!ec2) {
-      while (true) {
-        // Discard any received data.
-        io_buf_.Clear();
-        if (!socket_->Recv(io_buf_.AppendBuffer())) {
-          break;
-        }
+    while (!ec2) {
+      // Discard any received data.
+      io_buf_.Clear();
+      auto recv_sz = socket_->Recv(io_buf_.AppendBuffer());
+      if (!recv_sz || *recv_sz == 0) {
+        break;  // Peer closed connection.
       }
     }
   }
@@ -1220,15 +1224,9 @@ Connection::ParserStatus Connection::ParseRedis(unsigned max_busy_cycles) {
 
   auto dispatch_async = [this]() -> MessageHandle { return {FromArgs(tmp_parse_args_)}; };
 
-  ReadBuffer read_buffer;
-  read_buffer.slice = io_buf_.InputBuffer();
-  read_buffer.available_bytes = io_buf_.InputLen();
-
+  io::Bytes read_buffer = io_buf_.InputBuffer();
   do {
-    if (read_buffer.ShouldAdvance()) {  // can happen only with io_uring/bundles
-      read_buffer.slice = NextBundleBuffer(read_buffer.available_bytes);
-    }
-    result = redis_parser_->Parse(read_buffer.slice, &consumed, &tmp_parse_args_);
+    result = redis_parser_->Parse(read_buffer, &consumed, &tmp_parse_args_);
     request_consumed_bytes_ += consumed;
     if (result == RedisParser::OK && !tmp_parse_args_.empty()) {
       // If we get a non-STRING type (e.g., NIL, ARRAY), it's a protocol error.
@@ -1245,7 +1243,7 @@ Connection::ParserStatus Connection::ParseRedis(unsigned max_busy_cycles) {
       if (io_req_size_hist)
         io_req_size_hist->Add(request_consumed_bytes_);
       request_consumed_bytes_ = 0;
-      bool has_more = consumed < read_buffer.available_bytes;
+      bool has_more = consumed < read_buffer.size();
 
       if (tl_traffic_logger.log_file && IsMain() /* log only on the main interface */) {
         LogTraffic(id_, has_more, absl::MakeSpan(tmp_parse_args_),
@@ -1257,9 +1255,9 @@ Connection::ParserStatus Connection::ParseRedis(unsigned max_busy_cycles) {
     if (result != RedisParser::OK && result != RedisParser::INPUT_PENDING) {
       // We do not expect that a replica sends an invalid command so we log if it happens.
       LOG_IF(WARNING, cntx()->replica_conn)
-          << "Redis parser error: " << result << " during parse: " << ToSV(read_buffer.slice);
+          << "Redis parser error: " << result << " during parse: " << ToSV(read_buffer);
     }
-    read_buffer.Consume(consumed);
+    read_buffer.remove_prefix(consumed);
 
     // We must yield from time to time to allow other fibers to run.
     // Specifically, if a client sends a huge chunk of data resulting in a very long pipeline,
@@ -1268,8 +1266,7 @@ Connection::ParserStatus Connection::ParseRedis(unsigned max_busy_cycles) {
       stats_->num_read_yields++;
       ThisFiber::Yield();
     }
-  } while (RedisParser::OK == result && read_buffer.available_bytes > 0 &&
-           !reply_builder_->GetError());
+  } while (RedisParser::OK == result && read_buffer.size() > 0 && !reply_builder_->GetError());
 
   io_buf_.ConsumeInput(io_buf_.InputLen());
 
@@ -1278,7 +1275,7 @@ Connection::ParserStatus Connection::ParseRedis(unsigned max_busy_cycles) {
     return OK;
 
   if (result == RedisParser::INPUT_PENDING) {
-    DCHECK_EQ(read_buffer.available_bytes, 0u);
+    DCHECK_EQ(read_buffer.size(), 0u);
 
     return NEED_MORE;
   }
@@ -1411,28 +1408,26 @@ void Connection::HandleMigrateRequest() {
   }
 }
 
-error_code Connection::HandleRecvSocket() {
+io::Result<size_t> Connection::HandleRecvSocket() {
   phase_ = READ_SOCKET;
 
   io::MutableBytes append_buf = io_buf_.AppendBuffer();
   DCHECK(!append_buf.empty());
   ::io::Result<size_t> recv_sz = socket_->Recv(append_buf);
   last_interaction_ = time(nullptr);
 
-  if (!recv_sz) {
-    return recv_sz.error();
-  }
-
-  size_t commit_sz = *recv_sz;
-  io_buf_.CommitWrite(commit_sz);
-
-  stats_->io_read_bytes += commit_sz;
-  local_stats_.net_bytes_in += commit_sz;
+  // In case the socket was closed orderly, we get 0 bytes read.
+  if (recv_sz && *recv_sz) {
+    size_t commit_sz = *recv_sz;
+    io_buf_.CommitWrite(commit_sz);
 
-  ++stats_->io_read_cnt;
-  ++local_stats_.read_cnt;
+    stats_->io_read_bytes += commit_sz;
+    local_stats_.net_bytes_in += commit_sz;
 
-  return {};
+    ++stats_->io_read_cnt;
+    ++local_stats_.read_cnt;
+  }
+  return recv_sz;
 }
 
 auto Connection::IoLoop() -> variant<error_code, ParserStatus> {
@@ -1446,10 +1441,13 @@ auto Connection::IoLoop() -> variant<error_code, ParserStatus> {
 
   do {
     HandleMigrateRequest();
-    ec = HandleRecvSocket();
-    if (ec) {
-      LOG_IF(WARNING, cntx()->replica_conn) << "HandleRecvSocket() error: " << ec;
-      return ec;
+    auto recv_sz = HandleRecvSocket();
+    if (!recv_sz) {
+      LOG_IF(WARNING, cntx()->replica_conn) << "HandleRecvSocket() error: " << recv_sz.error();
+      return recv_sz.error();
+    }
+    if (*recv_sz == 0) {
+      break;
     }
 
     phase_ = PROCESS;
@@ -2065,19 +2063,6 @@ void Connection::BreakOnce(uint32_t ev_mask) {
   }
 }
 
-io::Bytes Connection::NextBundleBuffer(size_t total_len) {
-  ++recv_buf_.buf_pos;
-  io::Bytes res;
-#ifdef __linux__
-  fb2::UringProactor* up = static_cast<fb2::UringProactor*>(socket_->proactor());
-  recv_buf_.buf_id = up->GetBufIdByPos(kRecvSockGid, recv_buf_.buf_pos);
-  uint8_t* src = up->GetBufRingPtr(kRecvSockGid, recv_buf_.buf_id);
-  res = {src, std::min<size_t>(kRecvBufSize, total_len)};
-#endif
-
-  return res;
-}
-
 void Connection::IncrNumConns() {
   if (IsMainOrMemcache())
     ++stats_->num_conns_main;

src/facade/dragonfly_connection.h

@@ -382,7 +382,7 @@ class Connection : public util::Connection {
   PipelineMessagePtr GetFromPipelinePool();
 
   void HandleMigrateRequest();
-  std::error_code HandleRecvSocket();
+  io::Result<size_t> HandleRecvSocket();
 
   bool ShouldEndAsyncFiber(const MessageHandle& msg);
 
@@ -406,18 +406,12 @@ class Connection : public util::Connection {
     size_t available_bytes;
     io::Bytes slice;
 
-    bool ShouldAdvance() const {
-      return slice.empty();
-    }
-
     void Consume(size_t len) {
       available_bytes -= len;
       slice.remove_prefix(len);
     }
   };
 
-  io::Bytes NextBundleBuffer(size_t total_len);
-
   void IncrNumConns();
   void DecrNumConns();
 

src/server/protocol_client.cc

@@ -61,6 +61,11 @@ error_code ProtocolClient::Recv(FiberSocketBase* input, base::IoBuf* dest) {
     return exp_size.error();
   }
 
+  if (*exp_size == 0) {
+    VLOG(1) << "Connection closed by peer";
+    return make_error_code(errc::connection_aborted);
+  }
+
   TouchIoTime();
 
   dest->CommitWrite(*exp_size);
@@ -72,21 +77,21 @@ std::string ProtocolClient::ServerContext::Description() const {
 }
 
 void ValidateClientTlsFlags() {
-  if (!absl::GetFlag(FLAGS_tls_replication)) {
+  if (!GetFlag(FLAGS_tls_replication)) {
     return;
   }
 
   bool has_auth = false;
 
-  if (!absl::GetFlag(FLAGS_tls_key_file).empty()) {
-    if (absl::GetFlag(FLAGS_tls_cert_file).empty()) {
+  if (!GetFlag(FLAGS_tls_key_file).empty()) {
+    if (GetFlag(FLAGS_tls_cert_file).empty()) {
       LOG(ERROR) << "tls_cert_file flag should be set";
       exit(1);
     }
     has_auth = true;
   }
 
-  if (!absl::GetFlag(FLAGS_masterauth).empty())
+  if (!GetFlag(FLAGS_masterauth).empty())
     has_auth = true;
 
   if (!has_auth) {
@@ -97,7 +102,7 @@ void ValidateClientTlsFlags() {
 
 void ProtocolClient::MaybeInitSslCtx() {
 #ifdef DFLY_USE_SSL
-  if (absl::GetFlag(FLAGS_tls_replication)) {
+  if (GetFlag(FLAGS_tls_replication)) {
     ssl_ctx_ = CreateSslCntx(facade::TlsContextRole::CLIENT);
   }
 #endif
@@ -207,8 +212,8 @@ error_code ProtocolClient::ConnectAndAuth(std::chrono::milliseconds connect_time
 
   // CHECK_EQ(0, setsockopt(sock_->native_handle(), IPPROTO_TCP, TCP_NODELAY, &yes, sizeof(yes)));
 
-  auto masterauth = absl::GetFlag(FLAGS_masterauth);
-  auto masteruser = absl::GetFlag(FLAGS_masteruser);
+  auto masterauth = GetFlag(FLAGS_masterauth);
+  auto masteruser = GetFlag(FLAGS_masteruser);
   ResetParser(RedisParser::Mode::CLIENT);
   if (!masterauth.empty()) {
     auto cmd = masteruser.empty() ? StrCat("AUTH ", masterauth)