Detect what cause OOM for OSS (finos#381)

* Detect what cause OOM for OSS

Signed-off-by: gxu <georgexu420@163.com>

* Bump Version Number

Signed-off-by: gxu <georgexu420@163.com>

* Remove LARGEST_PAYLOAD_SIZE

Signed-off-by: gxu <georgexu420@163.com>

* Do not set Bytes::alloc to be noexcept

Signed-off-by: gxu <georgexu420@163.com>

---------

Signed-off-by: gxu <georgexu420@163.com>

Author: gxuu

scaler/io/ymq/bytes.h

@@ -101,7 +101,7 @@ class Bytes {
     }
 
     [[nodiscard("Allocated Bytes is not used, likely causing a memory leak")]]
-    static Bytes alloc(size_t len) noexcept
+    static Bytes alloc(size_t len)
     {
         auto ptr = new uint8_t[len];  // we just assume the allocation will succeed
         return Bytes {ptr, len};

scaler/io/ymq/configuration.h

@@ -34,8 +34,7 @@ template <typename T>
 using MoveOnlyFunction = std::function<T>;
 #endif
 
-constexpr const uint64_t IOCP_SOCKET_CLOSED   = 4;
-constexpr const uint64_t LARGEST_PAYLOAD_SIZE = 6000'000'000'000;  // 6TB
+constexpr const uint64_t IOCP_SOCKET_CLOSED = 4;
 
 struct Configuration {
 #ifdef __linux__

scaler/io/ymq/message_connection_tcp.cpp

@@ -1,6 +1,8 @@
 
 #include "scaler/io/ymq/message_connection_tcp.h"
 
+#include <new>
+
 #include "scaler/io/ymq/configuration.h"
 
 #ifdef __linux__
@@ -146,12 +148,25 @@ std::expected<void, MessageConnectionTCP::IOError> MessageConnectionTCP::tryRead
             readTo        = (char*)&message._header + message._cursor;
             remainingSize = HEADER_SIZE - message._cursor;
         } else if (message._cursor == HEADER_SIZE) {
-            if (message._header >= LARGEST_PAYLOAD_SIZE) {
+            // NOTE: We probably need a better protocol to solve this issue completely, but this should let us pin down
+            // why OSS sometimes throws bad_alloc
+            try {
+                // On Linux, this will never happen because this function is only called when
+                // new read comes in. On other platform, this might be different.
+                if (!message._payload.data()) {
+                    message._payload = Bytes::alloc(message._header);
+                }
+                readTo        = (char*)message._payload.data();
+                remainingSize = message._payload.len();
+            } catch (const std::bad_alloc& e) {
+                _logger.log(
+                    Logger::LoggingLevel::error,
+                    "Trying to allocate ",
+                    message._header,
+                    " bytes.",
+                    " bad_alloc caught, connection closed");
                 return std::unexpected {IOError::MessageTooLarge};
             }
-            message._payload = Bytes::alloc(message._header);
-            readTo           = (char*)message._payload.data();
-            remainingSize    = message._payload.len();
         } else {
             readTo        = (char*)message._payload.data() + (message._cursor - HEADER_SIZE);
             remainingSize = message._payload.len() - (message._cursor - HEADER_SIZE);

scaler/io/ymq/message_connection_tcp.h

@@ -7,6 +7,7 @@
 
 #include "scaler/io/ymq/configuration.h"
 #include "scaler/io/ymq/io_socket.h"
+#include "scaler/io/ymq/logging.h"
 #include "scaler/io/ymq/message_connection.h"
 #include "scaler/io/ymq/tcp_operations.h"
 
@@ -87,6 +88,7 @@ class MessageConnectionTCP: public MessageConnection {
     std::queue<TcpReadOperation> _receivedReadOperations;
 
     bool _disconnect;
+    Logger _logger;
 
     constexpr static bool isCompleteMessage(const TcpReadOperation& x);
     friend void IOSocket::onConnectionIdentityReceived(MessageConnectionTCP* conn) noexcept;

scaler/version.txt

@@ -1 +1 @@
-1.12.17
+1.12.18