Ubshm transport fix 2

.gitignore

@@ -53,6 +53,7 @@ graphify-out/
 /protoc-gen-mcpack*/
 
 # Ignore bazel files
+MODULE.bazel.lock
 bazel-bin
 bazel-out
 bazel-testlogs

docs/cn/ubring.md

@@ -187,7 +187,7 @@ UBRing 架构包含以下组件：
 
 ### 定时器管理
 
-UBRing 使用高精度定时器系统 (`timer_mgr.cpp`) 进行连接管理和超时处理，支持 epoll（Linux）和 kqueue（macOS）。
+UBRing 的连接管理和超时处理基于 bthread 定时器（`bthread_timer_add`/`bthread_timer_del`，由 `timer_mgr.cpp` 封装）。非阻塞的 `UbrTimerDel` 可在回调内自删，此时句柄会保留到该回调返回，因此并发的外部 `UbrTimerDelAndWait` 一定能等到它；销毁回调参数所属资源的外部路径用 `UbrTimerDelAndWait` 等待运行中回调退出；一次性定时器触发后自动清理句柄。进入关闭处理时会带上调用方看到的代际，并在 manager 锁内原子认领槽位，因此按旧快照发起的关闭（例如共享内存故障事件回调）不会误伤释放后被复用的新连接。定时器回调运行在进程全局的 bthread 定时线程上，必须快速返回。阻塞的共享内存清理不在这里进行：延迟清理与待回收映射的回调只仲裁清理所有权，并把任务投递给专用的清理 worker bthread（`ub_cleanup_worker.cpp`），由它把 UBS SDK 的 `unmap`/`deallocate` 调用移出定时线程。对一次性定时器，`UbrTimerDel` 返回 0 还额外保证 facade 已不再访问句柄槽位所在的存储（它会等待被定时线程派发但尚未竞争槽位的 wrapper 退出竞争），因此调用方可以释放承载该槽位的对象；延迟清理路径正是据此释放 `UbrCleanupCtl`。关闭检查定时器在链路空闲时按指数退避轮询（上限 `ub_event_queue_timer_interval_max_us`，默认 10ms），有流量或正在关闭时恢复快速轮询。
 
 ## 参考资料
 

docs/en/client.md

@@ -137,7 +137,7 @@ Similar with "http" prefix besides that the connections will be encrypted with S
 
 Get a list of servers with the specified service-name through consul. The default address of consul is localhost:8500, which can be modified by setting -consul\_agent\_addr in gflags. The connection timeout of consul is 200ms by default, which can be modified by -consul\_connect\_timeout\_ms. 
 
