tmp

Signed-off-by: Kostas Kyrimis <[email protected]>

Author: kostasrim

src/facade/dragonfly_connection.cc

@@ -10,11 +10,13 @@
 #include <absl/strings/str_cat.h>
 #include <absl/time/time.h>
 
+#include <condition_variable>
 #include <numeric>
 #include <utility>
 #include <variant>
 
 #include "absl/cleanup/cleanup.h"
+#include "absl/types/span.h"
 #include "base/cycle_clock.h"
 #include "base/flag_utils.h"
 #include "base/flags.h"
@@ -699,6 +701,8 @@ void Connection::OnShutdown() {
   VLOG(1) << "Connection::OnShutdown";
 
   BreakOnce(POLLHUP);
+  io_ec_ = make_error_code(errc::connection_aborted);
+  io_event_.notify();
 }
 
 void Connection::OnPreMigrateThread() {
@@ -1237,10 +1241,9 @@ Connection::ParserStatus Connection::ParseRedis(unsigned max_busy_cycles) {
   read_buffer.slice = io_buf_.InputBuffer();
   read_buffer.available_bytes = io_buf_.InputLen();
 
+  size_t total = 0;
+
   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_);
     request_consumed_bytes_ += consumed;
     if (result == RedisParser::OK && !tmp_parse_args_.empty()) {
@@ -1269,9 +1272,11 @@ 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);
+      LOG(WARNING) << "Redis parser error: " << result
+                   << " during parse: " << ToSV(read_buffer.slice);
     }
+
+    total += consumed;
     read_buffer.Consume(consumed);
 
     // We must yield from time to time to allow other fibers to run.
@@ -1284,7 +1289,7 @@ Connection::ParserStatus Connection::ParseRedis(unsigned max_busy_cycles) {
   } while (RedisParser::OK == result && read_buffer.available_bytes > 0 &&
            !reply_builder_->GetError());
 
-  io_buf_.ConsumeInput(io_buf_.InputLen());
+  io_buf_.ConsumeInput(total);
 
   parser_error_ = result;
   if (result == RedisParser::OK)
