Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

rust: add demo Tonic server #4318

Merged
merged 11 commits into from
Nov 12, 2020
Merged
Show file tree
Hide file tree
Changes from 9 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions tensorboard/data/server/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,10 @@ licenses(["notice"])

# Protocol buffer packages (as in `package foo.bar;` directives) that we need
# to compile to Rust bindings.
_proto_packages = ["tensorboard"]
_proto_packages = [
"demo",
"tensorboard",
]

# Generated files with Rust protobuf bindings. These only exist in the build
# graph, not in the source tree. The file name pattern is specified by Prost
Expand Down Expand Up @@ -35,6 +38,7 @@ rust_library(
"//third_party/rust:prost",
"//third_party/rust:rand",
"//third_party/rust:rand_chacha",
"//third_party/rust:tonic",
],
)

Expand Down Expand Up @@ -68,11 +72,17 @@ rust_binary(
name = "server",
srcs = ["main.rs"],
edition = "2018",
deps = [
":rustboard_core",
"//third_party/rust:tokio",
"//third_party/rust:tonic",
],
)

genrule(
name = "gen_protos",
srcs = [
":demo.proto",
"//tensorboard/compat/proto:proto_srcs",
],
outs = _genproto_files,
Expand All @@ -89,7 +99,7 @@ rust_binary(
edition = "2018",
visibility = ["//visibility:private"],
deps = [
"//third_party/rust:prost_build",
"//third_party/rust:tonic_build",
],
)

Expand Down
26 changes: 26 additions & 0 deletions tensorboard/data/server/demo.pb.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

15 changes: 15 additions & 0 deletions tensorboard/data/server/demo.proto
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
syntax = "proto3";

package demo;

service Demo {
rpc Add(AddRequest) returns (AddResponse);
}

message AddRequest {
repeated int32 term = 1;
}

message AddResponse {
int32 sum = 1;
}
11 changes: 9 additions & 2 deletions tensorboard/data/server/gen_protos_tool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,16 @@ fn main() -> std::io::Result<()> {
dir.push("genproto");
dir
};
prost_build::Config::new()
tonic_build::configure()
.out_dir(&out_dir)
.compile_protos(&["tensorboard/compat/proto/event.proto"], &["."])
.format(false) // don't run `rustfmt`; shouldn't be needed to build
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah :) I knew something was weird about the generated file

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah—I had format(true) (the default) originally, but that only worked
on my machine because I have rustfmt installed; it failed on Travis.
It would be nice for it to be formatted. We could format it pretty
easily if bazelbuild/rules_rust#388 were resolved, or we could copy what
bazelbuild/rules_rust#479 does and do it the hard way. I figured that to
start we could just accept the ugly generated code: we have the Rustdoc
server for normal API usage questions, and you can always run rustfmt
on it locally if you want to look at the implementation.

.compile(
&[
"tensorboard/compat/proto/event.proto",
"tensorboard/data/server/demo.proto",
],
&["."],
)
.expect("compile_protos");
Ok(())
}
4 changes: 4 additions & 0 deletions tensorboard/data/server/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,10 @@ mod scripted_reader;
/// Protocol buffer bindings.
#[allow(clippy::all)]
pub mod proto {
/// Bindings for `package demo`, used for a demo Tonic server.
pub mod demo {
include!("demo.pb.rs");
}
/// Bindings for `package tensorboard`, containing standard TensorFlow protos.
pub mod tensorboard {
include!("tensorboard.pb.rs");
Expand Down
36 changes: 35 additions & 1 deletion tensorboard/data/server/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,4 +13,38 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

fn main() {}
#![allow(clippy::needless_update)] // https://github.com/rust-lang/rust-clippy/issues/6323

use tonic::{transport::Server, Request, Response};

use rustboard_core::proto::demo as pb;

#[derive(Debug, Default)]
struct DemoHandler;

#[tonic::async_trait]
impl pb::demo_server::Demo for DemoHandler {
async fn add(
&self,
request: Request<pb::AddRequest>,
) -> Result<Response<pb::AddResponse>, tonic::Status> {
let request = request.into_inner();
let sum: i32 = request.term.into_iter().sum();
let response = pb::AddResponse {
sum,
..Default::default()
};
Ok(Response::new(response))
}
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let addr = "[::0]:6789".parse::<std::net::SocketAddr>()?;
let handler = DemoHandler::default();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One part that is iffy in my head is, how does it know to invoke async fn add for the Demo.Add service method?

Copy link
Contributor Author

@wchargin wchargin Nov 12, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At a high level, the answer is: “the Demo trait, which we implemented
on our handler”. That trait contains one Rust method for each gRPC
method. The trait is defined in the code generated by tonic_build;
that generated code has the hard-coded mapping between gRPC method names
and Rust methods and types.

I’ve never read the generated code before, but we can take a look. It’s
got some stuff going on to handle gRPC machinery and async internals,
but we shouldn’t usually need to care about those. Here it is, lightly
annotated (from tensorboard.demo.pb.rs after rustfmt):

