Make non-blocking peek optional

lib/cpp/src/thrift/server/TNonblockingServer.cpp

@@ -479,7 +479,7 @@ void TNonblockingServer::TConnection::workSocket() {
     // data sitting in their internal buffers and from libevent's perspective, there is no further data available. In
     // that case, not having this workSocket() call here would result in a hang as we will never get to work the socket,
     // despite having more data.
-    if (tSocket_->hasPendingDataToRead())
+    if (tSocket_->peek(true))
     {
         workSocket();
     }

lib/cpp/src/thrift/transport/TSSLSocket.cpp

@@ -249,17 +249,6 @@ TSSLSocket::~TSSLSocket() {
   close();
 }
 
-bool TSSLSocket::hasPendingDataToRead() {
-  if (!isOpen()) {
-    return false;
-  }
-  initializeHandshake();
-  if (!checkHandshake())
-    throw TSSLException("SSL_peek: Handshake is not completed");
-  // data may be available in SSL buffers (note: SSL_pending does not have a failure mode)
-  return TSocket::hasPendingDataToRead() || SSL_pending(ssl_) > 0;
-}
-
 void TSSLSocket::init() {
   handshakeCompleted_ = false;
   readRetryCount_ = 0;
@@ -283,13 +272,19 @@ bool TSSLSocket::isOpen() {
 /*
  * Note: This method is not libevent safe.
 */
-bool TSSLSocket::peek() {
+bool TSSLSocket::peek(bool nonBlocking) {
   if (!isOpen()) {
     return false;
   }
   initializeHandshake();
   if (!checkHandshake())
     throw TSSLException("SSL_peek: Handshake is not completed");
+  if (nonBlocking)
+  {
+    // data may be available in SSL buffers (note: SSL_pending does not have a failure mode)
+    return TSocket::peek(true) || SSL_pending(ssl_) > 0;
+  }
+
   int rc;
   uint8_t byte;
   do {

lib/cpp/src/thrift/transport/TSSLSocket.h

@@ -75,10 +75,9 @@ class TSSLSocket : public TSocket {
    * TTransport interface.
    */
   bool isOpen();
-  bool peek();
+  bool peek(bool nonBlocking=false);
   void open();
   void close();
-  bool hasPendingDataToRead();
   uint32_t read(uint8_t* buf, uint32_t len);
   void write(const uint8_t* buf, uint32_t len);
   uint32_t write_partial(const uint8_t* buf, uint32_t len);

lib/cpp/src/thrift/transport/TSocket.cpp

@@ -171,10 +171,6 @@ TSocket::~TSocket() {
 }
 
 bool TSocket::hasPendingDataToRead() {
-  if (!isOpen()) {
-    return false;
-  }
-
   int32_t retries = 0;
   THRIFT_IOCTL_SOCKET_NUM_BYTES_TYPE numBytesAvailable;
 try_again:
@@ -194,10 +190,14 @@ bool TSocket::isOpen() {
   return (socket_ != THRIFT_INVALID_SOCKET);
 }
 
-bool TSocket::peek() {
+bool TSocket::peek(bool nonBlocking) {
   if (!isOpen()) {
     return false;
   }
+  if (nonBlocking)
+  {
+    return hasPendingDataToRead();
+  }
   if (interruptListener_) {
     for (int retries = 0;;) {
       struct THRIFT_POLLFD fds[2];

lib/cpp/src/thrift/transport/TSocket.h

@@ -84,9 +84,11 @@ class TSocket : public TVirtualTransport<TSocket> {
   virtual bool isOpen();
 
   /**
-   * Calls select on the socket to see if there is more data available.
+   * Peeks at the socket to determine if there is at least 1 byte of data available to read
+   *
+   * @param nonBlocking whether the operation is non-blocking or not
    */
-  virtual bool peek();
+  virtual bool peek(bool nonBlocking=false);
 
   /**
    * Creates and opens the UNIX socket.
@@ -100,17 +102,6 @@ class TSocket : public TVirtualTransport<TSocket> {
    */
   virtual void close();
 
-  /**
-   * Determines whether there is pending data to read or not.
-   *
-   * This call does not block.
-   * \throws TTransportException of types:
-   *           NOT_OPEN means the socket has been closed
-   *           UNKNOWN means something unexpected happened
-   * \returns true if there is pending data to read, false otherwise
-   */
-  virtual bool hasPendingDataToRead();
-
   /**
    * Reads from the underlying socket.
    * \returns the number of bytes read or 0 indicates EOF
@@ -341,8 +332,9 @@ class TSocket : public TVirtualTransport<TSocket> {
   static bool useLowMinRto_;
 
 private:
-  void unix_open();
+  bool hasPendingDataToRead();
   void local_open();
+  void unix_open();
 };
 }
 }