Modify high water level callback

Author: an-tao

trantor/net/TcpConnection.h

@@ -359,7 +359,12 @@ class TRANTOR_EXPORT TcpConnection
     {
         sslErrorCallback_ = std::move(cb);
     }
-
+    /**
+     * @brief Get the data length in the sending buffer. The sending buffer is
+     * in the user memory space.
+     * @note This method should be called in the right event loop.
+     */
+    virtual size_t getBufferedDataLength() const = 0;
     // TODO: These should be internal APIs
     virtual void connectEstablished() = 0;
     virtual void connectDestroyed() = 0;

trantor/net/inner/TcpConnectionImpl.cc

@@ -320,6 +320,18 @@ void TcpConnectionImpl::setTcpNoDelay(bool on)
 {
     socketPtr_->setTcpNoDelay(on);
 }
+void TcpConnectionImpl::checkBufferedDataSize()
+{
+    loop_->assertInLoopThread();
+    if (highWaterMarkCallback_)
+    {
+        auto bufferedDataSize = getBufferedDataLength();
+        if (bufferedDataSize > highWaterMarkLen_)
+        {
+            highWaterMarkCallback_(shared_from_this(), bufferedDataSize);
+        }
+    }
+}
 void TcpConnectionImpl::connectDestroyed()
 {
     loop_->assertInLoopThread();
@@ -415,21 +427,7 @@ void TcpConnectionImpl::sendInLoop(const char *buffer, size_t length)
         writeBufferList_.back()->append(static_cast<const char *>(buffer) +
                                             sendLen,
                                         length);
-        if (highWaterMarkCallback_ &&
-            writeBufferList_.back()->remainingBytes() >
-                static_cast<long long>(highWaterMarkLen_))
-        {
-            highWaterMarkCallback_(shared_from_this(),
-                                   writeBufferList_.back()->remainingBytes());
-        }
-        if (highWaterMarkCallback_ && tlsProviderPtr_ &&
-            tlsProviderPtr_->getBufferedData().readableBytes() >
-                highWaterMarkLen_)
-        {
-            highWaterMarkCallback_(
-                shared_from_this(),
-                tlsProviderPtr_->getBufferedData().readableBytes());
-        }
+        checkBufferedDataSize();
     }
 }
 // The order of data sending should be same as the order of calls of send()
@@ -598,12 +596,16 @@ void TcpConnectionImpl::sendFile(BufferNodePtr &&fileNode)
         {
             auto n = sendNodeInLoop(fileNode);
             if (fileNode->remainingBytes() > 0 && n >= 0)
+            {
                 writeBufferList_.push_back(std::move(fileNode));
+                checkBufferedDataSize();
+            }
             return;
         }
         else
         {
             writeBufferList_.push_back(std::move(fileNode));
+            checkBufferedDataSize();
         }
     }
     else
@@ -614,11 +616,15 @@ void TcpConnectionImpl::sendFile(BufferNodePtr &&fileNode)
             {
                 auto n = thisPtr->sendNodeInLoop(node);
                 if (node->remainingBytes() > 0 && n >= 0)
+                {
                     thisPtr->writeBufferList_.push_back(std::move(node));
+                    thisPtr->checkBufferedDataSize();
+                }
             }
             else
             {
                 thisPtr->writeBufferList_.push_back(std::move(node));
+                thisPtr->checkBufferedDataSize();
             }
         });
     }
@@ -634,7 +640,10 @@ void TcpConnectionImpl::sendStream(
         {
             auto n = sendNodeInLoop(node);
             if (node->remainingBytes() > 0 && n >= 0)
+            {
                 writeBufferList_.push_back(std::move(node));
+                checkBufferedDataSize();
+            }
             return;
         }
     }
@@ -647,11 +656,15 @@ void TcpConnectionImpl::sendStream(
                 {
                     auto n = thisPtr->sendNodeInLoop(node);
                     if (node->remainingBytes() > 0 && n >= 0)
+                    {
                         thisPtr->writeBufferList_.push_back(std::move(node));
+                        thisPtr->checkBufferedDataSize();
+                    }
                 }
                 else
                 {
                     thisPtr->writeBufferList_.push_back(std::move(node));
+                    thisPtr->checkBufferedDataSize();
                 }
             });
     }
@@ -959,8 +972,8 @@ AsyncStreamPtr TcpConnectionImpl::sendAsyncStream(bool disableKickoff)
             idleTimeoutBackup_ = idleTimeout_;
             idleTimeout_ = 0;
         }
-
         writeBufferList_.push_back(asyncStreamNode);
+        checkBufferedDataSize();
     }
     else
     {
@@ -984,11 +997,15 @@ AsyncStreamPtr TcpConnectionImpl::sendAsyncStream(bool disableKickoff)
             {
                 auto n = thisPtr->sendNodeInLoop(node);
                 if (n >= 0 && (node->remainingBytes() > 0 || node->available()))
+                {
                     thisPtr->writeBufferList_.push_back(std::move(node));
+                    thisPtr->checkBufferedDataSize();
+                }
             }
             else
             {
                 thisPtr->writeBufferList_.push_back(std::move(node));
+                thisPtr->checkBufferedDataSize();
             }
         });
     }

trantor/net/inner/TcpConnectionImpl.h

@@ -202,17 +202,31 @@ class TcpConnectionImpl : public TcpConnection,
         idleTimeout_ = timeout;
         timingWheel->insertEntry(timeout, entry);
     }
+    size_t getBufferedDataLength() const override
+    {
+        loop_->assertInLoopThread();
+        size_t len = 0;
+        if (tlsProviderPtr_)
+        {
+            len += tlsProviderPtr_->getBufferedData().readableBytes();
+        }
+        for (auto &node : writeBufferList_)
+        {
+            len += node->remainingBytes();
+        }
+        return len;
+    }
 
   private:
     /// Internal use only.
-
     std::weak_ptr<KickoffEntry> kickoffEntry_;
     std::weak_ptr<TimingWheel> timingWheelWeakPtr_;
     size_t idleTimeout_{0};
     size_t idleTimeoutBackup_{0};
     Date lastTimingWheelUpdateTime_;
     void extendLife();
     void sendFile(BufferNodePtr &&fileNode);
+    void checkBufferedDataSize();
 
   protected:
     enum class ConnStatus