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); }