177 lines
7.3 KiB
C++
177 lines
7.3 KiB
C++
// 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.
|
|
|
|
|
|
#include <gflags/gflags.h>
|
|
#include <butil/logging.h>
|
|
#include <brpc/server.h>
|
|
#include <brpc/channel.h>
|
|
#include "view.pb.h"
|
|
|
|
DEFINE_int32(port, 8888, "TCP Port of this server");
|
|
DEFINE_string(target, "", "The server to view");
|
|
DEFINE_int32(timeout_ms, 5000, "Timeout for calling the server to view");
|
|
|
|
// handle HTTP response of accessing builtin services of the target server.
|
|
static void handle_response(brpc::Controller* client_cntl,
|
|
std::string target,
|
|
brpc::Controller* server_cntl,
|
|
google::protobuf::Closure* server_done) {
|
|
// Copy all headers. The "Content-Length" will be overwriteen.
|
|
server_cntl->http_response() = client_cntl->http_response();
|
|
// Copy content.
|
|
server_cntl->response_attachment() = client_cntl->response_attachment();
|
|
// Insert "rpc_view: <target>" before </body> so that users are always
|
|
// visually notified with target server w/o confusions.
|
|
butil::IOBuf& content = server_cntl->response_attachment();
|
|
butil::IOBuf before_body;
|
|
if (content.cut_until(&before_body, "</body>") == 0) {
|
|
before_body.append(
|
|
"<style type=\"text/css\">\n"
|
|
".rpcviewlogo {position: fixed; bottom: 0px; right: 0px;"
|
|
" color: #ffffff; background-color: #000000; }\n"
|
|
" </style>\n"
|
|
"<span class='rpcviewlogo'> rpc_view: ");
|
|
before_body.append(target);
|
|
before_body.append(" </span></body>");
|
|
before_body.append(content);
|
|
content = before_body;
|
|
}
|
|
// Notice that we don't set RPC to failed on http errors because we
|
|
// want to pass unchanged content to the users otherwise RPC replaces
|
|
// the content with ErrorText.
|
|
if (client_cntl->Failed() &&
|
|
client_cntl->ErrorCode() != brpc::EHTTP) {
|
|
server_cntl->SetFailed(client_cntl->ErrorCode(),
|
|
"%s", client_cntl->ErrorText().c_str());
|
|
}
|
|
delete client_cntl;
|
|
server_done->Run();
|
|
}
|
|
|
|
// A http_master_service.
|
|
class ViewServiceImpl : public ViewService {
|
|
public:
|
|
ViewServiceImpl() {}
|
|
virtual ~ViewServiceImpl() {}
|
|
virtual void default_method(google::protobuf::RpcController* cntl_base,
|
|
const HttpRequest*,
|
|
HttpResponse*,
|
|
google::protobuf::Closure* done) {
|
|
brpc::ClosureGuard done_guard(done);
|
|
brpc::Controller* server_cntl =
|
|
static_cast<brpc::Controller*>(cntl_base);
|
|
|
|
// Get or set target. Notice that we don't access FLAGS_target directly
|
|
// which is thread-unsafe (for string flags).
|
|
std::string target;
|
|
const std::string* newtarget =
|
|
server_cntl->http_request().uri().GetQuery("changetarget");
|
|
if (newtarget) {
|
|
if (GFLAGS_NS::SetCommandLineOption("target", newtarget->c_str()).empty()) {
|
|
server_cntl->SetFailed("Fail to change value of -target");
|
|
return;
|
|
}
|
|
target = *newtarget;
|
|
} else {
|
|
if (!GFLAGS_NS::GetCommandLineOption("target", &target)) {
|
|
server_cntl->SetFailed("Fail to get value of -target");
|
|
return;
|
|
}
|
|
}
|
|
|
|
// Create the http channel on-the-fly. Notice that we've set
|
|
// `defer_close_second' in main() so that dtor of channels do not
|
|
// close connections immediately and ad-hoc creation of channels
|
|
// often reuses the not-yet-closed connections.
|
|
brpc::Channel http_chan;
|
|
brpc::ChannelOptions http_chan_opt;
|
|
http_chan_opt.protocol = brpc::PROTOCOL_HTTP;
|
|
if (http_chan.Init(target.c_str(), &http_chan_opt) != 0) {
|
|
server_cntl->SetFailed(brpc::EINTERNAL,
|
|
"Fail to connect to %s", target.c_str());
|
|
return;
|
|
}
|
|
|
|
// Remove "Accept-Encoding". We need to insert additional texts
|
|
// before </body>, preventing the server from compressing the content
|
|
// simplifies our code. The additional bandwidth consumption shouldn't
|
|
// be an issue for infrequent checking out of builtin services pages.
|
|
server_cntl->http_request().RemoveHeader("Accept-Encoding");
|
|
|
|
brpc::Controller* client_cntl = new brpc::Controller;
|
|
client_cntl->http_request() = server_cntl->http_request();
|
|
// Remove "Host" so that RPC will laterly serialize the (correct)
|
|
// target server in.
|
|
client_cntl->http_request().RemoveHeader("host");
|
|
|
|
// Setup the URI.
|
|
const brpc::URI& server_uri = server_cntl->http_request().uri();
|
|
std::string uri = server_uri.path();
|
|
if (!server_uri.query().empty()) {
|
|
uri.push_back('?');
|
|
uri.append(server_uri.query());
|
|
}
|
|
if (!server_uri.fragment().empty()) {
|
|
uri.push_back('#');
|
|
uri.append(server_uri.fragment());
|
|
}
|
|
client_cntl->http_request().uri() = uri;
|
|
|
|
// /hotspots pages may take a long time to finish, since they all have
|
|
// query "seconds", we set the timeout to be longer than "seconds".
|
|
const std::string* seconds =
|
|
server_cntl->http_request().uri().GetQuery("seconds");
|
|
int64_t timeout_ms = FLAGS_timeout_ms;
|
|
if (seconds) {
|
|
timeout_ms += atoll(seconds->c_str()) * 1000;
|
|
}
|
|
client_cntl->set_timeout_ms(timeout_ms);
|
|
|
|
// Keep content as it is.
|
|
client_cntl->request_attachment() = server_cntl->request_attachment();
|
|
|
|
http_chan.CallMethod(NULL, client_cntl, NULL, NULL,
|
|
brpc::NewCallback(
|
|
handle_response, client_cntl, target,
|
|
server_cntl, done_guard.release()));
|
|
}
|
|
};
|
|
|
|
int main(int argc, char* argv[]) {
|
|
GFLAGS_NS::ParseCommandLineFlags(&argc, &argv, true);
|
|
if (FLAGS_target.empty() &&
|
|
(argc != 2 ||
|
|
GFLAGS_NS::SetCommandLineOption("target", argv[1]).empty())) {
|
|
LOG(ERROR) << "Usage: ./rpc_view <ip>:<port>";
|
|
return -1;
|
|
}
|
|
// This keeps ad-hoc creation of channels reuse previous connections.
|
|
GFLAGS_NS::SetCommandLineOption("defer_close_second", "10");
|
|
|
|
brpc::Server server;
|
|
server.set_version("rpc_view_server");
|
|
brpc::ServerOptions server_opt;
|
|
server_opt.http_master_service = new ViewServiceImpl;
|
|
if (server.Start(FLAGS_port, &server_opt) != 0) {
|
|
LOG(ERROR) << "Fail to start ViewServer";
|
|
return -1;
|
|
}
|
|
server.RunUntilAskedToQuit();
|
|
return 0;
|
|
}
|