Support backup request policy

src/brpc/backup_request_policy.h

@@ -0,0 +1,40 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+
+#ifndef BRPC_BACKUP_REQUEST_POLICY_H
+#define BRPC_BACKUP_REQUEST_POLICY_H
+
+#include "brpc/controller.h"
+
+namespace brpc {
+
+class BackupRequestPolicy {
+public:
+    virtual ~BackupRequestPolicy() = default;
+
+    // Return the time in milliseconds in which another request
+    // will be sent if RPC does not finish.
+    virtual int32_t GetBackupRequestMs() const = 0;
+
+    // Return true if the backup request should be sent.
+    virtual bool DoBackup(const Controller* controller) const = 0;
+};
+
+}
+
+#endif // BRPC_BACKUP_REQUEST_POLICY_H

src/brpc/channel.cpp

@@ -55,6 +55,7 @@ ChannelOptions::ChannelOptions()
     , log_succeed_without_server(true)
     , use_rdma(false)
     , auth(NULL)
+    , backup_request_policy(NULL)
     , retry_policy(NULL)
     , ns_filter(NULL)
 {}
@@ -497,6 +498,7 @@ void Channel::CallMethod(const google::protobuf::MethodDescriptor* method,
     cntl->_connect_timeout_ms = _options.connect_timeout_ms;
     if (cntl->backup_request_ms() == UNSET_MAGIC_NUM) {
         cntl->set_backup_request_ms(_options.backup_request_ms);
+        cntl->_backup_request_policy = _options.backup_request_policy;
     }
     if (cntl->connection_type() == CONNECTION_TYPE_UNKNOWN) {
         cntl->set_connection_type(_options.connection_type);
@@ -536,6 +538,7 @@ void Channel::CallMethod(const google::protobuf::MethodDescriptor* method,
         // Currently we cannot handle retry and backup request correctly
         cntl->set_max_retry(0);
         cntl->set_backup_request_ms(-1);
+        cntl->_backup_request_policy = NULL;
     }
 
     if (cntl->backup_request_ms() >= 0 &&

src/brpc/channel.h

@@ -34,6 +34,7 @@
 #include "brpc/controller.h"                // brpc::Controller
 #include "brpc/details/profiler_linker.h"
 #include "brpc/retry_policy.h"
+#include "brpc/backup_request_policy.h"
 #include "brpc/naming_service_filter.h"
 
 namespace brpc {
@@ -112,6 +113,12 @@ struct ChannelOptions {
     // Default: NULL
     const Authenticator* auth;
 
+    // Customize the backup request time and whether to send backup request.
+    // Priority: `backup_request_policy' > `backup_request_ms'.
+    // This object is NOT owned by channel and should remain valid when channel is used.
+    // Default: NULL
+    const BackupRequestPolicy* backup_request_policy;
+
     // Customize the error code that should be retried. The interface is
     // defined in src/brpc/retry_policy.h
     // This object is NOT owned by channel and should remain valid when

src/brpc/controller.cpp

@@ -258,6 +258,7 @@ void Controller::ResetPods() {
     _connection_type = CONNECTION_TYPE_UNKNOWN;
     _timeout_ms = UNSET_MAGIC_NUM;
     _backup_request_ms = UNSET_MAGIC_NUM;
+    _backup_request_policy = NULL;
     _connect_timeout_ms = UNSET_MAGIC_NUM;
     _real_timeout_ms = UNSET_MAGIC_NUM;
     _deadline_us = -1;
@@ -344,6 +345,11 @@ void Controller::set_backup_request_ms(int64_t timeout_ms) {
     }
 }
 
+int64_t Controller::backup_request_ms() const {
+    return NULL != _backup_request_policy ?
+           _backup_request_policy->GetBackupRequestMs() : _backup_request_ms;
+}
+
 void Controller::set_max_retry(int max_retry) {
     if (max_retry > MAX_RETRY_COUNT) {
         LOG(WARNING) << "Retry count can't be larger than "
@@ -1259,6 +1265,11 @@ int Controller::HandleSocketFailed(bthread_id_t id, void* data, int error_code,
                         cntl->timeout_ms(),
                         butil::endpoint2str(cntl->remote_side()).c_str());
     } else if (error_code == EBACKUPREQUEST) {
+        const BackupRequestPolicy* policy = cntl->_backup_request_policy;
+        if (NULL != policy && !policy->DoBackup(cntl)) {
+            // No need to do backup request.
+            return bthread_id_unlock(id);
+        }
         cntl->SetFailed(error_code, "Reached backup timeout=%" PRId64 "ms @%s",
                         cntl->backup_request_ms(),
                         butil::endpoint2str(cntl->remote_side()).c_str());

src/brpc/controller.h

@@ -47,6 +47,7 @@
 #include "brpc/grpc.h"
 #include "brpc/kvmap.h"
 #include "brpc/rpc_dump.h"
+#include "brpc/backup_request_policy.h"
 
 // EAUTH is defined in MAC
 #ifndef EAUTH
@@ -180,7 +181,7 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
     // Set/get the delay to send backup request in milliseconds. Use
     // ChannelOptions.backup_request_ms on unset.
     void set_backup_request_ms(int64_t timeout_ms);
-    int64_t backup_request_ms() const { return _backup_request_ms; }
+    int64_t backup_request_ms() const;
 
     // Set/get maximum times of retrying. Use ChannelOptions.max_retry on unset.
     // <=0 means no retry.
@@ -670,7 +671,7 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
     struct ClientSettings {
         int32_t timeout_ms;
         int32_t backup_request_ms;
-        int max_retry;                      
+        int max_retry;
         int32_t tos;
         ConnectionType connection_type;         
         CompressType request_compress_type;
@@ -800,6 +801,8 @@ friend void policy::ProcessThriftRequest(InputMessageBase*);
     int32_t _timeout_ms;
     int32_t _connect_timeout_ms;
     int32_t _backup_request_ms;
+    // Copied from `Channel' which might be destroyed after CallMethod.
+    const BackupRequestPolicy* _backup_request_policy;
     // If this rpc call has retry/backup request,this var save the real timeout for current call
     int64_t _real_timeout_ms;
     // Deadline of this RPC (since the Epoch in microseconds).

test/brpc_channel_unittest.cpp

@@ -310,14 +310,18 @@ class ChannelTest : public ::testing::Test{
                       bool single_server,
                       bool short_connection,
                       const brpc::Authenticator* auth = NULL,
-                      std::string connection_group = std::string()) {
+                      std::string connection_group = std::string(),
+                      bool use_backup_request_policy = false) {
         brpc::ChannelOptions opt;
         if (short_connection) {
             opt.connection_type = brpc::CONNECTION_TYPE_SHORT;
         }
         opt.auth = auth;
         opt.max_retry = 0;
         opt.connection_group = connection_group;
+        if (use_backup_request_policy) {
+            opt.backup_request_policy = &_backup_request_policy;
+        }
         if (single_server) {
             EXPECT_EQ(0, channel->Init(_ep, &opt)); 
         } else {                                                 
@@ -1918,6 +1922,110 @@ class ChannelTest : public ::testing::Test{
         StopAndJoin();
     }
 
+    void TestBackupRequest(bool single_server, bool async,
+                           bool short_connection) {
+        std::cout << " *** single=" << single_server
+                  << " async=" << async
+                  << " short=" << short_connection << std::endl;
+
+        ASSERT_EQ(0, StartAccept(_ep));
+        brpc::Channel channel;
+        SetUpChannel(&channel, single_server, short_connection);
+
+        const int RETRY_NUM = 1;
+        test::EchoRequest req;
+        test::EchoResponse res;
+        brpc::Controller cntl;
+        req.set_message(__FUNCTION__);
+
+        cntl.set_max_retry(RETRY_NUM);
+        cntl.set_backup_request_ms(10);  // 10ms
+        cntl.set_timeout_ms(100);  // 10ms
+        req.set_sleep_us(50000); // 100ms
+        CallMethod(&channel, &cntl, &req, &res, async);
+        ASSERT_EQ(0, cntl.ErrorCode()) << cntl.ErrorText();
+        ASSERT_TRUE(cntl.has_backup_request());
+        ASSERT_EQ(RETRY_NUM, cntl.retried_count());
+        bthread_usleep(70000);  // wait for the sleep task to finish
+
+        if (short_connection) {
+            // Sleep to let `_messenger' detect `Socket' being `SetFailed'
+            const int64_t start_time = butil::gettimeofday_us();
+            while (_messenger.ConnectionCount() != 0) {
+                EXPECT_LT(butil::gettimeofday_us(), start_time + 100000L/*100ms*/);
+                bthread_usleep(1000);
+            }
+        } else {
+            EXPECT_GE(1ul, _messenger.ConnectionCount());
+        }
+        StopAndJoin();
+    }
+
+    class BackupRequestPolicyImpl : public brpc::BackupRequestPolicy {
+    public:
+        int32_t GetBackupRequestMs() const override {
+            return 10;
+        }
+
+        // Return true if the backup request should be sent.
+        bool DoBackup(const brpc::Controller*) const override {
+            return backup;
+        }
+
+        bool backup{true};
+
+    };
+
+    void TestBackupRequestPolicy(bool single_server, bool async,
+                                 bool short_connection) {
+        ASSERT_EQ(0, StartAccept(_ep));
+        for (int i = 0; i < 3; ++i) {
+            bool backup = i != 1;
+            std::cout << " *** single=" << single_server
+                      << " async=" << async
+                      << " short=" << short_connection
+                      << " backup=" << backup
+                      << std::endl;
+
+            brpc::Channel channel;
+            SetUpChannel(&channel, single_server, short_connection, NULL, "", true);
+
+            const int RETRY_NUM = 1;
+            test::EchoRequest req;
+            test::EchoResponse res;
+            brpc::Controller cntl;
+            req.set_message(__FUNCTION__);
+
+            _backup_request_policy.backup = i == 0;
+            if (i == 2) {
+                // use `set_backup_request_ms'.
+                // Although _backup_request_policy.DoBackup return false, it is ignored.
+                cntl.set_backup_request_ms(10);  // 10ms
+            }
+            cntl.set_max_retry(RETRY_NUM);
+            cntl.set_timeout_ms(100);  // 100ms
+            req.set_sleep_us(50000); // 50ms
+            CallMethod(&channel, &cntl, &req, &res, async);
+            ASSERT_EQ(0, cntl.ErrorCode()) << cntl.ErrorText();
+            ASSERT_EQ(backup, cntl.has_backup_request());
+            ASSERT_EQ(backup ? RETRY_NUM : 0, cntl.retried_count());
+            bthread_usleep(70000);  // wait for the sleep task to finish
+
+            if (short_connection) {
+                // Sleep to let `_messenger' detect `Socket' being `SetFailed'
+                const int64_t start_time = butil::gettimeofday_us();
+                while (_messenger.ConnectionCount() != 0) {
+                    EXPECT_LT(butil::gettimeofday_us(), start_time + 100000L/*100ms*/);
+                    bthread_usleep(1000);
+                }
+            } else {
+                EXPECT_GE(1ul, _messenger.ConnectionCount());
+            }
+        }
+
+        StopAndJoin();
+    }
+
     butil::EndPoint _ep;
     butil::TempFile _server_list;                                        
     std::string _naming_url;
@@ -1930,6 +2038,7 @@ class ChannelTest : public ::testing::Test{
     bool _close_fd_once;
 
     MyEchoService _svc;
+    BackupRequestPolicyImpl _backup_request_policy;
 };
 
 class MyShared : public brpc::SharedObject {
@@ -2597,6 +2706,26 @@ TEST_F(ChannelTest, retry_backoff) {
     }
 }
 
+TEST_F(ChannelTest, backup_request) {
+    for (int i = 0; i <= 1; ++i) { // Flag SingleServer
+        for (int j = 0; j <= 1; ++j) { // Flag Asynchronous
+            for (int k = 0; k <= 1; ++k) { // Flag ShortConnection
+                TestBackupRequest(i, j, k);
+            }
+        }
+    }
+}
+
+TEST_F(ChannelTest, backup_request_policy) {
+    for (int i = 0; i <= 1; ++i) { // Flag SingleServer
+        for (int j = 0; j <= 1; ++j) { // Flag Asynchronous
+            for (int k = 0; k <= 1; ++k) { // Flag ShortConnection
+                TestBackupRequestPolicy(i, j, k);
+            }
+        }
+    }
+}
+
 TEST_F(ChannelTest, multiple_threads_single_channel) {
     srand(time(NULL));
     ASSERT_EQ(0, StartAccept(_ep));