1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
//! # A Library for Creating Norgopolis Modules
//! 
//! For information about Norgopolis, consult https://github.com/nvim-neorg/norgopolis.
//! 
//! This library exposes an API for creating and maintaining a connection to the Norgopolis router.
//! Norgopolis modules provide specific sets of functionality, for example multithreaded parsing, database
//! access, etc. All of the default modules created by the Neorg team are built on top of this library.
//!
//! # Usage
//! 
//! ### General Setup
//! 
//! First, create a struct for your module. Name it whatever you'd like:
//! 
//! ```rs
//! use norgopolis_module::{
//!     invoker_service::Service, module_communication::MessagePack, Code, Module, Status,
//! };
//! 
//! #[derive(Default)]
//! struct MyModule {
//!     // add any data or state you might need to maintain here...
//! }
//! ```
//! 
//! Second, implement the `norgopolis_module::invoker_service::Service` trait for your struct.
//! This forces you to implement a `call` function which will be invoked any time someone routes
//! a message to your module. Since async traits are not stabilized within Rust yet, tag your
//! trait implementation with `#[norgopolis_module::async_trait]`:
//! 
//! ```rs
//! use tokio_stream::wrappers::UnboundedReceiverStream;
//! 
//! #[norgopolis_module::async_trait]
//! impl Service for MyModule {
//!     type Stream = UnboundedReceiverStream<Result<MessagePack, Status>>;
//! 
//!     async fn call(
//!         &self,
//!         function: String,
//!         args: Option<MessagePack>,
//!     ) -> Result<Self::Stream, Status> {
//!         todo!()
//!     }
//! }
//! ```
//! 
//! ##### `Stream`
//! 
//! The `Stream` type defines what sort of data will be returned back via gRPC. We recommend
//! that you set it to `UnboundedReceiverStream<Result<MessagePack, Status>>`. This means that
//! given one request your module will be able to return an infinite amount of MessagePack responses,
//! or a status code in case something went wrong.
//! 
//! ##### `call`
//! 
//! The `call` function gets invoked whenever a client routes a message to you. The message contains:
//! - The function that they would like to invoke
//! - An optional set of parameters they would like to supply to the function.
//! 
//! ### Creating the Basic Glue
//! 
//! In the `call` function it's recommended to match over all possible function names that your module
//! supports and returning an error code if it's unsupported:
//! 
//! ```rs
//! match function.as_str() {
//!     "my-function" => todo!(),
//!     _ => Err(Status::new(Code::NotFound, "Requested function not found!")),
//! }
//! ```
//! 
//! > [!IMPORTANT]
//! > It's always better to return *some* sort of status code over panicking.
//! > Panicking will terminate the connection to Norgopolis and the user will not receive
//! > any sort of error or warning.
//! 
//! ### Decoding the Parameters
//! 
//! If your function takes in any amount of parameters then now is the time to decode them.
//! If your parameter is complex (e.g. a dictionary) then it's recommended to create a struct
//! designated for it. Be sure to derive `serde::Serialize`:
//! 
//! ```rs
//! #[derive(serde::Serialize)]
//! struct MyParameters {
//!     name: String,
//! }
//! ```
//! 
//! Aftewards, it's a simple matter of running `decode` on your arguments:
//! 
//! ```rs
//! match function.as_str() {
//!     "my-function" => {
//!         let args: MyParameters = args
//!             .unwrap() // WARNING: Don't actually use unwrap() in your code :)
//!             .decode()
//!             .map_err(|err| Status::new(Code::InvalidArgument, err.to_string()))?;
//! 
//!         // TODO: Do something with the parameters...
//!     },
//! }
//! ```
//! 
//! We manually provide the type of `args` so that Rust knows what type to serialize to.
//! Afterwards we wrap any possible errors into a status code which can be returned back to the client.
//! 
//! ### Sending Data back to the Client
//! 
//! Now that we have all of the input data in check we can process our data and return it back to the client.
//! The way we do this is in the form of a data stream. Thanks to data streams we can return long segments of
//! data over time instead of having to return the whole data upfront. When we return a segment of data, we
//! also return it in the form of a `Result<>`. This is because individual segments of data may contain errors,
//! but the whole process can complete succesfully. You should return errors from the `call` function when there
//! is an irrecoverable error, but should send back an error packet when a *portion* of the internal logic fails.
//! 
//! Let's showcase all of this via an example:
//! 
//! ```rs
//! match function.as_str() {
//!     "my-function" => {
//!         let args: MyParameters = args
//!             .unwrap() // WARNING: Don't actually use unwrap() in your code :)
//!             .decode()
//!             .map_err(|err| Status::new(Code::InvalidArgument, err.to_string()))?;
//! 
//!         let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
//! 
//!         // We send back an Ok() packet to the client with an encoded message of our choice
//!         // (it can be anything that's serializable with serde!)
//!         tx.send(Ok(MessagePack::encode(format!("Hello, {}!", args.name)))).unwrap();
//! 
//!         Ok(UnboundedReceiverStream::new(rx))
//!     },
//! }
//! ```
//! 
//! First, we create a sender and receiver via tokio's `unbounded_channel()`. This allows us to send data to the client
//! and for the client to read data from the module. All return messages have to be encoded via `MessagePack::encode`.
//! 
//! ### Running the Module
//! 
//! Now that we have all of the code set up, create an asynchronous main function. In here we will instantiate our
//! module and kick it into full gear:
//! 
//! ```rs
//! #[tokio::main]
//! async fn main() {
//!     Module::new().start(MyModule::default())
//!         .await
//!         .unwrap()
//! }
//! ```
//! 
//! Voila! You now have a fundamental understanding of how modules communicate with Norgopolis and how to write your own
//! norgopolis module. Happy coding!

