fix(sock): Use the updated cancellation cb interface.

dragonflydb · Sep 28, 2023 · 04c6c2b · 04c6c2b
1 parent 7630453
commit 04c6c2b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/src/facade/dragonfly_connection.cc b/src/facade/dragonfly_connection.cc
@@ -290,19 +290,19 @@ void Connection::OnShutdown() {
 
 void Connection::OnPreMigrateThread() {
   // If we migrating to another io_uring we should cancel any pending requests we have.
-  if (break_poll_id_ != UINT32_MAX) {
-    socket_->CancelPoll(break_poll_id_);
-    break_poll_id_ = UINT32_MAX;
+  if (break_cb_engaged_) {
+    socket_->CancelOnErrorCb();
+    break_cb_engaged_ = false;
   }
 }
 
 void Connection::OnPostMigrateThread() {
   // Once we migrated, we should rearm OnBreakCb callback.
   if (breaker_cb_) {
-    DCHECK_EQ(UINT32_MAX, break_poll_id_);
+    DCHECK(!break_cb_engaged_);
 
-    break_poll_id_ =
-        socket_->PollEvent(POLLERR | POLLHUP, [this](int32_t mask) { this->OnBreakCb(mask); });
+    socket_->RegisterOnErrorCb([this](int32_t mask) { this->OnBreakCb(mask); });
+    break_cb_engaged_ = true;
   }
 }
 
@@ -369,14 +369,14 @@ void Connection::HandleRequests() {
     } else {
       cc_.reset(service_->CreateContext(peer, this));
       if (breaker_cb_) {
-        break_poll_id_ =
-            socket_->PollEvent(POLLERR | POLLHUP, [this](int32_t mask) { this->OnBreakCb(mask); });
+        socket_->RegisterOnErrorCb([this](int32_t mask) { this->OnBreakCb(mask); });
+        break_cb_engaged_ = true;
       }
 
       ConnectionFlow(peer);
 
-      if (break_poll_id_ != UINT32_MAX) {
-        socket_->CancelPoll(break_poll_id_);
+      if (!break_cb_engaged_) {
+        socket_->CancelOnErrorCb();
       }
 
       cc_.reset();
@@ -729,12 +729,12 @@ void Connection::OnBreakCb(int32_t mask) {
   VLOG(1) << "Got event " << mask;
 
   if (!cc_) {
-    LOG(ERROR) << "Unexpected event " << mask << " " << break_poll_id_;
+    LOG(ERROR) << "Unexpected event " << mask;
     return;
   }
 
   cc_->conn_closing = true;
-  break_poll_id_ = UINT32_MAX;  // do not attempt to cancel it.
+  break_cb_engaged_ = false;  // do not attempt to cancel it.
 
   breaker_cb_(mask);
   evc_.notify();  // Notify dispatch fiber.

diff --git a/src/facade/dragonfly_connection.h b/src/facade/dragonfly_connection.h
@@ -261,7 +261,7 @@ class Connection : public util::Connection {
   std::unique_ptr<ConnectionContext> cc_;
 
   unsigned parser_error_ = 0;
-  uint32_t break_poll_id_ = UINT32_MAX;
+  bool break_cb_engaged_ = false;
 
   BreakerCb breaker_cb_;
   std::unique_ptr<Shutdown> shutdown_cb_;