CPP-914 Fix: WaitForHandler::on_set() called after error/timeout

mpenick · web-flow · commit bd325094efa1 · 2020-03-16T13:44:14.000-04:00
diff --git a/src/wait_for_handler.cpp b/src/wait_for_handler.cpp
@@ -51,6 +51,8 @@ void WaitForCallback::on_chain_write(Connection* connection) { handler_->start(c
 }}} // namespace datastax::internal::core
 
 void WaitForCallback::on_chain_set() {
+  if (handler_->is_finished_) return;
+
   if (handler_->on_set(Ptr(this))) {
     handler_->finish();
   } else {
@@ -59,6 +61,8 @@ void WaitForCallback::on_chain_set() {
 }
 
 void WaitForCallback::on_chain_error(CassError code, const String& message) {
+  if (handler_->is_finished_) return;
+
   OStringStream ss;
   ss << message << " (0x" << std::hex << std::uppercase << std::setw(8) << std::setfill('0') << code
      << ")";
@@ -67,6 +71,8 @@ void WaitForCallback::on_chain_error(CassError code, const String& message) {
 }
 
 void WaitForCallback::on_chain_timeout() {
+  if (handler_->is_finished_) return;
+
   handler_->on_error(WaitForHandler::WAIT_FOR_ERROR_REQUEST_TIMEOUT, "Request timed out");
   handler_->schedule();
 }
@@ -102,26 +108,26 @@ void WaitForHandler::start(Connection* connection) {
 }
 
 void WaitForHandler::schedule() {
-  if (connection_ && !is_finished_) { // Don't schedule a retry if the handler is finished.
-    retry_timer_.start(connection_->loop(), retry_wait_time_ms_,
-                       bind_callback(&WaitForHandler::on_retry_timeout, this));
-  }
+  assert(!is_finished_ && "This shouldn't be called if the handler is finished");
+  retry_timer_.start(connection_->loop(), retry_wait_time_ms_,
+                     bind_callback(&WaitForHandler::on_retry_timeout, this));
 }
 
 void WaitForHandler::finish() {
-  if (!is_finished_) {
-    is_finished_ = true;
-    request_handler_->set_response(current_host_, response_);
-    if (connection_) {
-      connection_.reset();
-      retry_timer_.stop();
-      timer_.stop();
-      dec_ref();
-    }
+  assert(!is_finished_ && "This shouldn't be called more than once");
+  is_finished_ = true;
+  request_handler_->set_response(current_host_, response_);
+  if (connection_) {
+    connection_.reset();
+    retry_timer_.stop();
+    timer_.stop();
+    dec_ref();
   }
 }
 
 void WaitForHandler::on_retry_timeout(Timer* timer) {
+  if (is_finished_) return;
+
   if (connection_->is_closing()) {
     on_error(WaitForHandler::WAIT_FOR_ERROR_CONNECTION_CLOSED, "Connection closed");
     finish();
@@ -133,6 +139,8 @@ void WaitForHandler::on_retry_timeout(Timer* timer) {
 }
 
 void WaitForHandler::on_timeout(Timer* timer) {
+  if (is_finished_) return;
+
   on_error(WaitForHandler::WAIT_FOR_ERROR_TIMEOUT, "Timed out");
   finish();
 }
diff --git a/tests/src/unit/tests/test_wait_for_handler.cpp b/tests/src/unit/tests/test_wait_for_handler.cpp
@@ -22,9 +22,25 @@
 #include "timer.hpp"
 #include "wait_for_handler.hpp"
 
+#include <ostream>
+
 using namespace datastax::internal;
 using namespace datastax::internal::core;
 
+typedef Vector<WaitForHandler::WaitForError> Errors;
+
+namespace std {
+
+std::ostream& operator<<(std::ostream& os, const Errors& errors) {
+  for (Errors::const_iterator it = errors.begin(); it != errors.end(); ++it) {
+    if (it != errors.begin()) os << ", ";
+    os << *it;
+  }
+  return os;
+}
+
+} // namespace std
+
 class WaitForHandlerUnitTest : public LoopTest {
 public:
   class TestWaitForHandler : public WaitForHandler {
@@ -35,54 +51,70 @@ class WaitForHandlerUnitTest : public LoopTest {
         : WaitForHandler(
               RequestHandler::Ptr(new RequestHandler(QueryRequest::Ptr(new QueryRequest("")),
                                                      ResponseFuture::Ptr(new ResponseFuture()))),
-              Host::Ptr(new Host(Address())), Response::Ptr(), max_wait_time, retry_wait_time) {}
-
-    virtual RequestCallback::Ptr callback() = 0;
+              Host::Ptr(new Host(Address())), Response::Ptr(), max_wait_time, retry_wait_time)
+        , count_on_set_(0)
+        , count_on_error_(0)
+        , is_idempotent_(false) {}
+
+    Ptr with_is_idempotent(bool is_idempotent) {
+      is_idempotent_ = is_idempotent;
+      return Ptr(this);
+    }
 
-  private:
-    virtual bool on_set(const ChainedRequestCallback::Ptr& callback) {
-      return false; // Never complete
+    Ptr with_expected_error(WaitForHandler::WaitForError error) {
+      expected_.push_back(error);
+      return Ptr(this);
     }
-  };
 
-  class RegularQueryHandler : public TestWaitForHandler {
-  public:
     virtual RequestCallback::Ptr callback() {
       WaitforRequestVec requests;
-      requests.push_back(make_request("local", "SELECT * FROM system.local WHERE key='local'"));
-      requests.push_back(make_request("peers", "SELECT * FROM system.peers"));
+      QueryRequest::Ptr table1_request(new QueryRequest("SELECT * FROM test.table1"));
+      QueryRequest::Ptr table2_request(new QueryRequest("SELECT * FROM test.table2"));
+      table1_request->set_is_idempotent(is_idempotent_);
+      table2_request->set_is_idempotent(is_idempotent_);
+      requests.push_back(WaitForRequest("table1", table1_request));
+      requests.push_back(WaitForRequest("table2", table2_request));
       return WaitForHandler::callback(requests);
     }
 
-  private:
-    virtual void on_error(WaitForError code, const String& message) {
-      EXPECT_TRUE(WAIT_FOR_ERROR_CONNECTION_CLOSED == code || WAIT_FOR_ERROR_REQUEST_ERROR == code);
-    }
-  };
+    int count_on_set() const { return count_on_set_; }
 
-  class IdempotentQueryHandler : public TestWaitForHandler {
-  public:
-    virtual RequestCallback::Ptr callback() {
-      WaitforRequestVec requests;
-      QueryRequest::Ptr local_request(
-          new QueryRequest("SELECT * FROM system.local WHERE key='local'"));
-      QueryRequest::Ptr peers_request(new QueryRequest("SELECT * FROM system.peers"));
-      local_request->set_is_idempotent(true);
-      peers_request->set_is_idempotent(true);
-      requests.push_back(WaitForRequest("local", local_request));
-      requests.push_back(WaitForRequest("peers", peers_request));
-      return WaitForHandler::callback(requests);
+  protected:
+    virtual bool on_set(const ChainedRequestCallback::Ptr& callback) {
+      EXPECT_EQ(0, count_on_error_); // Set shouldn't be called after an error
+      count_on_set_++;
+      return false; // Never complete
     }
 
-  private:
     virtual void on_error(WaitForError code, const String& message) {
-      EXPECT_TRUE(WAIT_FOR_ERROR_CONNECTION_CLOSED == code ||
-                  WAIT_FOR_ERROR_REQUEST_TIMEOUT == code);
+      ASSERT_NE(0, expected_.size());
+      bool found_expected = false;
+      for (Errors::const_iterator it = expected_.begin(), end = expected_.end();
+           !found_expected && it != end; ++it) {
+        if (*it == code) {
+          found_expected = true;
+        }
+      }
+      EXPECT_TRUE(found_expected) << "Expected error codes [ " << expected_
+                                  << " ], but received error " << code;
+      count_on_error_++;
     }
+
+  private:
+    Errors expected_;
+    int count_on_set_;
+    int count_on_error_;
+    bool is_idempotent_;
   };
 
   void run(const TestWaitForHandler::Ptr& handler, uint64_t timeout = 0) {
-    mockssandra::SimpleCluster cluster(simple());
+    mockssandra::SimpleRequestHandlerBuilder builder;
+    run(handler, builder, timeout);
+  }
+
+  void run(const TestWaitForHandler::Ptr& handler,
+           mockssandra::SimpleRequestHandlerBuilder& builder, uint64_t timeout = 0) {
+    mockssandra::SimpleCluster cluster(builder.build());
     ASSERT_EQ(cluster.start_all(), 0);
 
     handler_ = handler;
@@ -140,17 +172,46 @@ class WaitForHandlerUnitTest : public LoopTest {
 };
 
 TEST_F(WaitForHandlerUnitTest, CloseImmediatelyWhileWaiting) {
-  run(TestWaitForHandler::Ptr(new RegularQueryHandler()));
+  run(TestWaitForHandler::Ptr(new TestWaitForHandler())
+          ->with_expected_error(WaitForHandler::WAIT_FOR_ERROR_REQUEST_ERROR)
+          ->with_expected_error(WaitForHandler::WAIT_FOR_ERROR_CONNECTION_CLOSED));
 }
 
 TEST_F(WaitForHandlerUnitTest, CloseAfterTimeoutWhileWaiting) {
-  run(TestWaitForHandler::Ptr(new RegularQueryHandler()), 500);
+  run(TestWaitForHandler::Ptr(new TestWaitForHandler())
+          ->with_expected_error(WaitForHandler::WAIT_FOR_ERROR_REQUEST_ERROR)
+          ->with_expected_error(WaitForHandler::WAIT_FOR_ERROR_CONNECTION_CLOSED),
+      500);
 }
 
 TEST_F(WaitForHandlerUnitTest, CloseIdempotentImmediatelyWhileWaiting) {
-  run(TestWaitForHandler::Ptr(new IdempotentQueryHandler()));
+  run(TestWaitForHandler::Ptr(new TestWaitForHandler())
+          ->with_is_idempotent(true)
+          ->with_expected_error(WaitForHandler::WAIT_FOR_ERROR_REQUEST_TIMEOUT)
+          ->with_expected_error(WaitForHandler::WAIT_FOR_ERROR_CONNECTION_CLOSED));
 }
 
 TEST_F(WaitForHandlerUnitTest, CloseIdempotentAfterTimeoutWhileWaiting) {
-  run(TestWaitForHandler::Ptr(new IdempotentQueryHandler()), 500);
+  run(TestWaitForHandler::Ptr(new TestWaitForHandler())
+          ->with_is_idempotent(true)
+          ->with_expected_error(WaitForHandler::WAIT_FOR_ERROR_REQUEST_TIMEOUT)
+          ->with_expected_error(WaitForHandler::WAIT_FOR_ERROR_CONNECTION_CLOSED),
+      500);
+}
+
+TEST_F(WaitForHandlerUnitTest, EnsureOnSetNotCalledAfterTimeout) {
+  TestWaitForHandler::Ptr handler(
+      new TestWaitForHandler(1)); // Timeout handler before query returns
+
+  // Make sure the query doesn't complete before the handler times out
+  mockssandra::SimpleRequestHandlerBuilder builder;
+  builder.on(mockssandra::OPCODE_QUERY)
+      .system_local()
+      .system_peers()
+      .wait(200)
+      .empty_rows_result(1);
+
+  run(handler->with_expected_error(WaitForHandler::WAIT_FOR_ERROR_TIMEOUT), builder, 500);
+
+  EXPECT_EQ(0, handler->count_on_set()); // Ensure on_set() never called
 }