@@ -2189,43 +2194,52 @@ bool Connection::WeakRef::operator==(const WeakRef& other) const {
 }
 
 void Connection::DoReadOnRecv(const util::FiberSocketBase::RecvNotification& n) {
-  absl::Cleanup on_exit([this]() { io_event_.notify(); });
-
   if (std::holds_alternative<std::error_code>(n.read_result)) {
     io_ec_ = std::get<std::error_code>(n.read_result);
     return;
   }
 
-  // EnableRecvMultishot path
-  // if (std::holds_alternative<io::MutableBytes>(n.read_result)) {
-  //   auto buf = std::get<io::MutableBytes>(n.read_result);
-  //   // TODO get rid of the copy
-  //   io_buf_.WriteAndCommit(buf.data(), buf.size());
-  //   return;
-  // }
+  // TODO non epoll API via EnableRecvMultishot
+  // if (std::holds_alternative<io::MutableBytes>(n.read_result))
 
   if (std::holds_alternative<monostate>(n.read_result)) {
-    while (true) {
-      io::MutableBytes buf = io_buf_.AppendBuffer();
-      size_t buf_sz = buf.size();
-      int res = recv(socket_->native_handle(), buf.data(), buf.size(), 0);
-      if (res < 0) {
-        if (errno == EAGAIN || errno == EWOULDBLOCK) {
-          // we are done, there is no more socket data available to be read
-          return;
-        }
-        LOG(FATAL) << "Recv error: " << strerror(-res) << " errno " << errno;
+    if (io_buf_.AppendLen() == 0) {
+      // We will regrow in IoLoop
+      return;
+    }
+
+    io::MutableBytes buf = io_buf_.AppendBuffer();
+    int res = recv(socket_->native_handle(), buf.data(), buf.size(), 0);
+
+    // error path
+    if (res < 0) {
+      // LOG(INFO) << "ERROR";
+      if (errno == EAGAIN || errno == EWOULDBLOCK) {
+        // we are done, there is no more socket data available to be read.
+        LOG(INFO) << "EAGAIN on: " << this;
         return;
       }
-      if (res == 0) {
+
+      if (errno == ECONNRESET) {
+        // The peer can shutdown the connection abruptly.
         io_ec_ = make_error_code(errc::connection_aborted);
-        return;
-      }
-      io_buf_.CommitWrite(res);
-      if (buf_sz == size_t(res)) {
-        return;
       }
+
+      LOG_IF(FATAL, !io_ec_) << "Recv error: " << strerror(-res) << " errno " << errno;
+      return;
+    }
+
+    if (res == 0) {
+      LOG(INFO) << "ZERO";
+      io_ec_ = make_error_code(errc::connection_aborted);
+      return;
     }
+    // A recv call can return fewer bytes than requested even if the
+    // socket buffer actually contains enough data to satisfy the full request.
+    // TODO maybe worth looping here and try another recv call until it fails
+    // with EAGAIN or EWOULDBLOCK.
+    io_buf_.CommitWrite(res);
+    return;
   }
 
   DCHECK(false) << "Sould not reach here";
@@ -2240,18 +2254,47 @@ variant<error_code, Connection::ParserStatus> Connection::IoLoopV2() {
   auto* peer = socket_.get();
   recv_buf_.res_len = 0;
 
-  // TODO EnableRecvMultishot on iourin ?
+  // TODO EnableRecvMultishot
+  size_t j = 0;
 
   // Breaks with TLS. RegisterOnRecv is unimplemented.
-  peer->RegisterOnRecv([this](const FiberSocketBase::RecvNotification& n) {
+  peer->RegisterOnRecv([this, &j](const FiberSocketBase::RecvNotification& n) {
+    LOG(INFO) << "Called from hook " << this;
+    ++j;
     DoReadOnRecv(n);
     io_event_.notify();
   });
 
+  size_t i = 0;
+
   do {
+    ++i;
+    if (i == 1) {
+      LOG(INFO) << "Hit " << i << " with j " << j;
+    }
     HandleMigrateRequest();
 
-    io_event_.await([this]() { return io_buf_.InputLen() > 0 || io_ec_; });
+    // We *must* poll again for readiness. The event handler we registered above
+    // with RegisterOnRecv() will get called *once* for each socket readiness event.
+    // So, when we get notified below in io_event_.wait() we might read less data
+    // than it is available because io_buf_ does not have enough capacity. If we loop,
+    // and do not attempt to read from the socket again we can deadlock. To avoid this,
+    // we poll once for readiness before preempting.
+    DoReadOnRecv(FiberSocketBase::RecvNotification{std::monostate()});
+    LOG(INFO) << "Before block " << this;
+    std::cv_status res = io_event_.await_until(
+        [this]() { return io_buf_.InputLen() > 0 || io_ec_ || io_buf_.AppendLen() == 0; },
+        chrono::steady_clock::now() + 10000ms);
+    LOG(INFO) << "After block " << this;
+
+    if (res == std::cv_status::timeout) {
+      LOG(INFO) << "TIMEOUT DETECTED FOR CONNECTION " << this << "\nTotal iterations " << i
+                << "\nTotal registered on recv calls " << j;
+      LOG(INFO) << "^ Before " << io_buf_.InputLen() << " " << this;
+      DoReadOnRecv(FiberSocketBase::RecvNotification{std::monostate()});
+      LOG(INFO) << "^ After " << io_buf_.InputLen() << " " << this;
+      // DCHECK(io_buf_.InputLen() != 0);
+    }
 
     if (io_ec_) {
       LOG_IF(WARNING, cntx()->replica_conn) << "async io error: " << ec;
@@ -2268,6 +2311,8 @@ variant<error_code, Connection::ParserStatus> Connection::IoLoopV2() {
       parse_status = ParseMemcache();
     }
 
+    // LOG(INFO) << "parse status is " << size_t(parse_status);
+
     if (reply_builder_->GetError()) {
       return reply_builder_->GetError();
     }

tests/dragonfly/connection_test.py

@@ -408,23 +408,33 @@ async def pub_task():
 
 
 @pytest.mark.slow
-@dfly_args({"proactor_threads": "4"})
+@dfly_args({"proactor_threads": "4", "migrate_connections": False})
 async def test_pubsub_busy_connections(df_server: DflyInstance):
-    sleep = 60
+    sleep = 10
+
+    idd = 0
 
     async def sub_thread():
         i = 0
 
         async def sub_task():
             nonlocal i
+            nonlocal idd
             sleep_task = asyncio.create_task(asyncio.sleep(sleep))
+            j = idd
+            idd = idd + 1
             while not sleep_task.done():
                 client = df_server.client()
                 pubsub = client.pubsub()
-                await pubsub.subscribe("channel")
+                try:
+                    await pubsub.subscribe("channel")
+                except Exception as e:
+                    logging.info(f"ERRRRRRRROR {j}")
+                    pass
                 # await pubsub.unsubscribe("channel")
                 i = i + 1
                 await client.close()
+            logging.info(f"SUB DONE {j}")
 
         subs = [asyncio.create_task(sub_task()) for _ in range(10)]
         for s in subs:
@@ -436,24 +446,29 @@ async def pub_task():
         i = 0
         sleep_task = asyncio.create_task(asyncio.sleep(sleep))
         while not sleep_task.done():
+            # logging.info("before")
             await pub.publish("channel", f"message-{i}")
             i = i + 1
+            # logging.info("after")
+        logging.info("DONE")
 
     def run_in_thread():
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
         loop.run_until_complete(sub_thread())
 
     threads = []
-    for _ in range(10):
+    for _ in range(1):
         thread = Thread(target=run_in_thread)
         thread.start()
         threads.append(thread)
 
     await pub_task()
 
+    logging.info("==================")
     for thread in threads:
         thread.join()
+    logging.info("==================")
 
 
 async def test_subscribers_with_active_publisher(df_server: DflyInstance, max_connections=100):

tests/dragonfly/instance.py

@@ -436,6 +436,8 @@ def create(self, existing_port=None, path=None, version=100, **kwargs) -> DflyIn
         if version >= 1.26:
             args.setdefault("fiber_safety_margin=4096")
 
+        args.setdefault("expiremental_io_loop_v2=true")
+
         for k, v in args.items():
             args[k] = v.format(**self.params.env) if isinstance(v, str) else v