    impl<T, B> Service<http::Request<B>> for DemoServer<T>
    where
////////// Trait implemented for all handlers that implement the `Demo` trait,
////////// which includes our `DemoHandler` struct:
        T: Demo,
////////// The request and error types need to be moveable between threads,
////////// shareable across threads, and with only static borrowed data; this
////////// way, the server has more opportunities to use threads for handling:
        B: HttpBody + Send + Sync + 'static,
        B::Error: Into<StdError> + Send + 'static,
    {
        // ...
        fn call(&mut self, req: http::Request<B>) -> Self::Future {
            let inner = self.inner.clone();
            match req.uri().path() {
////////////////// Here's where we match against the gRPC request path...
                "/demo.Demo/Add" => {
                    #[allow(non_camel_case_types)]
                    struct AddSvc<T: Demo>(pub Arc<T>);
                    impl<T: Demo> tonic::server::UnaryService<super::AddRequest> for AddSvc<T> {
                        type Response = super::AddResponse;
                        type Future = BoxFuture<tonic::Response<Self::Response>, tonic::Status>;
                        fn call(
                            &mut self,
                            request: tonic::Request<super::AddRequest>,
                        ) -> Self::Future {
                            let inner = self.0.clone();
////////////////////////////// ...and here's where we call our `async fn add`:
                            let fut = async move { (*inner).add(request).await };
                            Box::pin(fut)
                        }
                    }
                    let inner = self.inner.clone();
                    let fut = async move {
                        let interceptor = inner.1.clone();
                        let inner = inner.0;
                        let method = AddSvc(inner);
                        let codec = tonic::codec::ProstCodec::default();
                        let mut grpc = if let Some(interceptor) = interceptor {
                            tonic::server::Grpc::with_interceptor(codec, interceptor)
                        } else {
                            tonic::server::Grpc::new(codec)
                        };
                        let res = grpc.unary(method, req).await;
                        Ok(res)
                    };
                    Box::pin(fut)
                }
////////////////// Here's the catch-all, with `Code.UNIMPLEMENTED == 12`:
                _ => Box::pin(async move {
                    Ok(http::Response::builder()
                        .status(200)
                        .header("grpc-status", "12")
                        .body(tonic::body::BoxBody::empty())
                        .unwrap())
                }),
            }
        }
    }

Server::builder()
.add_service(pb::demo_server::DemoServer::new(handler))
.serve(addr)
.await?;
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You didn’t comment on the postfix await syntax, so I will! :-)

When I saw that Rust was considering implementing await as a postfix
operator, like foo.await rather than await foo, I was pretty
skeptical, since it deviates from JS, Python, C#, etc. But then I came
across this example. Compare these two: which is easier to read?

// JavaScript: prefix `await` syntax
const data = (await (await fetch("data/runs")).json()).field;
// Rust: postfix `await` syntax
let data = fetch("data/runs").await?.json().await?.field;

Once I saw this, I was immediately sold. With prefix await, control
flow bounces back and forth:

const data = (await (await fetch("data/runs")).json()).field;  /*
                           ------------------ 1
                     ----- 2
                                               ------ 3
              ----- 4
                                                       ----- 5 */

and you have to deal with the associated parentheses: the first thing
that you see is await (await (...)). Getting the nesting right is
particularly annoying when manually sending fetches in the console.
With postfix await, it’s nice and linear, like a fluent method chain:

let data = fetch("data/runs").await?.json().await?.field;  /*
           --------------- 1  --- 2  --- 3  --- 4  --- 5   */

Ok(())
}