implementing protocol handlers' logic

Author: Tianyi Chen

src/include/network/connection_handle.h

@@ -87,7 +87,10 @@ class ConnectionHandle {
   Transition ProcessWrite();
   Transition GetResult();
   Transition CloseSocket();
-
+  /**
+   * Flush out all the responses and do real SSL handshake
+   */
+  Transition ProcessWrite_SSLHandshake();
  private:
   /**
    * A state machine is defined to be a set of states, a set of symbols it
@@ -171,6 +174,12 @@ class ConnectionHandle {
    */
   WriteState FlushWriteBuffer();
 
+  /**
+   * Process SSL handshake to generate valid SSL connection context
+   * for further communications
+   */
+  Transition ConnectionHandle::SSLHandshake();
+
   /**
    * Set the socket to non-blocking mode
    */
@@ -190,6 +199,13 @@ class ConnectionHandle {
     setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof one);
   }
 
+  /**
+   * Determine if there is still responses in the buffer
+   */
+  inline bool HasResponse() {
+    return (protocol_handler_->responses.size() != 0) || (wbuf_->buf_size != 0);
+  }
+
   int sock_fd_;                            // socket file descriptor
   struct event *network_event = nullptr;   // something to read from network
   struct event *workpool_event = nullptr;  // worker thread done the job

src/include/network/network_state.h

@@ -24,7 +24,8 @@ enum class ConnState {
   PROCESS,     // State that runs the network protocol on received data
   CLOSING,     // State for closing the client connection
   GET_RESULT,  // State when triggered by worker thread that completes the task.
-  SSL_HANDSHAKE// State when doing (Real) SSL handshake
+  PROCESS_WRITE_SSL_HANDSHAKE  // State to flush out responses and  doing (Real)
+                               // SSL handshake
 };
 
 // TODO(tianyu): Convert use cases of this to just return Transition
@@ -49,5 +50,5 @@ enum class Transition {
   RETRY,
   NEED_SSL_HANDSHAKE
 };
-}
-}
+}  // namespace network
+}  // namespace peloton

src/include/network/postgres_protocol_handler.h

