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

//! ump server running in an async task.
//!
//! # Example
//! ```
//! # tokio_test::block_on(async {
//! use std::ops::ControlFlow;
//! use ump_server::{
//!   async_trait,
//!   task::{Handler, spawn},
//!   ReplyContext, WeakClient
//! };
//!
//! enum Request {
//!   Add(usize, usize)
//! }
//! enum Reply {
//!   Sum(usize)
//! }
//! #[derive(Debug)]
//! enum MyError { }
//!
//! struct MyHandler {
//!   wclnt: WeakClient<Request, Reply, MyError>
//! };
//! #[async_trait]
//! impl Handler<Request, Reply, MyError, ()> for MyHandler {
//!   async fn proc_req(
//!     &mut self,
//!     msg: Request,
//!     rctx: ReplyContext<Reply, MyError>
//!   ) -> ControlFlow<(), ()> {
//!     match msg {
//!       Request::Add(a, b) => {
//!         rctx.reply(Reply::Sum(a + b));
//!         ControlFlow::Continue(())
//!       }
//!     }
//!   }
//! }
//!
//! let (clnt, jh) = spawn(|clnt| {
//!   // Store a weak client in the handler so it doesn't keep the dispatch
//!   // loop alive when the Client returned to the application is dropped.
//!   Ok(MyHandler {
//!     wclnt: clnt.weak()
//!   })
//! }).unwrap();
//!
//! let Ok(Reply::Sum(sum)) = clnt.areq(Request::Add(3, 7)).await else {
//!   panic!("Unexpected reply");
//! };
//! assert_eq!(sum, 10);
//!
//! // Dropping the only client will terminate the dispatch loop
//! drop(clnt);
//!
//! let _ = jh.await;
//! # });
//! ```

use std::ops::ControlFlow;

use tokio::task::{self, JoinHandle};

use async_trait::async_trait;

use super::{channel, Client, ReplyContext};

/// Message processing trait for an async handler.
#[async_trait]
pub trait Handler<S, R, E, RV> {
  /// Optional initialization callback.
  ///
  /// This is called on the dispatcher task before the main message dispatch
  /// loop is entered.
  #[allow(unused_variables)]
  fn init(&mut self, weak_client: ump::WeakClient<S, R, E>) {}

  /// Message processing callback.
  ///
  /// The callback must return `ControlFlow::Continue(())` to keep the
  /// dispatcher loop going.  Returning `ControlFlow::Break(RV)` will cause the
  /// dispatcher loop to abort and returns the value in `RV` from the task.
  async fn proc_req(
    &mut self,
    msg: S,
    rctx: ReplyContext<R, E>
  ) -> ControlFlow<RV, ()>;

  /// Optional termination callback.
  ///
  /// This is called on the dispatcher task just after the main message
  /// processing loop has been terminated.
  ///
  /// The `rv` argument is set to the return value returned from the dispatcher
  /// loop.  It will be set to `Some()` value if a request handler returned
  /// `ControlFlow::Break(RV)`.  If will be set to `None` if the dispatch loop
  /// terminated because the queue is empty and all of the linked clients have
  /// been dropped.
  ///
  /// The value returned from this callback is returned from the dispatcher
  /// task when it is joined.
  ///
  /// The default implementation simply returns the `rv` parameter.
  fn term(&mut self, rv: Option<RV>) -> Option<RV> {
    rv
  }
}

/// Run a task which will process incoming messages from an ump server
/// end-point.
///
/// See top module's documentation for an overview of the [dispatch
/// loop](crate#dispatch-loop).
#[allow(clippy::type_complexity)]
pub fn spawn<S, R, E, RV, F>(
  hbldr: impl FnOnce(&Client<S, R, E>) -> Result<F, E>
) -> Result<(Client<S, R, E>, JoinHandle<Option<RV>>), E>
where
  S: 'static + Send,
  R: 'static + Send,
  E: 'static + Send,
  RV: 'static + Send,
  F: Handler<S, R, E, RV> + Send + 'static
{
  let (server, client) = channel();

  let mut handler = hbldr(&client)?;

  #[cfg(feature = "watchdog")]
  let wdog = crate::wdog::run();

  let weak_client = client.weak();
  let jh = task::spawn(async move {
    handler.init(weak_client);
    let ret = loop {
      let (msg, rctx) = match server.async_wait().await {
        Ok(d) => d,
        Err(_) => break None
      };

      #[cfg(feature = "watchdog")]
      wdog.begin_process();

      let res = handler.proc_req(msg, rctx).await;

      #[cfg(feature = "watchdog")]
      wdog.end_process();

      match res {
        ControlFlow::Continue(_) => {}
        ControlFlow::Break(rv) => break Some(rv)
      }
    };

    #[cfg(feature = "watchdog")]
    let _ = wdog.kill();

    handler.term(ret)
  });

  Ok((client, jh))
}

// vim: set ft=rust et sw=2 ts=2 sts=2 cinoptions=2 tw=79 :