From c81811bc58813cc99c3d274e42559bad32f4c419 Mon Sep 17 00:00:00 2001 From: frumious Date: Tue, 5 Mar 2019 22:36:03 -0800 Subject: [PATCH] add server instrumentation for method calls --- Cargo.toml | 8 +++ compiler/src/codegen.rs | 54 +++++++++++---- examples/hello_world_instrument/client.rs | 38 ++++++++++ examples/hello_world_instrument/server.rs | 84 +++++++++++++++++++++++ src/lib.rs | 4 +- src/server.rs | 5 ++ 6 files changed, 180 insertions(+), 13 deletions(-) create mode 100644 examples/hello_world_instrument/client.rs create mode 100644 examples/hello_world_instrument/server.rs diff --git a/Cargo.toml b/Cargo.toml index d7f674e0a..646447a5d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,6 +48,14 @@ path = "examples/hello_world/client.rs" name = "greeter_server" path = "examples/hello_world/server.rs" +[[example]] +name = "greeter_client_instrument" +path = "examples/hello_world_instrument/client.rs" + +[[example]] +name = "greeter_server_instrument" +path = "examples/hello_world_instrument/server.rs" + [dev-dependencies] serde_json = "1.0" serde = "1.0" diff --git a/compiler/src/codegen.rs b/compiler/src/codegen.rs index bfd36695c..9b8b07a11 100644 --- a/compiler/src/codegen.rs +++ b/compiler/src/codegen.rs @@ -379,24 +379,35 @@ impl<'a> MethodGen<'a> { w.fn_def(&sig); } - fn write_bind(&self, w: &mut CodeWriter) { + fn write_bind(&self, w: &mut CodeWriter, instrument: bool) { let add = match self.method_type().0 { MethodType::Unary => "add_unary_handler", MethodType::ClientStreaming => "add_client_streaming_handler", MethodType::ServerStreaming => "add_server_streaming_handler", MethodType::Duplex => "add_duplex_streaming_handler", }; - w.block( - &format!( - "builder = builder.{}(&{}, move |ctx, req, resp| {{", - add, - self.const_method_name() - ), - "});", - |w| { - w.write_line(&format!("instance.{}(ctx, req, resp)", self.name())); - }, + + let line = &format!( + "builder = builder.{}(&{}, move |ctx, req, resp| {{", + add, + self.const_method_name() ); + + if instrument { + w.block(line, "});", |w| { + w.write_line("&p.before();"); + w.write_line(&format!( + "let res = instance.{}(ctx, req, resp);", + self.name() + )); + w.write_line("&p.after();"); + w.write_line("res"); + }); + } else { + w.block(line, "});", |w| { + w.write_line(&format!("instance.{}(ctx, req, resp)", self.name())); + }); + } } } @@ -489,7 +500,26 @@ impl<'a> ServiceGen<'a> { w.write_line("let mut builder = ::grpcio::ServiceBuilder::new();"); for method in &self.methods { w.write_line("let mut instance = s.clone();"); - method.write_bind(w); + method.write_bind(w, false); + } + w.write_line("builder.build()"); + }); + + // server with instrumentation + w.write_line(""); + + let s = format!( + "create_instrumented_{}(s: S, p: P) -> {}", + to_snake_case(&self.service_name()), + self.service_name(), + fq_grpc("ServerInstrumenter"), + fq_grpc("Service") + ); + w.pub_fn(&s, |w| { + w.write_line("let mut builder = ::grpcio::ServiceBuilder::new();"); + for method in &self.methods { + w.write_line("let mut instance = s.clone();"); + method.write_bind(w, true); } w.write_line("builder.build()"); }); diff --git a/examples/hello_world_instrument/client.rs b/examples/hello_world_instrument/client.rs new file mode 100644 index 000000000..d860602e6 --- /dev/null +++ b/examples/hello_world_instrument/client.rs @@ -0,0 +1,38 @@ +// Copyright 2017 PingCAP, Inc. +// +// Licensed 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, +// See the License for the specific language governing permissions and +// limitations under the License. + +extern crate grpcio; +extern crate grpcio_proto; +#[macro_use] +extern crate log; + +#[path = "../log_util.rs"] +mod log_util; + +use std::sync::Arc; + +use grpcio::{ChannelBuilder, EnvBuilder}; +use grpcio_proto::example::helloworld::HelloRequest; +use grpcio_proto::example::helloworld_grpc::GreeterClient; + +fn main() { + let _guard = log_util::init_log(None); + let env = Arc::new(EnvBuilder::new().build()); + let ch = ChannelBuilder::new(env).connect("localhost:50051"); + let client = GreeterClient::new(ch); + + let mut req = HelloRequest::new(); + req.set_name("world".to_owned()); + let reply = client.say_hello(&req).expect("rpc"); + info!("Greeter received: {}", reply.get_message()); +} diff --git a/examples/hello_world_instrument/server.rs b/examples/hello_world_instrument/server.rs new file mode 100644 index 000000000..6f50b202d --- /dev/null +++ b/examples/hello_world_instrument/server.rs @@ -0,0 +1,84 @@ +// Copyright 2017 PingCAP, Inc. +// +// Licensed 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, +// See the License for the specific language governing permissions and +// limitations under the License. + +extern crate futures; +extern crate grpcio; +extern crate grpcio_proto; +#[macro_use] +extern crate log; + +#[path = "../log_util.rs"] +mod log_util; + +use std::io::Read; +use std::sync::Arc; +use std::{io, thread}; + +use futures::sync::oneshot; +use futures::Future; +use grpcio::{Environment, RpcContext, ServerBuilder, ServerInstrumenter, UnarySink}; + +use grpcio_proto::example::helloworld::{HelloReply, HelloRequest}; +use grpcio_proto::example::helloworld_grpc::{self, Greeter}; + +#[derive(Clone)] +struct GreeterService; + +impl Greeter for GreeterService { + fn say_hello(&mut self, ctx: RpcContext, req: HelloRequest, sink: UnarySink) { + let msg = format!("Hello {}", req.get_name()); + let mut resp = HelloReply::new(); + resp.set_message(msg); + let f = sink + .success(resp) + .map_err(move |e| error!("failed to reply {:?}: {:?}", req, e)); + ctx.spawn(f) + } +} + +#[derive(Copy, Clone)] +struct MyServerInstrumenter; + +impl ServerInstrumenter for MyServerInstrumenter { + fn before(&self) { + println!("Hello world!"); + } + + fn after(&self) { + println!("Goodbye."); + } +} + +fn main() { + let _guard = log_util::init_log(None); + let env = Arc::new(Environment::new(1)); + let service = + helloworld_grpc::create_instrumented_greeter(GreeterService, MyServerInstrumenter); + let mut server = ServerBuilder::new(env) + .register_service(service) + .bind("127.0.0.1", 50_051) + .build() + .unwrap(); + server.start(); + for &(ref host, port) in server.bind_addrs() { + info!("listening on {}:{}", host, port); + } + let (tx, rx) = oneshot::channel(); + thread::spawn(move || { + info!("Press ENTER to exit..."); + let _ = io::stdin().read(&mut [0]).unwrap(); + tx.send(()) + }); + let _ = rx.wait(); + let _ = server.shutdown().wait(); +} diff --git a/src/lib.rs b/src/lib.rs index ee39207d0..8ef79f3d6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -83,4 +83,6 @@ pub use crate::env::{EnvBuilder, Environment}; pub use crate::error::{Error, Result}; pub use crate::log_util::redirect_log; pub use crate::metadata::{Metadata, MetadataBuilder, MetadataIter}; -pub use crate::server::{Server, ServerBuilder, Service, ServiceBuilder, ShutdownFuture}; +pub use crate::server::{ + Server, ServerBuilder, ServerInstrumenter, Service, ServiceBuilder, ShutdownFuture, +}; diff --git a/src/server.rs b/src/server.rs index 4e8904fd8..e15ecadae 100644 --- a/src/server.rs +++ b/src/server.rs @@ -47,6 +47,11 @@ impl Handler { } } +pub trait ServerInstrumenter { + fn before(&self); + fn after(&self); +} + pub trait CloneableHandler: Send { fn handle(&mut self, ctx: RpcContext, reqs: Option); fn box_clone(&self) -> Box;