uint32_t HeaderClientChannel::sendRequest(
const RpcOptions& rpcOptions,
std::unique_ptr<RequestCallback> cb,
std::unique_ptr<apache::thrift::ContextStack> ctx,
unique_ptr<IOBuf> buf) {
// cb is not allowed to be null.
DCHECK(cb);
cb->context_ = RequestContext::saveContext();
if (isSecurityPending()) {
cb->securityStart_ = std::chrono::duration_cast<Us>(
HResClock::now().time_since_epoch()).count();
afterSecurity_.push_back(
std::make_tuple(static_cast<AfterSecurityMethod>(
&HeaderClientChannel::sendRequest),
RpcOptions(rpcOptions),
std::move(cb),
std::move(ctx),
std::move(buf),
header_->releaseWriteHeaders()));
// Security always happens at the beginning of the channel, with seq id 0.
// Return sequence id expected to be generated when security is done.
if (++sendSecurityPendingSeqId_ == ResponseChannel::ONEWAY_REQUEST_ID) {
++sendSecurityPendingSeqId_;
}
return sendSecurityPendingSeqId_;
}
DestructorGuard dg(this);
// Oneway requests use a special sequence id.
// Make sure this non-oneway request doesn't use
// the oneway request ID.
if (++sendSeqId_ == ResponseChannel::ONEWAY_REQUEST_ID) {
++sendSeqId_;
}
std::chrono::milliseconds timeout(timeout_);
if (rpcOptions.getTimeout() > std::chrono::milliseconds(0)) {
timeout = rpcOptions.getTimeout();
}
auto twcb = new TwowayCallback(this,
sendSeqId_,
header_->getProtocolId(),
std::move(cb),
std::move(ctx),
timer_.get(),
timeout,
rpcOptions.getChunkTimeout());
maybeSetPriorityHeader(rpcOptions);
maybeSetTimeoutHeader(rpcOptions);
if (header_->getClientType() != THRIFT_HEADER_CLIENT_TYPE &&
header_->getClientType() != THRIFT_HEADER_SASL_CLIENT_TYPE) {
recvCallbackOrder_.push_back(sendSeqId_);
}
recvCallbacks_[sendSeqId_] = twcb;
setBaseReceivedCallback();
sendMessage(twcb, std::move(buf));
return sendSeqId_;
}
TEST(ThriftServer, FailureInjection) {
enum ExpectedFailure { NONE = 0, ERROR, TIMEOUT, DISCONNECT, END };
std::atomic<ExpectedFailure> expected(NONE);
using apache::thrift::transport::TTransportException;
class Callback : public RequestCallback {
public:
explicit Callback(const std::atomic<ExpectedFailure>* expected)
: expected_(expected) {}
private:
void requestSent() override {}
void replyReceived(ClientReceiveState&& state) override {
std::string response;
try {
TestServiceAsyncClient::recv_sendResponse(response, state);
EXPECT_EQ(NONE, *expected_);
} catch (const apache::thrift::TApplicationException& ex) {
const auto& headers = state.header()->getHeaders();
EXPECT_TRUE(headers.find("ex") != headers.end() &&
headers.find("ex")->second == kInjectedFailureErrorCode);
EXPECT_EQ(ERROR, *expected_);
} catch (...) {
ADD_FAILURE() << "Unexpected exception thrown";
}
// Now do it again with exception_wrappers.
auto ew =
TestServiceAsyncClient::recv_wrapped_sendResponse(response, state);
if (ew) {
EXPECT_TRUE(
ew.is_compatible_with<apache::thrift::TApplicationException>());
EXPECT_EQ(ERROR, *expected_);
} else {
EXPECT_EQ(NONE, *expected_);
}
}
void requestError(ClientReceiveState&& state) override {
try {
std::rethrow_exception(state.exception());
} catch (const TTransportException& ex) {
if (ex.getType() == TTransportException::TIMED_OUT) {
EXPECT_EQ(TIMEOUT, *expected_);
} else {
EXPECT_EQ(DISCONNECT, *expected_);
}
} catch (...) {
ADD_FAILURE() << "Unexpected exception thrown";
}
}
const std::atomic<ExpectedFailure>* expected_;
};
TestThriftServerFactory<TestInterface> factory;
ScopedServerThread sst(factory.create());
folly::EventBase base;
std::shared_ptr<TAsyncSocket> socket(
TAsyncSocket::newSocket(&base, *sst.getAddress()));
TestServiceAsyncClient client(HeaderClientChannel::newChannel(socket));
auto server = std::dynamic_pointer_cast<ThriftServer>(sst.getServer().lock());
CHECK(server);
SCOPE_EXIT { server->setFailureInjection(ThriftServer::FailureInjection()); };
RpcOptions rpcOptions;
rpcOptions.setTimeout(std::chrono::milliseconds(100));
for (int i = 0; i < END; ++i) {
auto exp = static_cast<ExpectedFailure>(i);
ThriftServer::FailureInjection fi;
switch (exp) {
case NONE:
break;
case ERROR:
fi.errorFraction = 1;
break;
case TIMEOUT:
fi.dropFraction = 1;
break;
case DISCONNECT:
fi.disconnectFraction = 1;
break;
case END:
LOG(FATAL) << "unreached";
break;
}
server->setFailureInjection(std::move(fi));
expected = exp;
auto callback = folly::make_unique<Callback>(&expected);
client.sendResponse(rpcOptions, std::move(callback), 1);
base.loop();
}
}
uint32_t HTTPClientChannel::sendRequest(
RpcOptions& rpcOptions,
std::unique_ptr<RequestCallback> cb,
std::unique_ptr<apache::thrift::ContextStack> ctx,
unique_ptr<IOBuf> buf,
std::shared_ptr<THeader> header) {
// cb is not allowed to be null.
DCHECK(cb);
DestructorGuard dg(this);
cb->context_ = RequestContext::saveContext();
std::chrono::milliseconds timeout(timeout_);
if (rpcOptions.getTimeout() > std::chrono::milliseconds(0)) {
timeout = rpcOptions.getTimeout();
}
auto twcb =
new HTTPTransactionTwowayCallback(std::move(cb),
std::move(ctx),
isSecurityActive(),
protocolId_,
timer_.get(),
std::chrono::milliseconds(timeout_));
if (!httpSession_) {
TTransportException ex(TTransportException::NOT_OPEN,
"HTTPSession is not open");
twcb->messageSendError(
folly::make_exception_wrapper<TTransportException>(std::move(ex)));
delete twcb;
return -1;
}
auto txn = httpSession_->newTransaction(twcb);
if (!txn) {
TTransportException ex(TTransportException::NOT_OPEN,
"Too many active requests on connection");
// Might be able to create another transaction soon
ex.setOptions(TTransportException::CHANNEL_IS_VALID);
twcb->messageSendError(
folly::make_exception_wrapper<TTransportException>(std::move(ex)));
delete twcb;
return -1;
}
auto streamId = txn->getID();
setRequestHeaderOptions(header.get());
addRpcOptionHeaders(header.get(), rpcOptions);
auto msg = buildHTTPMessage(header.get());
txn->sendHeaders(msg);
txn->sendBody(std::move(buf));
txn->sendEOM();
twcb->sendQueued();
return (uint32_t)streamId;
}
TEST(ThriftServer, ClientTimeoutTest) {
TestThriftServerFactory<TestInterface> factory;
auto server = factory.create();
ScopedServerThread sst(server);
folly::EventBase base;
auto getClient = [&base, &sst]() {
std::shared_ptr<TAsyncSocket> socket(
TAsyncSocket::newSocket(&base, *sst.getAddress()));
return std::make_shared<TestServiceAsyncClient>(
HeaderClientChannel::newChannel(socket));
};
int cbCtor = 0;
int cbCall = 0;
auto callback = [&cbCall, &cbCtor](
std::shared_ptr<TestServiceAsyncClient> client, bool& timeout) {
cbCtor++;
return std::unique_ptr<RequestCallback>(new FunctionReplyCallback(
[&cbCall, client, &timeout](ClientReceiveState&& state) {
cbCall++;
if (state.exception()) {
timeout = true;
try {
std::rethrow_exception(state.exception());
} catch (const TTransportException& e) {
EXPECT_EQ(int(TTransportException::TIMED_OUT), int(e.getType()));
}
return;
}
try {
std::string resp;
client->recv_sendResponse(resp, state);
} catch (const TApplicationException& e) {
timeout = true;
EXPECT_EQ(int(TApplicationException::TIMEOUT), int(e.getType()));
EXPECT_TRUE(state.header()->getFlags() &
HEADER_FLAG_SUPPORT_OUT_OF_ORDER);
return;
}
timeout = false;
}));
};
// Set the timeout to be 5 milliseconds, but the call will take 10 ms.
// The server should send a timeout after 5 milliseconds
RpcOptions options;
options.setTimeout(std::chrono::milliseconds(5));
auto client1 = getClient();
bool timeout1;
client1->sendResponse(options, callback(client1, timeout1), 10000);
base.loop();
EXPECT_TRUE(timeout1);
usleep(10000);
// This time we set the timeout to be 100 millseconds. The server
// should not time out
options.setTimeout(std::chrono::milliseconds(100));
client1->sendResponse(options, callback(client1, timeout1), 10000);
base.loop();
EXPECT_FALSE(timeout1);
usleep(10000);
// This time we set server timeout to be 5 millseconds. However, the
// task should start processing within that millisecond, so we should
// not see an exception because the client timeout should be used after
// processing is started
server->setTaskExpireTime(std::chrono::milliseconds(5));
client1->sendResponse(options, callback(client1, timeout1), 10000);
base.loop();
usleep(10000);
// The server timeout stays at 5 ms, but we put the client timeout at
// 5 ms. We should timeout even though the server starts processing within
// 5ms.
options.setTimeout(std::chrono::milliseconds(5));
client1->sendResponse(options, callback(client1, timeout1), 10000);
base.loop();
EXPECT_TRUE(timeout1);
usleep(50000);
// And finally, with the server timeout at 50 ms, we send 2 requests at
// once. Because the first request will take more than 50 ms to finish
// processing (the server only has 1 worker thread), the second request
// won't start processing until after 50ms, and will timeout, despite the
// very high client timeout.
// We don't know which one will timeout (race conditions) so we just check
// the xor
auto client2 = getClient();
bool timeout2;
server->setTaskExpireTime(std::chrono::milliseconds(50));
options.setTimeout(std::chrono::milliseconds(110));
client1->sendResponse(options, callback(client1, timeout1), 100000);
client2->sendResponse(options, callback(client2, timeout2), 100000);
base.loop();
EXPECT_TRUE(timeout1 || timeout2);
EXPECT_FALSE(timeout1 && timeout2);
EXPECT_EQ(cbCall, cbCtor);
}