pub mod invoker_service;
mod stdio_service;

use std::time::Duration;

use futures::FutureExt;
use invoker_service::InvokerService;
use invoker_service::Service;
use module_communication::invoker_server::InvokerServer;
use stdio_service::StdioService;
use tokio::time::sleep;
use tokio_stream::wrappers::ReceiverStream;
use tonic::transport::Server;

pub use norgopolis_protos::module_communication;
pub use tonic::async_trait;
pub use tonic::{Code, Status};

/// Describes a module that can communicate with Norgopolis
/// over stdin/stdout.
pub struct Module {
    /// Timeout duration for the module. If no messages are received by the module after this time
    /// has passed the module will automatically shut down.
    ///
    /// Default is 5 minutes.
    pub timeout: Duration,
}

impl Default for Module {
    fn default() -> Self {
        Self::new()
    }
}

impl Module {
    pub fn new() -> Self {
        Module {
            timeout: Duration::from_secs(60 * 5),
        }
    }

    pub fn timeout(self, timeout: Duration) -> Self {
        Module { timeout }
    }

    pub async fn start<S>(self, service: S) -> Result<(), anyhow::Error>
    where
        S: Service + Sync + Send + 'static,
    {
        let (keepalive_tx, mut keepalive_rx) = tokio::sync::mpsc::unbounded_channel::<()>();

        tokio::spawn(async move {
            sleep(self.timeout).await;

            if keepalive_rx.recv().now_or_never().is_none() {
                std::process::exit(0);
            }

            // Drain the remained of the messages.
            while keepalive_rx.recv().now_or_never().is_some() {}
        });

        let (stdin, stdout) = (tokio::io::stdin(), tokio::io::stdout());
        let stdio_service = StdioService { stdin, stdout };

        // TODO: Do this in a better way
        // `once_stream` doesn't work :/
        let (tx, rx) = tokio::sync::mpsc::channel::<Result<StdioService, anyhow::Error>>(1);
        tx.send(Ok(stdio_service)).await?;

        Ok(Server::builder()
            .add_service(InvokerServer::new(InvokerService::new(
                service,
                keepalive_tx,
            )))
            .serve_with_incoming(ReceiverStream::new(rx))
            .await?)
    }
}