chore(reply_builder): add dcheck that each command invocation has rep… (#2067)

* chore(reply_builder): add dcheck that each command invocation has replied

src/facade/reply_builder.cc

@@ -37,7 +37,7 @@ DoubleToStringConverter dfly_conv(kConvFlags, "inf", "nan", 'e', -6, 21, 6, 0);
 }  // namespace
 
 SinkReplyBuilder::SinkReplyBuilder(::io::Sink* sink)
-    : sink_(sink), should_batch_(false), should_aggregate_(false) {
+    : sink_(sink), should_batch_(false), should_aggregate_(false), has_replied_(true) {
 }
 
 void SinkReplyBuilder::CloseConnection() {
@@ -46,6 +46,7 @@ void SinkReplyBuilder::CloseConnection() {
 }
 
 void SinkReplyBuilder::Send(const iovec* v, uint32_t len) {
+  has_replied_ = true;
   DCHECK(sink_);
   constexpr size_t kMaxBatchSize = 1024;
 
@@ -98,6 +99,14 @@ void SinkReplyBuilder::SendRaw(std::string_view raw) {
   Send(&v, 1);
 }
 
+void SinkReplyBuilder::ExpectReply() {
+  has_replied_ = false;
+}
+
+bool SinkReplyBuilder::HasReplied() const {
+  return has_replied_;
+}
+
 void SinkReplyBuilder::SendError(ErrorReply error) {
   if (error.status)
     return SendError(*error.status);

src/facade/reply_builder.h

@@ -113,6 +113,9 @@ class SinkReplyBuilder {
     bool is_nested_ = true;
   };
 
+  void ExpectReply();
+  bool HasReplied() const;
+
  protected:
   void SendRaw(std::string_view str);  // Sends raw without any formatting.
   void SendRawVec(absl::Span<const std::string_view> msg_vec);
@@ -134,6 +137,7 @@ class SinkReplyBuilder {
 
   // Similarly to batch mode but is controlled by at operation level.
   bool should_aggregate_ : 1;
+  bool has_replied_ : 1;
 };
 
 class MCReplyBuilder : public SinkReplyBuilder {

src/facade/reply_capture.cc

@@ -7,6 +7,7 @@
 #include "reply_capture.h"
 
 #define SKIP_LESS(needed)     \
+  has_replied_ = true;        \
   if (reply_mode_ < needed) { \
     current_ = monostate{};   \
     return;                   \

src/server/main_service.cc

@@ -1047,6 +1047,30 @@ void Service::DispatchCommand(CmdArgList args, facade::ConnectionContext* cntx)
   }
 }
 
+class ReplyGuard {
+ public:
+  ReplyGuard(ConnectionContext* cntx, std::string_view cid_name) {
+    auto* builder = cntx->reply_builder();
+    const bool is_script = bool(cntx->conn_state.script_info);
+    const bool is_one_of =
+        absl::flat_hash_set<std::string_view>({"REPLCONF", "DFLY", "EXEC", "EVALSHA"})
+            .contains(cid_name);
+    const bool should_dcheck = !is_one_of && !is_script;
+    if (should_dcheck) {
+      builder->ExpectReply();
+    }
+  }
+
+  ~ReplyGuard() {
+    if (builder) {
+      DCHECK(builder->HasReplied());
+    }
+  }
+
+ private:
+  RedisReplyBuilder* builder = nullptr;
+};
+
 bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args, ConnectionContext* cntx,
                         bool record_stats) {
   DCHECK(cid);
@@ -1066,6 +1090,9 @@ bool Service::InvokeCmd(const CommandId* cid, CmdArgList tail_args, ConnectionCo
     DispatchMonitor(cntx, cid, tail_args);
   }
 
+#ifndef NDEBUG
+  ReplyGuard reply_guard(cntx, cid->name());
+#endif
   try {
     cid->Invoke(tail_args, cntx);
   } catch (std::exception& e) {