Use Router in Trusted Database (#1748)

This change updates Trusted Database example to use Router pattern.

Ref #1066

docs/programming-oak.md

@@ -403,17 +403,20 @@ from the implicit incoming channel, usually by implementing the
 trait setting the
 [`Command`](https://project-oak.github.io/oak/sdk/doc/oak/trait.CommandHandler.html#associatedtype.Command)
 associated type to
-[`ConfigMap`](https://project-oak.github.io/oak/sdk/doc/oak/proto/oak/application/struct.ConfigMap.html),
-or manually reading from the initial `Receiver`:
+[`ConfigMap`](https://project-oak.github.io/oak/sdk/doc/oak/proto/oak/application/struct.ConfigMap.html):
 
 <!-- prettier-ignore-start -->
-[embedmd]:# (../examples/trusted_database/module/rust/src/lib.rs Rust /oak::entrypoint/ /.*let config_map =.*/)
+[embedmd]:# (../examples/aggregator/module/rust/src/lib.rs Rust /impl oak::CommandHandler/ /.*context.*/)
 ```Rust
-oak::entrypoint!(oak_main<ConfigMap> => |receiver: Receiver<ConfigMap>| {
-    let log_sender = oak::logger::create().unwrap();
-    oak::logger::init(log_sender, log::Level::Debug).unwrap();
+impl oak::CommandHandler for Main {
+    type Command = ConfigMap;
 
-    let config_map = receiver.receive().expect("Couldn't read config map");
+    fn handle_command(&mut self, command: ConfigMap) -> anyhow::Result<()> {
+        let log_sender = oak::logger::create()?;
+        oak::logger::init(log_sender.clone(), log::Level::Debug)?;
+        let config: Config =
+            toml::from_slice(&command.items.get("config").expect("Couldn't find config"))
+                .context("Couldn't parse TOML config file")?;
 ```
 <!-- prettier-ignore-end -->
 

examples/trusted_database/module/rust/Cargo.toml

@@ -14,6 +14,7 @@ log = "*"
 oak = "=0.1.0"
 oak_abi = "=0.1.0"
 oak_io = "=0.1.0"
+oak_services = "=0.1.0"
 prost = "*"
 serde = "*"
 quick-xml = { version = "*", features = ["serialize"] }

examples/trusted_database/module/rust/build.rs

@@ -18,7 +18,7 @@ fn main() {
     oak_utils::compile_protos(
         &[
             "../../proto/trusted_database.proto",
-            "../../proto/trusted_database_command.proto",
+            "../../proto/trusted_database_init.proto",
         ],
         &["../../proto", "../../../../"],
     );

examples/trusted_database/module/rust/src/handler.rs

@@ -14,39 +14,37 @@
 // limitations under the License.
 //
 
-//! Trusted Database Handler Node.
-//!
-//! In the current implementation clients send their location coordinates (latitude and longitude)
-//! and the Handler Node returns the location of the closest Point Of Interest.
-//!
-//! The Handler Node searches for the closest Point Of Interest in the database received from the
-//! Main Node.
-
 use crate::proto::oak::examples::trusted_database::{
     GetPointOfInterestRequest, GetPointOfInterestResponse, ListPointsOfInterestRequest,
     ListPointsOfInterestResponse, Location, PointOfInterestMap, TrustedDatabase,
-    TrustedDatabaseCommand, TrustedDatabaseDispatcher,
+    TrustedDatabaseDispatcher, TrustedDatabaseInit,
 };
 use log::{debug, error, warn};
-use oak::{
-    grpc,
-    io::{Receiver, ReceiverExt},
-};
+use oak::grpc;
 
 // Error messages.
 const NO_LOCATION_ERROR: &str = "Location is not specified";
 const ID_NOT_FOUND_ERROR: &str = "ID not found";
 const EMPTY_DATABASE_ERROR: &str = "Empty database";
 
-/// Oak Node that contains a copy of the database.
-pub struct TrustedDatabaseHandlerNode {
+/// Oak Handler Node that contains a copy of the database and handles client requests.
+#[derive(Default)]
+pub struct Handler {
     points_of_interest: PointOfInterestMap,
 }
 
-oak::impl_dispatcher!(impl TrustedDatabaseHandlerNode : TrustedDatabaseDispatcher);
+impl oak::WithInit for Handler {
+    type Init = TrustedDatabaseInit;
+
+    fn create(init: Self::Init) -> Self {
+        oak::logger::init(init.log_sender.unwrap(), log::Level::Debug).unwrap();
+        let points_of_interest = init.points_of_interest.expect("Couldn't receive database");
+        Self { points_of_interest }
+    }
+}
 
 /// A gRPC service implementation for the Private Information Retrieval example.
-impl TrustedDatabase for TrustedDatabaseHandlerNode {
+impl TrustedDatabase for Handler {
     // Find Point Of Interest based on id.
     fn get_point_of_interest(
         &mut self,
@@ -73,7 +71,7 @@ impl TrustedDatabase for TrustedDatabaseHandlerNode {
         &mut self,
         request: ListPointsOfInterestRequest,
     ) -> grpc::Result<ListPointsOfInterestResponse> {
-        debug!("Received request: {:?}", request);
+        error!("Received request: {:?}", request);
         let request_location = request.location.ok_or_else(|| {
             let err = grpc::build_status(grpc::Code::InvalidArgument, &NO_LOCATION_ERROR);
             warn!("{:?}", err);
@@ -138,18 +136,6 @@ pub fn distance(first: Location, second: Location) -> f32 {
     EARTH_RADIUS * central_angle
 }
 
-oak::entrypoint!(handler_oak_main<TrustedDatabaseCommand> => |command_receiver: Receiver<TrustedDatabaseCommand>| {
-    let log_sender = oak::logger::create().unwrap();
-    oak::logger::init(log_sender, log::Level::Debug).unwrap();
-
-    // Receive command.
-    let command: TrustedDatabaseCommand =
-        command_receiver.receive().expect("Couldn't receive command");
-    let receiver = command.invocation_receiver.expect("Couldn't receive gRPC invocation receiver");
+oak::entrypoint_command_handler_init!(handler => Handler);
 
-    // Run event loop and handle incoming invocations.
-    let node = TrustedDatabaseHandlerNode { points_of_interest: command.points_of_interest.expect("No database entries") };
-    let invocation_receiver = receiver.receiver.expect("Empty gRPC invocation receiver");
-    // The event loop only runs once because the `Main` Node sends a single invocation.
-    oak::run_command_loop(node, invocation_receiver.iter());
-});
+oak::impl_dispatcher!(impl Handler : TrustedDatabaseDispatcher);

examples/trusted_database/module/rust/src/lib.rs

@@ -35,6 +35,7 @@
 pub mod proto {
     pub mod oak {
         pub use oak::proto::oak::invocation;
+        pub use oak_services::proto::oak::log;
         pub mod examples {
             pub mod trusted_database {
                 include!(concat!(
@@ -43,7 +44,7 @@ pub mod proto {
                 ));
                 include!(concat!(
                     env!("OUT_DIR"),
-                    "/oak.examples.trusted_database_command.rs"
+                    "/oak.examples.trusted_database_init.rs"
                 ));
             }
         }
@@ -52,82 +53,45 @@ pub mod proto {
 
 mod database;
 mod handler;
+mod router;
 #[cfg(test)]
 mod tests;
 
+use crate::{proto::oak::examples::trusted_database::TrustedDatabaseInit, router::Router};
 use anyhow::Context;
 use database::load_database;
-use log::debug;
-use oak::{
-    grpc,
-    io::{Receiver, ReceiverExt, SenderExt},
-    proto::oak::invocation::GrpcInvocationReceiver,
-    CommandHandler,
-};
-use oak_abi::{label::Label, proto::oak::application::ConfigMap};
-use proto::oak::examples::trusted_database::{PointOfInterestMap, TrustedDatabaseCommand};
+use oak::proto::oak::application::ConfigMap;
+use oak_abi::label::Label;
 
-/// Oak Node that contains an in-memory database.
-pub struct TrustedDatabaseNode {
-    points_of_interest: PointOfInterestMap,
-}
+/// Main entrypoint of the Trusted Database application.
+///
+/// This node is in charge of creating the other top-level nodes, but does not process any request.
+#[derive(Default)]
+struct Main;
 
-impl CommandHandler for TrustedDatabaseNode {
-    type Command = grpc::Invocation;
+impl oak::CommandHandler for Main {
+    type Command = ConfigMap;
 
-    fn handle_command(&mut self, invocation: grpc::Invocation) -> anyhow::Result<()> {
-        // Create a client request handler Node.
-        debug!("Creating handler Node");
-        // TODO(#1406): Use client assigned label for creating a new handler Node.
-        let sender = oak::io::node_create(
-            "handler",
-            &Label::public_untrusted(),
-            &oak::node_config::wasm("app", "handler_oak_main"),
-        )
-        .context("Couldn't create handler Node")?;
+    fn handle_command(&mut self, config_map: Self::Command) -> anyhow::Result<()> {
+        let log_sender = oak::logger::create()?;
+        oak::logger::init(log_sender.clone(), log::Level::Debug)?;
+        let points_of_interest = load_database(config_map).expect("Couldn't load database");
 
-        // Create a gRPC invocation channel for forwarding requests to the
-        // `TrustedDatabaseHandlerNode`.
-        let (invocation_sender, invocation_receiver) = oak::io::channel_create::<grpc::Invocation>(
-            "gRPC invocation",
+        let init = TrustedDatabaseInit {
+            log_sender: Some(log_sender),
+            points_of_interest: Some(points_of_interest),
+        };
+        let router_sender = oak::io::entrypoint_node_create::<Router, _, _>(
+            "router",
             &Label::public_untrusted(),
+            "app",
+            init,
         )
-        .context("Couldn't create gRPC invocation channel")?;
-
-        // Create a command message that contains a copy of the database.
-        let command = TrustedDatabaseCommand {
-            invocation_receiver: Some(GrpcInvocationReceiver {
-                receiver: Some(invocation_receiver),
-            }),
-            points_of_interest: Some(self.points_of_interest.clone()),
-        };
-
-        // Send the command massage to create a `TrustedDatabaseHandlerNode`
-        debug!("Sending command message to handler Node");
-        sender
-            .send(&command)
-            .context("Couldn't send command to handler Node")?;
-        oak::channel_close(sender.handle.handle).context("Couldn't close sender channel")?;
-
-        // Send the original gRPC invocation to the `TrustedDatabaseHandlerNode`
-        debug!("Sending gRPC invocation to handler Node");
-        invocation_sender
-            .send(&invocation)
-            .context("Couldn't send gRPC invocation to handler Node")?;
-        oak::channel_close(invocation_sender.handle.handle)
-            .context("Couldn't close sender channel")?;
-
+        .context("Couldn't create router node")?;
+        oak::grpc::server::init_with_sender("[::]:8080", router_sender)
+            .context("Couldn't create gRPC server pseudo-Node")?;
         Ok(())
     }
 }
 
-oak::entrypoint!(oak_main<ConfigMap> => |receiver: Receiver<ConfigMap>| {
-    let log_sender = oak::logger::create().unwrap();
-    oak::logger::init(log_sender, log::Level::Debug).unwrap();
-
-    let config_map = receiver.receive().expect("Couldn't read config map");
-    let points_of_interest = load_database(config_map).expect("Couldn't load database");
-    let grpc_channel =
-        oak::grpc::server::init("[::]:8080").expect("Couldn't create gRPC server pseudo-Node");
-    oak::run_command_loop(TrustedDatabaseNode { points_of_interest }, grpc_channel.iter());
-});
+oak::entrypoint_command_handler!(oak_main => Main);

examples/trusted_database/module/rust/src/router.rs

@@ -0,0 +1,65 @@
+//
+// Copyright 2020 The Project Oak Authors
+//
+// 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,
+// 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.
+//
+
+use crate::{handler::Handler, proto::oak::examples::trusted_database::TrustedDatabaseInit};
+use anyhow::Context;
+use oak::{
+    grpc,
+    io::{ReceiverExt, SenderExt},
+    CommandHandler,
+};
+
+/// Oak Router Node that contains an in-memory database (saved in the `init`).
+#[derive(Default)]
+pub struct Router {
+    /// Init message to be sent to every newly created Handler Node.
+    init: TrustedDatabaseInit,
+}
+
+impl oak::WithInit for Router {
+    type Init = TrustedDatabaseInit;
+
+    fn create(init: Self::Init) -> Self {
+        Self { init }
+    }
+}
+
+/// Processes client requests and creates individual Handler Nodes.
+/// Each newly created Handler Node receives a copy of the database stored in [`Router::init`].
+impl CommandHandler for Router {
+    type Command = grpc::Invocation;
+
+    fn handle_command(&mut self, invocation: Self::Command) -> anyhow::Result<()> {
+        let label = invocation
+            .receiver
+            .as_ref()
+            .context("Couldn't get receiver")?
+            .label()
+            .context("Couldn't get label")?;
+        let handler_invocation_sender = oak::io::entrypoint_node_create::<Handler, _, _>(
+            "handler",
+            &label,
+            "app",
+            self.init.clone(),
+        )
+        .context("Couldn't create handler node")?;
+        handler_invocation_sender
+            .send(&invocation)
+            .context("Couldn't send invocation to handler node")
+    }
+}
+
+oak::entrypoint_command_handler_init!(router => Router);

examples/trusted_database/proto/trusted_database_command.proto → examples/trusted_database/proto/trusted_database_init.proto

@@ -16,12 +16,16 @@
 
 syntax = "proto3";
 
-package oak.examples.trusted_database_command;
+package oak.examples.trusted_database_init;
 
 import "trusted_database.proto";
 import "oak_services/proto/grpc_invocation.proto";
+import "oak_services/proto/log.proto";
+import "proto/handle.proto";
 
-message TrustedDatabaseCommand {
-  oak.invocation.GrpcInvocationReceiver invocation_receiver = 1;
+// Initialization message that should be sent to Router Oak Node and Handler Oak Node.
+message TrustedDatabaseInit {
+  oak.handle.Sender log_sender = 1 [(oak.handle.message_type) = ".oak.log.LogMessage"];
+  // Copy of the database.
   oak.examples.trusted_database.PointOfInterestMap points_of_interest = 2;
 }