-By default, [stale](https://www.consul.io/api/index.html#consistency-modes) and passing(only servers with passing in statuses are returned) are added to [parameters of the consul request]((https://www.consul.io/api/health.html#parameters-2)), which can be modified by -consul\_url\_parameter in gflags.
+By default, [stale](https://www.consul.io/api/index.html#consistency-modes) and passing(only servers with passing in statuses are returned) are added to [parameters of the consul request](https://www.consul.io/api/health.html#parameters-2), which can be modified by -consul\_url\_parameter in gflags.
 
 Except the first request to consul, the follow-up requests use the [long polling](https://www.consul.io/api/index.html#blocking-queries) feature. That is, the consul responds only when the server list is updated or the request times out. The timeout defaults to 60 seconds, which can be modified by -consul\_blocking\_query\_wait\_secs.
 

docs/en/ubring.md

@@ -186,7 +186,7 @@ The shared memory manager (`shm_mgr.cpp`) provides a unified interface for diffe
 
 ### Timer Management
 
-UBRing uses a high-precision timer system (`timer_mgr.cpp`) for connection management and timeout handling, supporting both epoll (Linux) and kqueue (macOS).
+UBRing connection management and timeout handling are built on bthread timers (`bthread_timer_add`/`bthread_timer_del`, wrapped in `timer_mgr.cpp`). The non-blocking `UbrTimerDel` is safe to call from within a timer callback; a periodic callback that deletes its own timer keeps its handle anchored until the callback returns, so a concurrent `UbrTimerDelAndWait` from a teardown path always waits for it. External teardown that frees resources reachable from the callback argument uses `UbrTimerDelAndWait` to wait out a possibly running callback; one-shot timers clean up their handle automatically when they fire. Entering a close carries the generation the caller saw and claims the pool slot atomically under the manager lock, so a close started from a stale snapshot (the shared-memory fault event, for instance) cannot hit a connection that reused the slot in the meantime. Timer callbacks run on the process-wide bthread timer thread and must return quickly. Blocking shared-memory cleanup is therefore not done there: the delayed-clear and pending-unmap callbacks only settle the cleanup ownership and post a job to a dedicated cleanup worker bthread (`ub_cleanup_worker.cpp`), which runs the UBS SDK `unmap`/`deallocate` calls off the timer thread. For a one-shot timer, a `UbrTimerDel` that returns 0 additionally guarantees that the facade is done with the storage the handle slot lives in -- it waits out a wrapper the timer thread already dispatched -- so the caller may free the object that holds the slot; the delayed-clear path frees `UbrCleanupCtl` that way. The close-check timer backs off exponentially while the link is idle (capped by `ub_event_queue_timer_interval_max_us`, 10ms by default) and returns to fast polling once there is traffic or a close in progress.
 
 ## References
 

src/brpc/global.cpp

@@ -575,7 +575,7 @@ static void GlobalInitializeOrDieImpl() {
     Protocol mongo_protocol = { ParseMongoMessage,
                                 nullptr, nullptr,
                                 ProcessMongoRequest, nullptr,
-                                nullptr, nullptr, nullptr,
+                                VerifyMongoRequest, nullptr, nullptr,
                                 CONNECTION_TYPE_POOLED, "mongo" };
     if (RegisterProtocol(PROTOCOL_MONGO, mongo_protocol) != 0) {
         exit(1);

src/brpc/grpc.cpp

@@ -169,6 +169,11 @@ void PercentDecode(const std::string& str, std::string* str_out) {
     }
 }
 
+// The gRPC wire format limits TimeoutValue to "an ASCII string of at most 8
+// digits", so this is the largest value a conforming peer can send. See
+// doc/PROTOCOL-HTTP2.md in grpc/grpc.
+static const int64_t MAX_GRPC_TIMEOUT_VALUE = 99999999;
+
 int64_t ConvertGrpcTimeoutToUS(const std::string* grpc_timeout) {
     if (!grpc_timeout || grpc_timeout->empty()) {
         return -1;
@@ -186,6 +191,14 @@ int64_t ConvertGrpcTimeoutToUS(const std::string* grpc_timeout) {
     if ((size_t)(endptr - grpc_timeout->data()) != grpc_timeout->size() - 1) {
         return -1;
     }
+    // A value outside this range is not a valid timeout, and a large one makes
+    // the conversions below overflow int64_t, so treat it as no deadline like
+    // any other malformed value. strtol() saturates at LONG_MAX/LONG_MIN for
+    // input that does not fit, which is rejected here as well. Zero keeps
+    // going through and yields an already-expired deadline, as before.
+    if (timeout_value < 0 || timeout_value > MAX_GRPC_TIMEOUT_VALUE) {
+        return -1;
+    }
     switch (*endptr) {
         case 'H':
             return timeout_value * 3600 * 1000000;

src/brpc/policy/mongo_protocol.cpp

@@ -171,6 +171,17 @@ ParseResult ParseMongoMessage(butil::IOBuf* source,
     return MakeMessage(msg);
 }
 
+bool VerifyMongoRequest(const InputMessageBase* msg_base) {
+    Server* server = (Server*)msg_base->arg();
+    if (server->options().auth) {
+        // The mongo protocol carries no credential that ServerOptions.auth
+        // could verify.
+        LOG(WARNING) << "mongo does not support authentication";
+        return false;
+    }
+    return true;
+}
+
 // Defined in baidu_rpc_protocol.cpp
 void EndRunningCallMethodInPool(
     ::google::protobuf::Service* service,

src/brpc/policy/mongo_protocol.h

@@ -31,6 +31,9 @@ ParseResult ParseMongoMessage(butil::IOBuf* source, Socket* socket, bool read_eo
 // Actions to a (client) request in mongo format
 void ProcessMongoRequest(InputMessageBase* msg);
 
+// Verify authentication information in mongo format
+bool VerifyMongoRequest(const InputMessageBase* msg);
+
 } // namespace policy
 } // namespace brpc
 

src/brpc/policy/redis_protocol.cpp

@@ -133,6 +133,15 @@ ParseResult ParseRedisMessage(butil::IOBuf* source, Socket* socket,
             // was not enabled on this port at all.
             return MakeParseError(PARSE_ERROR_TRY_OTHERS);
         }
+        if (server->options().auth != nullptr && !socket->IsAuthenticated()) {
+            // Redis commands are handled right here rather than in
+            // ProcessRedisRequest(), so they never reach InputMessenger's
+            // authentication hook, and the protocol carries no credential
+            // that ServerOptions.auth could verify.
+            LOG(WARNING) << "Reject redis request from unauthenticated "
+                         << *socket;
+            return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
+        }
         RedisConnContext* ctx = static_cast<RedisConnContext*>(socket->parsing_context());
         if (ctx == nullptr) {
             ctx = new RedisConnContext(rs);

src/brpc/policy/rtmp_protocol.cpp

@@ -3490,6 +3490,15 @@ ParseResult ParseRtmpMessage(butil::IOBuf* source, Socket *socket, bool read_eof
         if (!r.is_ok()) {
             return r;
         }
+        if (server->options().auth != nullptr && !socket->IsAuthenticated()) {
+            // RTMP messages are handled inside Feed() rather than in a
+            // Process function, so they never reach InputMessenger's
+            // authentication hook, and the protocol carries no credential
+            // that ServerOptions.auth could verify.
+            LOG(WARNING) << "Reject rtmp connection from unauthenticated "
+                         << *socket;
+            return MakeParseError(PARSE_ERROR_ABSOLUTELY_WRONG);
+        }
         rtmp_ctx = new RtmpContext(nullptr, server);
         socket->reset_parsing_context(rtmp_ctx);
         // We don't need to customize app_connect at server-side.

src/brpc/redis_command.cpp

@@ -68,29 +68,39 @@ inline size_t AppendDecimal(char* outbuf, unsigned long d) {
 // short or does not have many %. In a 100K-time call to formating of
 // "GET key1", the time spent on RedisRequest.AddCommand() are ~700ns
 // vs. ~400ns while using snprintf() vs. AppendDecimal() respectively.
-inline void AppendHeader(std::string& buf, char fc, unsigned long value) {
+inline void AppendHeader(butil::IOBuf& buf, char fc, unsigned long value) {
     char header[32];
     header[0] = fc;
     size_t len = AppendDecimal(header + 1, value);
     header[len + 1] = '\r';
     header[len + 2] = '\n';
     buf.append(header, len + 3);
 }
-inline void AppendHeader(butil::IOBuf& buf, char fc, unsigned long value) {
+
+inline int AppendHeader(butil::IOBufAppender& buf, char fc,
+                        unsigned long value) {
     char header[32];
     header[0] = fc;
     size_t len = AppendDecimal(header + 1, value);
     header[len + 1] = '\r';
     header[len + 2] = '\n';
-    buf.append(header, len + 3);
+    return buf.append(header, len + 3);
 }
 
-static void FlushComponent(std::string* out, std::string* compbuf, int* ncomp) {
-    AppendHeader(*out, '$', compbuf->size());
-    out->append(*compbuf);
-    out->append("\r\n", 2);
+static int FlushComponent(butil::IOBufAppender* out, std::string* compbuf,
+                          int* ncomp) {
+    if (AppendHeader(*out, '$', compbuf->size()) != 0 ||
+        out->append(compbuf->data(), compbuf->size()) != 0 ||
+        out->append("\r\n", 2) != 0) {
+        return -1;
+    }
     compbuf->clear();
     ++*ncomp;
+    return 0;
+}
+
+inline butil::Status AppendFailed() {
+    return butil::Status(ENOMEM, "Fail to append Redis command");
 }
 
 // Support hiredis-style format, namely everything is same with printf except
@@ -105,8 +115,7 @@ RedisCommandFormatV(butil::IOBuf* outbuf, const char* fmt, va_list ap) {
         return butil::Status(EINVAL, "Param[outbuf] or [fmt] is NULL");
     }
     const size_t fmt_len = strlen(fmt);
-    std::string nocount_buf;
-    nocount_buf.reserve(fmt_len * 3 / 2 + 16);
+    butil::IOBufAppender nocount_buf;
     std::string compbuf;  // A component
     compbuf.reserve(fmt_len + 16);
     const char* c = fmt;
@@ -123,15 +132,19 @@ RedisCommandFormatV(butil::IOBuf* outbuf, const char* fmt, va_list ap) {
                 if (quote_char) {
                     compbuf.push_back(*c);
                 } else if (!compbuf.empty() || format_arg_in_comp) {
-                    FlushComponent(&nocount_buf, &compbuf, &ncomponent);
+                    if (FlushComponent(&nocount_buf, &compbuf, &ncomponent) != 0) {
+                        return AppendFailed();
+                    }
                     format_arg_in_comp = false;
                 }
             } else if (*c == '"' || *c == '\'') {  // Check quotation.
                 if (!quote_char) {  // begin quote
                     quote_char = *c;
                     quote_pos = c;
                     if (!compbuf.empty() || format_arg_in_comp) {
-                        FlushComponent(&nocount_buf, &compbuf, &ncomponent);
+                        if (FlushComponent(&nocount_buf, &compbuf, &ncomponent) != 0) {
+                            return AppendFailed();
+                        }
                         format_arg_in_comp = false;
                     }
                 } else if (quote_char == *c) {
@@ -144,7 +157,9 @@ RedisCommandFormatV(butil::IOBuf* outbuf, const char* fmt, va_list ap) {
                         compbuf.push_back(*c);
                     } else { // end quote
                         quote_char = 0;
-                        FlushComponent(&nocount_buf, &compbuf, &ncomponent);
+                        if (FlushComponent(&nocount_buf, &compbuf, &ncomponent) != 0) {
+                            return AppendFailed();
+                        }
                         format_arg_in_comp = false;
                     }
                 } else {
@@ -323,7 +338,9 @@ RedisCommandFormatV(butil::IOBuf* outbuf, const char* fmt, va_list ap) {
     }
 
     if (!compbuf.empty() || format_arg_in_comp) {
-        FlushComponent(&nocount_buf, &compbuf, &ncomponent);
+        if (FlushComponent(&nocount_buf, &compbuf, &ncomponent) != 0) {
+            return AppendFailed();
+        }
         format_arg_in_comp = false;
     }
 
@@ -332,7 +349,9 @@ RedisCommandFormatV(butil::IOBuf* outbuf, const char* fmt, va_list ap) {
         "formatting of conversion specifiers)";
 
     AppendHeader(*outbuf, '*', ncomponent);
-    outbuf->append(nocount_buf);
+    if (ncomponent != 0) {
+        outbuf->append(nocount_buf.buf().movable());
+    }
     return butil::Status::OK();
 }
 
@@ -350,8 +369,7 @@ RedisCommandNoFormat(butil::IOBuf* outbuf, const butil::StringPiece& cmd) {
         return butil::Status(EINVAL, "Param[outbuf] or [cmd] is NULL");
     }
     const size_t cmd_len = cmd.size();
-    std::string nocount_buf;
-    nocount_buf.reserve(cmd_len * 3 / 2 + 16);
+    butil::IOBufAppender nocount_buf;
     std::string compbuf;  // A component
     compbuf.reserve(cmd_len + 16);
     int ncomponent = 0;
@@ -362,14 +380,18 @@ RedisCommandNoFormat(butil::IOBuf* outbuf, const butil::StringPiece& cmd) {
             if (quote_char) {
                 compbuf.push_back(*c);
             } else if (!compbuf.empty()) {
-                FlushComponent(&nocount_buf, &compbuf, &ncomponent);
+                if (FlushComponent(&nocount_buf, &compbuf, &ncomponent) != 0) {
+                    return AppendFailed();
+                }
             }
         } else if (*c == '"' || *c == '\'') {  // Check quotation.
             if (!quote_char) {  // begin quote
                 quote_char = *c;
                 quote_pos = c;
                 if (!compbuf.empty()) {
-                    FlushComponent(&nocount_buf, &compbuf, &ncomponent);
+                    if (FlushComponent(&nocount_buf, &compbuf, &ncomponent) != 0) {
+                        return AppendFailed();
+                    }
                 }
             } else if (quote_char == *c) {
                 const char last_char = (compbuf.empty() ? 0 : compbuf.back());
@@ -381,7 +403,9 @@ RedisCommandNoFormat(butil::IOBuf* outbuf, const butil::StringPiece& cmd) {
                     compbuf.push_back(*c);
                 } else { // end quote
                     quote_char = 0;
-                    FlushComponent(&nocount_buf, &compbuf, &ncomponent);
+                    if (FlushComponent(&nocount_buf, &compbuf, &ncomponent) != 0) {
+                        return AppendFailed();
+                    }
                 }
             } else {
                 compbuf.push_back(*c);
@@ -400,11 +424,15 @@ RedisCommandNoFormat(butil::IOBuf* outbuf, const butil::StringPiece& cmd) {
     }
 
     if (!compbuf.empty()) {
-        FlushComponent(&nocount_buf, &compbuf, &ncomponent);
+        if (FlushComponent(&nocount_buf, &compbuf, &ncomponent) != 0) {
+            return AppendFailed();
+        }
     }
 
     AppendHeader(*outbuf, '*', ncomponent);
-    outbuf->append(nocount_buf);
+    if (ncomponent != 0) {
+        outbuf->append(nocount_buf.buf().movable());
+    }
     return butil::Status::OK();
 }
 

src/brpc/selective_channel.cpp

@@ -142,6 +142,9 @@ friend class SubDone;
     ~Sender() { Clear(); }
     int IssueRPC(int64_t start_realtime_us);
     Resource PopFree();
+    // Called with the main call-id lock held. On false, the lock has been
+    // released or cleanup completed; the caller must return immediately
+    // without accessing the Sender, SubDone, or main Controller.
     bool PushFree(const Resource& r);
     const Controller* SubController(int index) const;
     void Run();
@@ -485,8 +488,15 @@ inline Resource Sender::PopFree() {
 inline bool Sender::PushFree(const Resource& r) {
     if (_nfree < 2) {
         _free_resources[_nfree++] = r;
-        if (_finished && _nfree == _nalloc) {
-            Clear();
+        if (_finished) {
+            if (_nfree == _nalloc) {
+                Clear();
+            } else {
+                // Let the remaining sub callbacks return their resources.
+                const CallId cid = _main_cntl->call_id();
+                CHECK_EQ(0, bthread_id_unlock(cid));
+            }
+            // Clear() or another callback after unlock may destroy us.
             return false;
         }
         return true;

src/brpc/ubshm/common/common.h

@@ -24,9 +24,6 @@
 #include "butil/compiler_specific.h"
 #include "butil/logging.h"
 
-#define LIKELY(x) BAIDU_LIKELY(x)
-#define UNLIKELY(x) BAIDU_UNLIKELY(x)
-
 #ifndef UNREFERENCE_PARAM
 #define UNREFERENCE_PARAM(x) ((void)(x))
 #endif
@@ -42,11 +39,11 @@
 #endif
 
 #ifdef __cplusplus
-#include <atomic>
-using AtomicInt = std::atomic<int>;
-using AtomicBool = std::atomic<bool>;
-using AtomicUintFast64 = std::atomic<uint_fast64_t>;
-using AtomicUintFast8 = std::atomic<uint_fast8_t>;
+#include "butil/atomicops.h"
+using AtomicInt = butil::atomic<int>;
+using AtomicBool = butil::atomic<bool>;
+using AtomicUintFast64 = butil::atomic<uint_fast64_t>;
+using AtomicUintFast8 = butil::atomic<uint_fast8_t>;
 #define ATOMIC_INIT(var, value) var.store(value)
 #define ATOMIC_STORE(var, value) var.store(value)
 #define ATOMIC_LOAD(var) var.load()
@@ -116,6 +113,7 @@ static inline int Copy64Byte(int8_t *dst, int8_t *src) {
 
 #define SEC_TO_NSEC 1000000000
 #define MSEC_TO_NSEC 1000000
+#define SEC_TO_USEC 1000000
 #define USEC_TO_NSEC 1000
 #define MSEC_TO_SEC 1000
 #define MAX_IP_PORT_STR_LEN 23

src/brpc/ubshm/common/thread_lock.h

@@ -30,7 +30,7 @@ extern "C" {
 
 static inline void UnlockMutex(pthread_mutex_t **mtx)
 {
-    if (LIKELY(mtx != nullptr && *mtx != nullptr)) {
+    if (BAIDU_LIKELY(mtx != nullptr && *mtx != nullptr)) {
         pthread_mutex_unlock(*mtx);
     } else {
         LOG(ERROR) << "Invalid input for mtx.";