feat(server): Improved cancellation (#599)

src/server/common.cc

@@ -263,6 +263,10 @@ std::string GenericError::Format() const {
     return absl::StrCat(ec_.message(), ":", details_);
 }
 
+Context::~Context() {
+  JoinErrorHandler();
+}
+
 GenericError Context::GetError() {
   std::lock_guard lk(mu_);
   return err_;
@@ -273,20 +277,45 @@ const Cancellation* Context::GetCancellation() const {
 }
 
 void Context::Cancel() {
-  Error(std::make_error_code(errc::operation_canceled), "Context cancelled");
+  ReportError(std::make_error_code(errc::operation_canceled), "Context cancelled");
 }
 
 void Context::Reset(ErrHandler handler) {
   std::lock_guard lk{mu_};
+  JoinErrorHandler();
   err_ = {};
   err_handler_ = std::move(handler);
   Cancellation::flag_.store(false, std::memory_order_relaxed);
 }
 
-GenericError Context::Switch(ErrHandler handler) {
+GenericError Context::SwitchErrorHandler(ErrHandler handler) {
   std::lock_guard lk{mu_};
-  if (!err_)
+  if (!err_) {
+    // No need to check for the error handler - it can't be running
+    // if no error is set.
     err_handler_ = std::move(handler);
+  }
+  return err_;
+}
+
+void Context::JoinErrorHandler() {
+  if (err_handler_fb_.IsJoinable())
+    err_handler_fb_.Join();
+}
+
+GenericError Context::ReportErrorInternal(GenericError&& err) {
+  std::lock_guard lk{mu_};
+  if (err_)
+    return err_;
+  err_ = std::move(err);
+
+  // This context is either new or was Reset, where the handler was joined
+  CHECK(!err_handler_fb_.IsJoinable());
+
+  if (err_handler_)
+    err_handler_fb_ = util::fibers_ext::Fiber{err_handler_, err_};
+
+  Cancellation::Cancel();
   return err_;
 }
 

src/server/common.h

@@ -15,6 +15,7 @@
 
 #include "facade/facade_types.h"
 #include "facade/op_status.h"
+#include "util/fibers/fiber.h"
 
 namespace dfly {
 
@@ -243,7 +244,8 @@ using AggregateGenericError = AggregateValue<GenericError>;
 // Context is a utility for managing error reporting and cancellation for complex tasks.
 //
 // When submitting an error with `Error`, only the first is stored (as in aggregate values).
-// Then a special error handler is run, if present, and the context is cancelled.
+// Then a special error handler is run, if present, and the context is cancelled. The error handler
+// is run in a separate handler to free up the caller.
 //
 // Manual cancellation with `Cancel` is simulated by reporting an `errc::operation_canceled` error.
 // This allows running the error handler and representing this scenario as an error.
@@ -255,38 +257,22 @@ class Context : protected Cancellation {
   Context(ErrHandler err_handler) : Cancellation{}, err_{}, err_handler_{std::move(err_handler)} {
   }
 
-  // Cancels the context by submitting an `errc::operation_canceled` error.
-  void Cancel();
-  using Cancellation::IsCancelled;
+  ~Context();
 
+  void Cancel();  // Cancels the context by submitting an `errc::operation_canceled` error.
+  using Cancellation::IsCancelled;
   const Cancellation* GetCancellation() const;
 
   GenericError GetError();
 
   // Report an error by submitting arguments for GenericError.
   // If this is the first error that occured, then the error handler is run
   // and the context is cancelled.
-  //
-  // Note: this function blocks when called from inside an error handler.
-  template <typename... T> GenericError Error(T... ts) {
-    if (!mu_.try_lock())  // TODO: Maybe use two separate locks.
-      return GenericError{std::forward<T>(ts)...};
-
-    std::lock_guard lk{mu_, std::adopt_lock};
-    if (err_)
-      return err_;
-
-    GenericError new_err{std::forward<T>(ts)...};
-    if (err_handler_)
-      err_handler_(new_err);
-
-    err_ = std::move(new_err);
-    Cancellation::Cancel();
-
-    return err_;
+  template <typename... T> GenericError ReportError(T... ts) {
+    return ReportErrorInternal(GenericError{std::forward<T>(ts)...});
   }
 
-  // Reset error and cancellation flag, assign new error handler.
+  // Wait for error handler to stop, reset error and cancellation flag, assign new error handler.
   void Reset(ErrHandler handler);
 
   // Atomically replace the error handler if no error is present, and return the
@@ -295,12 +281,21 @@ class Context : protected Cancellation {
   // Beware, never do this manually in two steps. If you check for cancellation,
   // set the error handler and initialize resources, then the new error handler
   // will never run if the context was cancelled between the first two steps.
-  GenericError Switch(ErrHandler handler);
+  GenericError SwitchErrorHandler(ErrHandler handler);
+
+  // If any error handler is running, wait for it to stop.
+  void JoinErrorHandler();
+
+ private:
+  // Report error.
+  GenericError ReportErrorInternal(GenericError&& err);
 
  private:
   GenericError err_;
-  ErrHandler err_handler_;
   ::boost::fibers::mutex mu_;
+
+  ErrHandler err_handler_;
+  ::util::fibers_ext::Fiber err_handler_fb_;
 };
 
 struct ScanOpts {

src/server/dflycmd.cc

@@ -314,7 +314,7 @@ void DflyCmd::Expire(CmdArgList args, ConnectionContext* cntx) {
 }
 
 OpStatus DflyCmd::StartFullSyncInThread(FlowInfo* flow, Context* cntx, EngineShard* shard) {
-  DCHECK(!flow->full_sync_fb.joinable());
+  DCHECK(!flow->full_sync_fb.IsJoinable());
 
   SaveMode save_mode = shard == nullptr ? SaveMode::SUMMARY : SaveMode::SINGLE_SHARD;
   flow->saver.reset(new RdbSaver(flow->conn->socket(), save_mode, false));
@@ -341,8 +341,8 @@ void DflyCmd::StopFullSyncInThread(FlowInfo* flow, EngineShard* shard) {
   }
 
   // Wait for full sync to finish.
-  if (flow->full_sync_fb.joinable()) {
-    flow->full_sync_fb.join();
+  if (flow->full_sync_fb.IsJoinable()) {
+    flow->full_sync_fb.Join();
   }
 
   // Reset cleanup and saver
@@ -382,18 +382,18 @@ void DflyCmd::FullSyncFb(FlowInfo* flow, Context* cntx) {
   }
 
   if (ec) {
-    cntx->Error(ec);
+    cntx->ReportError(ec);
     return;
   }
 
   if (ec = saver->SaveBody(cntx->GetCancellation(), nullptr); ec) {
-    cntx->Error(ec);
+    cntx->ReportError(ec);
     return;
   }
 
   ec = flow->conn->socket()->Write(io::Buffer(flow->eof_token));
   if (ec) {
-    cntx->Error(ec);
+    cntx->ReportError(ec);
     return;
   }
 }
@@ -406,9 +406,8 @@ uint32_t DflyCmd::CreateSyncSession() {
   auto err_handler = [this, sync_id](const GenericError& err) {
     LOG(INFO) << "Replication error: " << err.Format();
 
-    // Stop replication in case of error.
-    // StopReplication needs to run async to prevent blocking
-    // the error handler.
+    // Spawn external fiber to allow destructing the context from outside
+    // and return from the handler immediately.
     ::boost::fibers::fiber{&DflyCmd::StopReplication, this, sync_id}.detach();
   };
 
@@ -473,8 +472,8 @@ void DflyCmd::CancelReplication(uint32_t sync_id, shared_ptr<ReplicaInfo> replic
       }
     }
 
-    if (flow->full_sync_fb.joinable()) {
-      flow->full_sync_fb.join();
+    if (flow->full_sync_fb.IsJoinable()) {
+      flow->full_sync_fb.Join();
     }
   });
 
@@ -484,6 +483,9 @@ void DflyCmd::CancelReplication(uint32_t sync_id, shared_ptr<ReplicaInfo> replic
     replica_infos_.erase(sync_id);
   }
 
+  // Wait for error handler to quit.
+  replica_ptr->cntx.JoinErrorHandler();
+
   LOG(INFO) << "Evicted sync session " << sync_id;
 }
 

src/server/dflycmd.h

@@ -12,6 +12,7 @@
 #include <memory>
 
 #include "server/conn_context.h"
+#include "util/fibers/fiber.h"
 
 namespace facade {
 class RedisReplyBuilder;
@@ -91,8 +92,8 @@ class DflyCmd {
 
     facade::Connection* conn;
 
-    ::boost::fibers::fiber full_sync_fb;  // Full sync fiber.
-    std::unique_ptr<RdbSaver> saver;      // Saver used by the full sync phase.
+    util::fibers_ext::Fiber full_sync_fb;  // Full sync fiber.
+    std::unique_ptr<RdbSaver> saver;       // Saver used by the full sync phase.
     std::string eof_token;
 
     std::function<void()> cleanup;  // Optional cleanup for cancellation.