third_party/brpc
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases 1 Wiki Activity
0cf6c335b8
brpc/test/brpc_socket_unittest.cpp
yangzuhao 0cf6c335b8 brpc源码
2022-12-14 20:13:26 +08:00

1019 lines
35 KiB
C++
Raw Blame History

// 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.
// brpc - A framework to host and access services throughout Baidu.
// Date: Sun Jul 13 15:04:18 CST 2014
#include <sys/types.h>
#include <sys/socket.h>
#include <fcntl.h> // F_GETFD
#include <gtest/gtest.h>
#include <gflags/gflags.h>
#include "butil/gperftools_profiler.h"
#include "butil/time.h"
#include "butil/macros.h"
#include "butil/fd_utility.h"
#include "bthread/unstable.h"
#include "bthread/task_control.h"
#include "brpc/socket.h"
#include "brpc/errno.pb.h"
#include "brpc/acceptor.h"
#include "brpc/policy/hulu_pbrpc_protocol.h"
#include "brpc/policy/most_common_message.h"
#include "brpc/policy/http_rpc_protocol.h"
#include "brpc/nshead.h"
#include "brpc/server.h"
#include "brpc/channel.h"
#include "brpc/controller.h"
#include "health_check.pb.h"
#if defined(OS_MACOSX)
#include <sys/event.h>
#endif
#define CONNECT_IN_KEEPWRITE 1;
namespace bthread {
extern TaskControl* g_task_control;
}
namespace brpc {
DECLARE_int32(health_check_interval);
}
void EchoProcessHuluRequest(brpc::InputMessageBase* msg_base);
int main(int argc, char* argv[]) {
testing::InitGoogleTest(&argc, argv);
GFLAGS_NS::ParseCommandLineFlags(&argc, &argv, true);
brpc::Protocol dummy_protocol =
{ brpc::policy::ParseHuluMessage,
brpc::SerializeRequestDefault,
brpc::policy::PackHuluRequest,
EchoProcessHuluRequest, EchoProcessHuluRequest,
NULL, NULL, NULL,
brpc::CONNECTION_TYPE_ALL, "dummy_hulu" };
EXPECT_EQ(0, RegisterProtocol((brpc::ProtocolType)30, dummy_protocol));
return RUN_ALL_TESTS();
}
struct WaitData {
bthread_id_t id;
int error_code;
std::string error_text;
WaitData() : id(INVALID_BTHREAD_ID), error_code(0) {}
};
int OnWaitIdReset(bthread_id_t id, void* data, int error_code,
const std::string& error_text) {
static_cast<WaitData*>(data)->id = id;
static_cast<WaitData*>(data)->error_code = error_code;
static_cast<WaitData*>(data)->error_text = error_text;
return bthread_id_unlock_and_destroy(id);
}
class SocketTest : public ::testing::Test{
protected:
SocketTest(){
};
virtual ~SocketTest(){};
virtual void SetUp() {
};
virtual void TearDown() {
};
};
brpc::Socket* global_sock = NULL;
class CheckRecycle : public brpc::SocketUser {
void BeforeRecycle(brpc::Socket* s) {
ASSERT_TRUE(global_sock);
ASSERT_EQ(global_sock, s);
global_sock = NULL;
delete this;
}
};
TEST_F(SocketTest, not_recycle_until_zero_nref) {
std::cout << "sizeof(Socket)=" << sizeof(brpc::Socket) << std::endl;
int fds[2];
ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds));
brpc::SocketId id = 8888;
butil::EndPoint dummy;
ASSERT_EQ(0, str2endpoint("192.168.1.26:8080", &dummy));
brpc::SocketOptions options;
options.fd = fds[1];
options.remote_side = dummy;
options.user = new CheckRecycle;
ASSERT_EQ(0, brpc::Socket::Create(options, &id));
{
brpc::SocketUniquePtr s;
ASSERT_EQ(0, brpc::Socket::Address(id, &s));
global_sock = s.get();
ASSERT_TRUE(s.get());
ASSERT_EQ(fds[1], s->fd());
ASSERT_EQ(dummy, s->remote_side());
ASSERT_EQ(id, s->id());
ASSERT_EQ(0, s->SetFailed());
ASSERT_EQ(s.get(), global_sock);
}
ASSERT_EQ((brpc::Socket*)NULL, global_sock);
close(fds[0]);
brpc::SocketUniquePtr ptr;
ASSERT_EQ(-1, brpc::Socket::Address(id, &ptr));
}
butil::atomic<int> winner_count(0);
const int AUTH_ERR = -9;
void* auth_fighter(void* arg) {
bthread_usleep(10000);
int auth_error = 0;
brpc::Socket* s = (brpc::Socket*)arg;
if (s->FightAuthentication(&auth_error) == 0) {
winner_count.fetch_add(1);
s->SetAuthentication(AUTH_ERR);
} else {
EXPECT_EQ(AUTH_ERR, auth_error);
}
return NULL;
}
TEST_F(SocketTest, authentication) {
brpc::SocketId id;
brpc::SocketOptions options;
ASSERT_EQ(0, brpc::Socket::Create(options, &id));
brpc::SocketUniquePtr s;
ASSERT_EQ(0, brpc::Socket::Address(id, &s));
bthread_t th[64];
for (size_t i = 0; i < ARRAY_SIZE(th); ++i) {
ASSERT_EQ(0, bthread_start_urgent(&th[i], NULL, auth_fighter, s.get()));
}
for (size_t i = 0; i < ARRAY_SIZE(th); ++i) {
ASSERT_EQ(0, bthread_join(th[i], NULL));
}
// Only one fighter wins
ASSERT_EQ(1, winner_count.load());
// Fight after signal is OK
int auth_error = 0;
ASSERT_NE(0, s->FightAuthentication(&auth_error));
ASSERT_EQ(AUTH_ERR, auth_error);
// Socket has been `SetFailed' when authentication failed
ASSERT_TRUE(brpc::Socket::Address(s->id(), NULL));
}
static butil::atomic<int> g_called_seq(1);
class MyMessage : public brpc::SocketMessage {
public:
MyMessage(const char* str, size_t len, int* called = NULL)
: _str(str), _len(len), _called(called) {}
private:
butil::Status AppendAndDestroySelf(butil::IOBuf* out_buf, brpc::Socket*) {
out_buf->append(_str, _len);
if (_called) {
*_called = g_called_seq.fetch_add(1, butil::memory_order_relaxed);
}
delete this;
return butil::Status::OK();
};
const char* _str;
size_t _len;
int* _called;
};
class MyErrorMessage : public brpc::SocketMessage {
public:
explicit MyErrorMessage(const butil::Status& st) : _status(st) {}
private:
butil::Status AppendAndDestroySelf(butil::IOBuf*, brpc::Socket*) {
return _status;
};
butil::Status _status;
};
TEST_F(SocketTest, single_threaded_write) {
int fds[2];
ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds));
brpc::SocketId id = 8888;
butil::EndPoint dummy;
ASSERT_EQ(0, str2endpoint("192.168.1.26:8080", &dummy));
brpc::SocketOptions options;
options.fd = fds[1];
options.remote_side = dummy;
options.user = new CheckRecycle;
ASSERT_EQ(0, brpc::Socket::Create(options, &id));
{
brpc::SocketUniquePtr s;
ASSERT_EQ(0, brpc::Socket::Address(id, &s));
global_sock = s.get();
ASSERT_TRUE(s.get());
ASSERT_EQ(fds[1], s->fd());
ASSERT_EQ(dummy, s->remote_side());
ASSERT_EQ(id, s->id());
const int BATCH = 5;
for (size_t i = 0; i < 20; ++i) {
char buf[32 * BATCH];
size_t len = snprintf(buf, sizeof(buf), "hello world! %lu", i);
if (i % 4 == 0) {
brpc::SocketMessagePtr<MyMessage> msg(new MyMessage(buf, len));
ASSERT_EQ(0, s->Write(msg));
} else if (i % 4 == 1) {
brpc::SocketMessagePtr<MyErrorMessage> msg(
new MyErrorMessage(butil::Status(EINVAL, "Invalid input")));
bthread_id_t wait_id;
WaitData data;
ASSERT_EQ(0, bthread_id_create2(&wait_id, &data, OnWaitIdReset));
brpc::Socket::WriteOptions wopt;
wopt.id_wait = wait_id;
ASSERT_EQ(0, s->Write(msg, &wopt));
ASSERT_EQ(0, bthread_id_join(wait_id));
ASSERT_EQ(wait_id.value, data.id.value);
ASSERT_EQ(EINVAL, data.error_code);
ASSERT_EQ("Invalid input", data.error_text);
continue;
} else if (i % 4 == 2) {
int seq[BATCH] = {};
brpc::SocketMessagePtr<MyMessage> msgs[BATCH];
// re-print the buffer.
len = 0;
for (int j = 0; j < BATCH; ++j) {
if (j % 2 == 0) {
// Empty message, should be skipped.
msgs[j].reset(new MyMessage(buf+len, 0, &seq[j]));
} else {
size_t sub_len = snprintf(
buf+len, sizeof(buf)-len, "hello world! %lu.%d", i, j);
msgs[j].reset(new MyMessage(buf+len, sub_len, &seq[j]));
len += sub_len;
}
}
for (size_t i = 0; i < BATCH; ++i) {
ASSERT_EQ(0, s->Write(msgs[i]));
}
for (int j = 1; j < BATCH; ++j) {
ASSERT_LT(seq[j-1], seq[j]) << "j=" << j;
}
} else {
butil::IOBuf src;
src.append(buf);
ASSERT_EQ(len, src.length());
ASSERT_EQ(0, s->Write(&src));
ASSERT_TRUE(src.empty());
}
char dest[sizeof(buf)];
ASSERT_EQ(len, (size_t)read(fds[0], dest, sizeof(dest)));
ASSERT_EQ(0, memcmp(buf, dest, len));
}
ASSERT_EQ(0, s->SetFailed());
}
ASSERT_EQ((brpc::Socket*)NULL, global_sock);
close(fds[0]);
}
void EchoProcessHuluRequest(brpc::InputMessageBase* msg_base) {
brpc::DestroyingPtr<brpc::policy::MostCommonMessage> msg(
static_cast<brpc::policy::MostCommonMessage*>(msg_base));
butil::IOBuf buf;
buf.append(msg->meta);
buf.append(msg->payload);
ASSERT_EQ(0, msg->socket()->Write(&buf));
}
class MyConnect : public brpc::AppConnect {
public:
MyConnect() : _done(NULL), _data(NULL), _called_start_connect(false) {}
void StartConnect(const brpc::Socket*,
void (*done)(int err, void* data),
void* data) {
LOG(INFO) << "Start application-level connect";
_done = done;
_data = data;
_called_start_connect = true;
}
void StopConnect(brpc::Socket*) {
LOG(INFO) << "Stop application-level connect";
}
void MakeConnectDone() {
_done(0, _data);
}
bool is_start_connect_called() const { return _called_start_connect; }
private:
void (*_done)(int err, void* data);
void* _data;
bool _called_start_connect;
};
TEST_F(SocketTest, single_threaded_connect_and_write) {
// FIXME(gejun): Messenger has to be new otherwise quitting may crash.
brpc::Acceptor* messenger = new brpc::Acceptor;
const brpc::InputMessageHandler pairs[] = {
{ brpc::policy::ParseHuluMessage,
EchoProcessHuluRequest, NULL, NULL, "dummy_hulu" }
};
butil::EndPoint point(butil::IP_ANY, 7878);
int listening_fd = tcp_listen(point);
ASSERT_TRUE(listening_fd > 0);
butil::make_non_blocking(listening_fd);
ASSERT_EQ(0, messenger->AddHandler(pairs[0]));
ASSERT_EQ(0, messenger->StartAccept(listening_fd, -1, NULL));
brpc::SocketId id = 8888;
brpc::SocketOptions options;
options.remote_side = point;
std::shared_ptr<MyConnect> my_connect = std::make_shared<MyConnect>();
options.app_connect = my_connect;
options.user = new CheckRecycle;
ASSERT_EQ(0, brpc::Socket::Create(options, &id));
{
brpc::SocketUniquePtr s;
ASSERT_EQ(0, brpc::Socket::Address(id, &s));
global_sock = s.get();
ASSERT_TRUE(s.get());
ASSERT_EQ(-1, s->fd());
ASSERT_EQ(point, s->remote_side());
ASSERT_EQ(id, s->id());
for (size_t i = 0; i < 20; ++i) {
char buf[64];
const size_t meta_len = 4;
*(uint32_t*)(buf + 12) = *(uint32_t*)"Meta";
const size_t len = snprintf(buf + 12 + meta_len,
sizeof(buf) - 12 - meta_len,
"hello world! %lu", i);
memcpy(buf, "HULU", 4);
// HULU uses host byte order directly...
*(uint32_t*)(buf + 4) = len + meta_len;
*(uint32_t*)(buf + 8) = meta_len;
int called = 0;
if (i % 2 == 0) {
brpc::SocketMessagePtr<MyMessage> msg(
new MyMessage(buf, 12 + meta_len + len, &called));
ASSERT_EQ(0, s->Write(msg));
} else {
butil::IOBuf src;
src.append(buf, 12 + meta_len + len);
ASSERT_EQ(12 + meta_len + len, src.length());
ASSERT_EQ(0, s->Write(&src));
ASSERT_TRUE(src.empty());
}
if (i == 0) {
// connection needs to be established at first time.
// Should be intentionally blocked in app_connect.
bthread_usleep(10000);
ASSERT_TRUE(my_connect->is_start_connect_called());
ASSERT_LT(0, s->fd()); // already tcp connected
ASSERT_EQ(0, called); // request is not serialized yet.
my_connect->MakeConnectDone();
ASSERT_LT(0, called); // serialized
}
int64_t start_time = butil::gettimeofday_us();
while (s->fd() < 0) {
bthread_usleep(1000);
ASSERT_LT(butil::gettimeofday_us(), start_time + 1000000L) << "Too long!";
}
#if defined(OS_LINUX)
ASSERT_EQ(0, bthread_fd_wait(s->fd(), EPOLLIN));
#elif defined(OS_MACOSX)
ASSERT_EQ(0, bthread_fd_wait(s->fd(), EVFILT_READ));
#endif
char dest[sizeof(buf)];
ASSERT_EQ(meta_len + len, (size_t)read(s->fd(), dest, sizeof(dest)));
ASSERT_EQ(0, memcmp(buf + 12, dest, meta_len + len));
}
ASSERT_EQ(0, s->SetFailed());
}
ASSERT_EQ((brpc::Socket*)NULL, global_sock);
// The id is invalid.
brpc::SocketUniquePtr ptr;
ASSERT_EQ(-1, brpc::Socket::Address(id, &ptr));
messenger->StopAccept(0);
ASSERT_EQ(-1, messenger->listened_fd());
ASSERT_EQ(-1, fcntl(listening_fd, F_GETFD));
ASSERT_EQ(EBADF, errno);
}
#define NUMBER_WIDTH 16
struct WriterArg {
size_t times;
size_t offset;
brpc::SocketId socket_id;
};
void* FailedWriter(void* void_arg) {
WriterArg* arg = static_cast<WriterArg*>(void_arg);
brpc::SocketUniquePtr sock;
if (brpc::Socket::Address(arg->socket_id, &sock) < 0) {
printf("Fail to address SocketId=%" PRIu64 "\n", arg->socket_id);
return NULL;
}
char buf[32];
for (size_t i = 0; i < arg->times; ++i) {
bthread_id_t id;
EXPECT_EQ(0, bthread_id_create(&id, NULL, NULL));
snprintf(buf, sizeof(buf), "%0" BAIDU_SYMBOLSTR(NUMBER_WIDTH) "lu",
i + arg->offset);
butil::IOBuf src;
src.append(buf);
brpc::Socket::WriteOptions wopt;
wopt.id_wait = id;
sock->Write(&src, &wopt);
EXPECT_EQ(0, bthread_id_join(id));
// Only the first connect can see ECONNREFUSED and then
// calls `SetFailed' making others' error_code=EINVAL
//EXPECT_EQ(ECONNREFUSED, error_code);
}
return NULL;
}
TEST_F(SocketTest, fail_to_connect) {
const size_t REP = 10;
butil::EndPoint point(butil::IP_ANY, 7563/*not listened*/);
brpc::SocketId id = 8888;
brpc::SocketOptions options;
options.remote_side = point;
options.user = new CheckRecycle;
ASSERT_EQ(0, brpc::Socket::Create(options, &id));
{
brpc::SocketUniquePtr s;
ASSERT_EQ(0, brpc::Socket::Address(id, &s));
global_sock = s.get();
ASSERT_TRUE(s.get());
ASSERT_EQ(-1, s->fd());
ASSERT_EQ(point, s->remote_side());
ASSERT_EQ(id, s->id());
pthread_t th[8];
WriterArg args[ARRAY_SIZE(th)];
for (size_t i = 0; i < ARRAY_SIZE(th); ++i) {
args[i].times = REP;
args[i].offset = i * REP;
args[i].socket_id = id;
ASSERT_EQ(0, pthread_create(&th[i], NULL, FailedWriter, &args[i]));
}
for (size_t i = 0; i < ARRAY_SIZE(th); ++i) {
ASSERT_EQ(0, pthread_join(th[i], NULL));
}
ASSERT_EQ(-1, s->SetFailed()); // already SetFailed
ASSERT_EQ(-1, s->fd());
}
// KeepWrite is possibly still running.
int64_t start_time = butil::gettimeofday_us();
while (global_sock != NULL) {
bthread_usleep(1000);
ASSERT_LT(butil::gettimeofday_us(), start_time + 1000000L) << "Too long!";
}
ASSERT_EQ(-1, brpc::Socket::Status(id));
// The id is invalid.
brpc::SocketUniquePtr ptr;
ASSERT_EQ(-1, brpc::Socket::Address(id, &ptr));
}
TEST_F(SocketTest, not_health_check_when_nref_hits_0) {
brpc::SocketId id = 8888;
butil::EndPoint point(butil::IP_ANY, 7584/*not listened*/);
brpc::SocketOptions options;
options.remote_side = point;
options.user = new CheckRecycle;
options.health_check_interval_s = 1/*s*/;
ASSERT_EQ(0, brpc::Socket::Create(options, &id));
{
brpc::SocketUniquePtr s;
ASSERT_EQ(0, brpc::Socket::Address(id, &s));
s->SetHCRelatedRefHeld(); // set held status
global_sock = s.get();
ASSERT_TRUE(s.get());
ASSERT_EQ(-1, s->fd());
ASSERT_EQ(point, s->remote_side());
ASSERT_EQ(id, s->id());
char buf[64];
const size_t meta_len = 4;
*(uint32_t*)(buf + 12) = *(uint32_t*)"Meta";
const size_t len = snprintf(buf + 12 + meta_len,
sizeof(buf) - 12 - meta_len,
"hello world!");
memcpy(buf, "HULU", 4);
// HULU uses host byte order directly...
*(uint32_t*)(buf + 4) = len + meta_len;
*(uint32_t*)(buf + 8) = meta_len;
butil::IOBuf src;
src.append(buf, 12 + meta_len + len);
ASSERT_EQ(12 + meta_len + len, src.length());
#ifdef CONNECT_IN_KEEPWRITE
bthread_id_t wait_id;
WaitData data;
ASSERT_EQ(0, bthread_id_create2(&wait_id, &data, OnWaitIdReset));
brpc::Socket::WriteOptions wopt;
wopt.id_wait = wait_id;
ASSERT_EQ(0, s->Write(&src, &wopt));
ASSERT_EQ(0, bthread_id_join(wait_id));
ASSERT_EQ(wait_id.value, data.id.value);
ASSERT_EQ(ECONNREFUSED, data.error_code);
ASSERT_TRUE(butil::StringPiece(data.error_text).starts_with(
"Fail to connect "));
#else
ASSERT_EQ(-1, s->Write(&src));
ASSERT_EQ(ECONNREFUSED, errno);
#endif
ASSERT_TRUE(src.empty());
ASSERT_EQ(-1, s->fd());
}
// HealthCheckThread is possibly still running. Spin until global_sock
// is NULL(set in CheckRecycle::BeforeRecycle). Notice that you should
// not spin until Socket::Status(id) becomes -1 and assert global_sock
// to be NULL because invalidating id happens before calling BeforeRecycle.
const int64_t start_time = butil::gettimeofday_us();
while (global_sock != NULL) {
bthread_usleep(1000);
ASSERT_LT(butil::gettimeofday_us(), start_time + 1000000L);
}
ASSERT_EQ(-1, brpc::Socket::Status(id));
}
class HealthCheckTestServiceImpl : public test::HealthCheckTestService {
public:
HealthCheckTestServiceImpl()
: _sleep_flag(true) {}
virtual ~HealthCheckTestServiceImpl() {}
virtual void default_method(google::protobuf::RpcController* cntl_base,
const test::HealthCheckRequest* request,
test::HealthCheckResponse* response,
google::protobuf::Closure* done) {
brpc::ClosureGuard done_guard(done);
brpc::Controller* cntl = (brpc::Controller*)cntl_base;
if (_sleep_flag) {
bthread_usleep(510000 /* 510ms, a little bit longer than the default
timeout of health check rpc */);
}
cntl->response_attachment().append("OK");
}
bool _sleep_flag;
};
TEST_F(SocketTest, app_level_health_check) {
int old_health_check_interval = brpc::FLAGS_health_check_interval;
GFLAGS_NS::SetCommandLineOption("health_check_path", "/HealthCheckTestService");
GFLAGS_NS::SetCommandLineOption("health_check_interval", "1");
butil::EndPoint point(butil::IP_ANY, 7777);
brpc::ChannelOptions options;
options.protocol = "http";
options.max_retry = 0;
brpc::Channel channel;
ASSERT_EQ(0, channel.Init(point, &options));
{
brpc::Controller cntl;
cntl.http_request().uri() = "/";
channel.CallMethod(NULL, &cntl, NULL, NULL, NULL);
EXPECT_TRUE(cntl.Failed());
ASSERT_EQ(ECONNREFUSED, cntl.ErrorCode());
}
// 2s to make sure remote is connected by HealthCheckTask and enter the
// sending-rpc state. Because the remote is not down, so hc rpc would keep
// sending.
int listening_fd = tcp_listen(point);
bthread_usleep(2000000);
// 2s to make sure HealthCheckTask find socket is failed and correct impl
// should trigger next round of hc
close(listening_fd);
bthread_usleep(2000000);
brpc::Server server;
HealthCheckTestServiceImpl hc_service;
ASSERT_EQ(0, server.AddService(&hc_service, brpc::SERVER_DOESNT_OWN_SERVICE));
ASSERT_EQ(0, server.Start(point, NULL));
for (int i = 0; i < 4; ++i) {
// although ::connect would succeed, the stall in hc_service makes
// the health check rpc fail.
brpc::Controller cntl;
cntl.http_request().uri() = "/";
channel.CallMethod(NULL, &cntl, NULL, NULL, NULL);
ASSERT_EQ(EHOSTDOWN, cntl.ErrorCode());
bthread_usleep(1000000 /*1s*/);
}
hc_service._sleep_flag = false;
bthread_usleep(2000000 /* a little bit longer than hc rpc timeout + hc interval */);
// should recover now
{
brpc::Controller cntl;
cntl.http_request().uri() = "/";
channel.CallMethod(NULL, &cntl, NULL, NULL, NULL);
ASSERT_FALSE(cntl.Failed());
ASSERT_GT(cntl.response_attachment().size(), (size_t)0);
}
GFLAGS_NS::SetCommandLineOption("health_check_path", "");
char hc_buf[8];
snprintf(hc_buf, sizeof(hc_buf), "%d", old_health_check_interval);
GFLAGS_NS::SetCommandLineOption("health_check_interval", hc_buf);
}
TEST_F(SocketTest, health_check) {
// FIXME(gejun): Messenger has to be new otherwise quitting may crash.
brpc::Acceptor* messenger = new brpc::Acceptor;
brpc::SocketId id = 8888;
butil::EndPoint point(butil::IP_ANY, 7878);
const int kCheckInteval = 1;
brpc::SocketOptions options;
options.remote_side = point;
options.user = new CheckRecycle;
options.health_check_interval_s = kCheckInteval/*s*/;
ASSERT_EQ(0, brpc::Socket::Create(options, &id));
brpc::SocketUniquePtr s;
ASSERT_EQ(0, brpc::Socket::Address(id, &s));
s->SetHCRelatedRefHeld(); // set held status
global_sock = s.get();
ASSERT_TRUE(s.get());
ASSERT_EQ(-1, s->fd());
ASSERT_EQ(point, s->remote_side());
ASSERT_EQ(id, s->id());
int32_t nref = -1;
ASSERT_EQ(0, brpc::Socket::Status(id, &nref));
ASSERT_EQ(2, nref);
char buf[64];
const size_t meta_len = 4;
*(uint32_t*)(buf + 12) = *(uint32_t*)"Meta";
const size_t len = snprintf(buf + 12 + meta_len,
sizeof(buf) - 12 - meta_len,
"hello world!");
memcpy(buf, "HULU", 4);
// HULU uses host byte order directly...
*(uint32_t*)(buf + 4) = len + meta_len;
*(uint32_t*)(buf + 8) = meta_len;
const bool use_my_message = (butil::fast_rand_less_than(2) == 0);
brpc::SocketMessagePtr<MyMessage> msg;
int appended_msg = 0;
butil::IOBuf src;
if (use_my_message) {
LOG(INFO) << "Use MyMessage";
msg.reset(new MyMessage(buf, 12 + meta_len + len, &appended_msg));
} else {
src.append(buf, 12 + meta_len + len);
ASSERT_EQ(12 + meta_len + len, src.length());
}
#ifdef CONNECT_IN_KEEPWRITE
bthread_id_t wait_id;
WaitData data;
ASSERT_EQ(0, bthread_id_create2(&wait_id, &data, OnWaitIdReset));
brpc::Socket::WriteOptions wopt;
wopt.id_wait = wait_id;
if (use_my_message) {
ASSERT_EQ(0, s->Write(msg, &wopt));
} else {
ASSERT_EQ(0, s->Write(&src, &wopt));
}
ASSERT_EQ(0, bthread_id_join(wait_id));
ASSERT_EQ(wait_id.value, data.id.value);
ASSERT_EQ(ECONNREFUSED, data.error_code);
ASSERT_TRUE(butil::StringPiece(data.error_text).starts_with(
"Fail to connect "));
if (use_my_message) {
ASSERT_TRUE(appended_msg);
}
#else
if (use_my_message) {
ASSERT_EQ(-1, s->Write(msg));
} else {
ASSERT_EQ(-1, s->Write(&src));
}
ASSERT_EQ(ECONNREFUSED, errno);
#endif
ASSERT_TRUE(src.empty());
ASSERT_EQ(-1, s->fd());
ASSERT_TRUE(global_sock);
brpc::SocketUniquePtr invalid_ptr;
ASSERT_EQ(-1, brpc::Socket::Address(id, &invalid_ptr));
ASSERT_EQ(1, brpc::Socket::Status(id));
const brpc::InputMessageHandler pairs[] = {
{ brpc::policy::ParseHuluMessage,
EchoProcessHuluRequest, NULL, NULL, "dummy_hulu" }
};
int listening_fd = tcp_listen(point);
ASSERT_TRUE(listening_fd > 0);
butil::make_non_blocking(listening_fd);
ASSERT_EQ(0, messenger->AddHandler(pairs[0]));
ASSERT_EQ(0, messenger->StartAccept(listening_fd, -1, NULL));
int64_t start_time = butil::gettimeofday_us();
nref = -1;
while (brpc::Socket::Status(id, &nref) != 0) {
bthread_usleep(1000);
ASSERT_LT(butil::gettimeofday_us(),
start_time + kCheckInteval * 1000000L + 100000L/*100ms*/);
}
//ASSERT_EQ(2, nref);
ASSERT_TRUE(global_sock);
int fd = 0;
{
brpc::SocketUniquePtr ptr;
ASSERT_EQ(0, brpc::Socket::Address(id, &ptr));
ASSERT_NE(0, ptr->fd());
fd = ptr->fd();
}
// SetFailed again, should reconnect and succeed soon.
ASSERT_EQ(0, s->SetFailed());
ASSERT_EQ(fd, s->fd());
start_time = butil::gettimeofday_us();
while (brpc::Socket::Status(id) != 0) {
bthread_usleep(1000);
ASSERT_LT(butil::gettimeofday_us(), start_time + 1200000L);
}
ASSERT_TRUE(global_sock);
{
brpc::SocketUniquePtr ptr;
ASSERT_EQ(0, brpc::Socket::Address(id, &ptr));
ASSERT_NE(0, ptr->fd());
}
s.release()->Dereference();
// Must stop messenger before SetFailed the id otherwise HealthCheckThread
// still has chance to get reconnected and revive the id.
messenger->StopAccept(0);
ASSERT_EQ(-1, messenger->listened_fd());
ASSERT_EQ(-1, fcntl(listening_fd, F_GETFD));
ASSERT_EQ(EBADF, errno);
ASSERT_EQ(0, brpc::Socket::SetFailed(id));
// HealthCheckThread is possibly still addressing the Socket.
start_time = butil::gettimeofday_us();
while (global_sock != NULL) {
bthread_usleep(1000);
ASSERT_LT(butil::gettimeofday_us(), start_time + 1000000L);
}
ASSERT_EQ(-1, brpc::Socket::Status(id));
// The id is invalid.
brpc::SocketUniquePtr ptr;
ASSERT_EQ(-1, brpc::Socket::Address(id, &ptr));
}
void* Writer(void* void_arg) {
WriterArg* arg = static_cast<WriterArg*>(void_arg);
brpc::SocketUniquePtr sock;
if (brpc::Socket::Address(arg->socket_id, &sock) < 0) {
printf("Fail to address SocketId=%" PRIu64 "\n", arg->socket_id);
return NULL;
}
char buf[32];
for (size_t i = 0; i < arg->times; ++i) {
snprintf(buf, sizeof(buf), "%0" BAIDU_SYMBOLSTR(NUMBER_WIDTH) "lu",
i + arg->offset);
butil::IOBuf src;
src.append(buf);
if (sock->Write(&src) != 0) {
if (errno == brpc::EOVERCROWDED) {
// The buf is full, sleep a while and retry.
bthread_usleep(1000);
--i;
continue;
}
printf("Fail to write into SocketId=%" PRIu64 ", %s\n",
arg->socket_id, berror());
break;
}
}
return NULL;
}
TEST_F(SocketTest, multi_threaded_write) {
const size_t REP = 20000;
int fds[2];
for (int k = 0; k < 2; ++k) {
printf("Round %d\n", k + 1);
ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds));
pthread_t th[8];
WriterArg args[ARRAY_SIZE(th)];
std::vector<size_t> result;
result.reserve(ARRAY_SIZE(th) * REP);
brpc::SocketId id = 8888;
butil::EndPoint dummy;
ASSERT_EQ(0, str2endpoint("192.168.1.26:8080", &dummy));
brpc::SocketOptions options;
options.fd = fds[1];
options.remote_side = dummy;
options.user = new CheckRecycle;
ASSERT_EQ(0, brpc::Socket::Create(options, &id));
brpc::SocketUniquePtr s;
ASSERT_EQ(0, brpc::Socket::Address(id, &s));
s->_ssl_state = brpc::SSL_OFF;
global_sock = s.get();
ASSERT_TRUE(s.get());
ASSERT_EQ(fds[1], s->fd());
ASSERT_EQ(dummy, s->remote_side());
ASSERT_EQ(id, s->id());
butil::make_non_blocking(fds[0]);
for (size_t i = 0; i < ARRAY_SIZE(th); ++i) {
args[i].times = REP;
args[i].offset = i * REP;
args[i].socket_id = id;
ASSERT_EQ(0, pthread_create(&th[i], NULL, Writer, &args[i]));
}
if (k == 1) {
printf("sleep 100ms to block writers\n");
bthread_usleep(100000);
}
butil::IOPortal dest;
const int64_t start_time = butil::gettimeofday_us();
for (;;) {
ssize_t nr = dest.append_from_file_descriptor(fds[0], 32768);
if (nr < 0) {
if (errno == EINTR) {
continue;
}
if (EAGAIN != errno) {
ASSERT_EQ(EAGAIN, errno) << berror();
}
bthread_usleep(1000);
if (butil::gettimeofday_us() >= start_time + 2000000L) {
LOG(FATAL) << "Wait too long!";
break;
}
continue;
}
while (dest.length() >= NUMBER_WIDTH) {
char buf[NUMBER_WIDTH + 1];
dest.copy_to(buf, NUMBER_WIDTH);
buf[sizeof(buf)-1] = 0;
result.push_back(strtol(buf, NULL, 10));
dest.pop_front(NUMBER_WIDTH);
}
if (result.size() >= REP * ARRAY_SIZE(th)) {
break;
}
}
for (size_t i = 0; i < ARRAY_SIZE(th); ++i) {
ASSERT_EQ(0, pthread_join(th[i], NULL));
}
ASSERT_TRUE(dest.empty());
bthread::g_task_control->print_rq_sizes(std::cout);
std::cout << std::endl;
ASSERT_EQ(REP * ARRAY_SIZE(th), result.size())
<< "write_head=" << s->_write_head;
std::sort(result.begin(), result.end());
result.resize(std::unique(result.begin(),
result.end()) - result.begin());
ASSERT_EQ(REP * ARRAY_SIZE(th), result.size());
ASSERT_EQ(0UL, *result.begin());
ASSERT_EQ(REP * ARRAY_SIZE(th) - 1, *(result.end() - 1));
ASSERT_EQ(0, s->SetFailed());
s.release()->Dereference();
ASSERT_EQ((brpc::Socket*)NULL, global_sock);
close(fds[0]);
}
}
void* FastWriter(void* void_arg) {
WriterArg* arg = static_cast<WriterArg*>(void_arg);
brpc::SocketUniquePtr sock;
if (brpc::Socket::Address(arg->socket_id, &sock) < 0) {
printf("Fail to address SocketId=%" PRIu64 "\n", arg->socket_id);
return NULL;
}
char buf[] = "hello reader side!";
int64_t begin_ts = butil::cpuwide_time_us();
int64_t nretry = 0;
size_t c = 0;
for (; c < arg->times; ++c) {
butil::IOBuf src;
src.append(buf, 16);
if (sock->Write(&src) != 0) {
if (errno == brpc::EOVERCROWDED) {
// The buf is full, sleep a while and retry.
bthread_usleep(1000);
--c;
++nretry;
continue;
}
printf("Fail to write into SocketId=%" PRIu64 ", %s\n",
arg->socket_id, berror());
break;
}
}
int64_t end_ts = butil::cpuwide_time_us();
int64_t total_time = end_ts - begin_ts;
printf("total=%ld count=%ld nretry=%ld\n",
(long)total_time * 1000/ c, (long)c, (long)nretry);
return NULL;
}
struct ReaderArg {
int fd;
size_t nread;
};
void* reader(void* void_arg) {
ReaderArg* arg = static_cast<ReaderArg*>(void_arg);
const size_t LEN = 32768;
char* buf = (char*)malloc(LEN);
while (1) {
ssize_t nr = read(arg->fd, buf, LEN);
if (nr < 0) {
printf("Fail to read, %m\n");
return NULL;
} else if (nr == 0) {
printf("Far end closed\n");
return NULL;
}
arg->nread += nr;
}
return NULL;
}
TEST_F(SocketTest, multi_threaded_write_perf) {
const size_t REP = 1000000000;
int fds[2];
ASSERT_EQ(0, socketpair(AF_UNIX, SOCK_STREAM, 0, fds));
bthread_t th[3];
WriterArg args[ARRAY_SIZE(th)];
brpc::SocketId id = 8888;
butil::EndPoint dummy;
ASSERT_EQ(0, str2endpoint("192.168.1.26:8080", &dummy));
brpc::SocketOptions options;
options.fd = fds[1];
options.remote_side = dummy;
options.user = new CheckRecycle;
ASSERT_EQ(0, brpc::Socket::Create(options, &id));
brpc::SocketUniquePtr s;
ASSERT_EQ(0, brpc::Socket::Address(id, &s));
s->_ssl_state = brpc::SSL_OFF;
ASSERT_EQ(2, brpc::NRefOfVRef(s->_versioned_ref));
global_sock = s.get();
ASSERT_TRUE(s.get());
ASSERT_EQ(fds[1], s->fd());
ASSERT_EQ(dummy, s->remote_side());
ASSERT_EQ(id, s->id());
for (size_t i = 0; i < ARRAY_SIZE(th); ++i) {
args[i].times = REP;
args[i].offset = i * REP;
args[i].socket_id = id;
bthread_start_background(&th[i], NULL, FastWriter, &args[i]);
}
pthread_t rth;
ReaderArg reader_arg = { fds[0], 0 };
pthread_create(&rth, NULL, reader, &reader_arg);
butil::Timer tm;
ProfilerStart("write.prof");
const uint64_t old_nread = reader_arg.nread;
tm.start();
sleep(2);
tm.stop();
const uint64_t new_nread = reader_arg.nread;
ProfilerStop();
printf("tp=%" PRIu64 "M/s\n", (new_nread - old_nread) / tm.u_elapsed());
for (size_t i = 0; i < ARRAY_SIZE(th); ++i) {
args[i].times = 0;
}
for (size_t i = 0; i < ARRAY_SIZE(th); ++i) {
ASSERT_EQ(0, bthread_join(th[i], NULL));
}
ASSERT_EQ(0, s->SetFailed());
s.release()->Dereference();
pthread_join(rth, NULL);
ASSERT_EQ((brpc::Socket*)NULL, global_sock);
close(fds[0]);
}
Reference in New Issue View Git Blame Copy Permalink