@@ -37,17 +37,26 @@ typedef std::vector<std::unique_ptr<OutputPacket>> ResponseBuffer;
 
 class PostgresProtocolHandler : public ProtocolHandler {
  public:
-  // TODO we need to somehow make this virtual?
   PostgresProtocolHandler(tcop::TrafficCop *traffic_cop);
 
   ~PostgresProtocolHandler();
+  /**
+   * Parse the content in the buffer and process to generate results.
+   * thread_id is the thread of current running thread. This is used
+   * to generate txn
+   */
+  ProcessResult Process(Buffer &rbuf, const bool ssl_able, const size_t thread_id);
 
-  ProcessResult Process(Buffer &rbuf, const size_t thread_id);
-
-  /* Main switch case wrapper to process every packet apart from the startup
-   * packet. Avoid flushing the response for extended protocols. */
-  ProcessResult ProcessNormalPacket(InputPacket *pkt, const size_t thread_id);
+  /**
+   * 
+   */
+  ProcessResult ProcessInitialPackets(
 
+  /** 
+   * Main switch case wrapper to process every general packet apart from the
+   * initial packets. Avoid flushing the response for extended protocols.
+   */
+  ProcessResult ProcessNormalPackets(InputPacket *pkt, const size_t thread_id);
   /* Manage the startup packet */
   //  bool ManageStartupPacket();
   void SendInitialResponse();
@@ -96,6 +105,9 @@ class PostgresProtocolHandler : public ProtocolHandler {
   //===--------------------------------------------------------------------===//
   // PROTOCOL HANDLING FUNCTIONS
   //===--------------------------------------------------------------------===//
+  /* Manage the startup packet */
+  virtual void SendInitialResponse();
+
   // Generic error protocol packet
   void SendErrorResponse(
       std::vector<std::pair<NetworkMessageType, std::string>> error_status);

src/include/network/protocol_handler.h

@@ -36,11 +36,7 @@ class ProtocolHandler {
   //  bool ManageStartupPacket();
   virtual void SendInitialResponse();
 
-  virtual bool ProcessInitialPackets(InputPacket *pkt, Client client,
-                                     bool ssl_able, bool &ssl_sent,
-                                     bool &finish_startup_packet) = 0;
-
-  virtual ProcessResult Process(Buffer &rbuf, const size_t thread_id);
+  virtual ProcessResult Process(Buffer &rbuf, const bool ssl_able, const size_t thread_id);
 
   virtual void Reset();
 

src/network/connection_handle.cpp

@@ -92,9 +92,10 @@ DEF_TRANSITION_GRAPH
     ON(FINISH) SET_STATE_TO(CLOSING) AND_INVOKE(CloseSocket) 
   END_DEF
 
-  DEFINE_STATE(SSL_HANDSHAKE)
-    ON(WAKEUP) SET_STATE_TO(SSL_HANDSHAKE) AND_INVOKE(SSL_handshake)
-    ON(NEED_DATA) SET_STATE_TO(SSL_HANDSHAKE) AND_WAIT
+  DEFINE_STATE(PROCESS_WRITE_SSL_HANDSHAKE)
+    ON(WAKEUP) SET_STATE_TO(PROCESS_WRITE_SSL_HANDSHAKE)
+      AND_INVOKE(ProcessWrite_SSLHandshake)
+    ON(NEED_DATA) SET_STATE_TO(PROCESS_WRITE_SSL_HANDSHAKE) AND_WAIT
     ON(FINISH) SET_STATE_TO(CLOSING) AND_INVOKE(CloseSocket)
     ON(PROCEED) SET_STATE_TO(PROCESS) AND_INVOKE(Process)
   END_DEF
@@ -104,7 +105,8 @@ DEF_TRANSITION_GRAPH
     ON(NEED_DATA) SET_STATE_TO(READ) AND_INVOKE(FillReadBuffer)
     ON(GET_RESULT) SET_STATE_TO(GET_RESULT) AND_WAIT
     ON(FINISH) SET_STATE_TO(CLOSING) AND_INVOKE(CloseSocket)
-    ON(SSL_HANDSHAKE) SET_STATE_TO(SSL_HANDSHKE) AND_INVOKE(SSL_handshake)
+    ON(NEED_SSL_HANDSHAKE) SET_STATE_TO(PROCESS_WRITE_SSL_HANDSHKE) 
+      AND_INVOKE(ProcessWrite_SSLHandshake)
   END_DEF
 
   DEFINE_STATE(WRITE) 
@@ -576,13 +578,19 @@ Transition ConnectionHandle::CloseSocket() {
   }
 }
 
-Transition ConnectionHandle::SSL_handshake() {
-  if (HaveResponse()) {
-    switch(ProcessWrite()) {
-      
-    } 
+Transition ConnectionHandle::ProcessWrite_SSLHandshake() {
+  // Flush out all the response first
+  if (HasResponse()) {
+    auto write_ret = ProcessWrite();
+    if (write_ret != Transition::PROCEED) {
+      return write_ret;
+    }
   }
 
+  return SSLHandshake();
+}
+
+Transition ConnectionHandle::SSLHandshake() {
   if (conn_SSL_context == nullptr) {
     conn_SSL_context = SSL_new(PelotonServer::ssl_context);
     if (conn_SSL_context == nullptr) {

src/network/postgres_protocol_handler.cpp

@@ -1035,12 +1035,11 @@ void PostgresProtocolHandler::ProcessStartupPacket(InputPacket *pkt,
 ProcessResult PostgresProtocolHandler::Process(Buffer &rbuf,
                                                const bool ssl_able,
                                                const size_t thread_id) {
-  // if the packet is of startup packet format
-  if (!ParseInputPacket(rbuf, rpkt_, startup_stage_))
+  if (!ParseInputPacket(rbuf, rpkt_, init_stage_))
     return ProcessResult::MORE_DATA_REQUIRED;
 
   ProcessResult process_status;
-  if (startup_stage) {
+  if (init_stage_) {
     process_status = ProcessInitialPacket();
   } else {
     process_status = ProcessNormalPacket();