소스 검색

reqrep.rs catch ctrl-c SIGINT signal and gracefully stop the service

ghassmo 5 년 전
부모
커밋
63c3e8c9af
4개의 변경된 파일83개의 추가작업 그리고 36개의 파일을 삭제
  1. 2 0
      Cargo.toml
  2. 6 4
      src/error.rs
  3. 12 12
      src/service/gateway.rs
  4. 63 20
      src/service/reqrep.rs

+ 2 - 0
Cargo.toml

@@ -48,6 +48,7 @@ log = "0.4"
 ctrlc = "3.1.7"
 serde_json = "1.0.61"
 owning_ref = "0.4.1"
+signal-hook = "0.3.8"
 
 smol = "1.2.4"
 futures = "0.3.5"
@@ -86,6 +87,7 @@ bytes = "1.0.1"
 # wallet deps
 rocksdb = "0.16.0"
 dirs = "2.0.2"
+
 [dependencies.rusqlite]
 version = "0.25.1"
 features = ["bundled", "sqlcipher"]

+ 6 - 4
src/error.rs

@@ -35,8 +35,9 @@ pub enum Error {
     Utf8Error,
     NoteDecryptionFailed,
     ServicesError(&'static str),
-    ZMQError,
+    ZMQError(String),
     VerifyFailed,
+    TryIntoError,
 }
 
 impl std::error::Error for Error {}
@@ -69,16 +70,17 @@ impl fmt::Display for Error {
             Error::Utf8Error => f.write_str("Malformed UTF8"),
             Error::NoteDecryptionFailed => f.write_str("Unable to decrypt mint note"),
             Error::ServicesError(ref err) => write!(f, "Services error: {}", err),
-            Error::ZMQError => f.write_str("ZMQ error"),
+            Error::ZMQError(ref err) =>  write!(f, "ZMQError: {}", err),
             Error::VerifyFailed => f.write_str("Verify failed"),
+            Error::TryIntoError => f.write_str("TryInto get an error"),
         }
     }
 }
 
 // TODO: Match statement to parse external errors into strings.
 impl From<zeromq::ZmqError> for Error {
-    fn from(_err: zeromq::ZmqError) -> Error {
-        Error::ZMQError
+    fn from(err: zeromq::ZmqError) -> Error {
+        Error::ZMQError(err.to_string())
     }
 }
 

+ 12 - 12
src/service/gateway.rs

@@ -6,7 +6,6 @@ use super::reqrep::{Publisher, RepProtocol, Reply, ReqProtocol, Request, Subscri
 use crate::{Error, Result};
 
 use async_executor::Executor;
-
 use log::*;
 
 pub type Slabs = Vec<Vec<u8>>;
@@ -36,30 +35,28 @@ impl GatewayService {
     }
 
     pub async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) -> Result<()> {
-        let mut socket = RepProtocol::new(self.addr.clone());
 
-        let (send, recv) = socket.start().await?;
-        info!("server started: bind to {}", self.addr.to_string());
+        let mut protocol = RepProtocol::new(String::from("GATEWAY"),self.addr.clone());
 
-        self.publisher.lock().await.start().await?;
+        let (send, recv) = protocol.start().await?;
 
-        info!("publisher started");
+        self.publisher.lock().await.start().await?;
 
         let handle_request_task = executor.spawn(self.handle_request(send.clone(), recv.clone()));
 
-        socket.run().await?;
+        protocol.run(executor.clone()).await?;
 
-        handle_request_task.cancel().await;
+        let _ = handle_request_task.cancel().await;
         Ok(())
     }
 
+
     async fn handle_request(
         self: Arc<Self>,
         send_queue: async_channel::Sender<Reply>,
         recv_queue: async_channel::Receiver<Request>,
     ) -> Result<()> {
         let data = vec![];
-
         loop {
             match recv_queue.recv().await {
                 Ok(request) => {
@@ -83,15 +80,18 @@ impl GatewayService {
                             info!("received getlastindex msg");
                         }
                         _ => {
-                            return Err(Error::ServicesError("wrong command"));
+                            return Err(Error::ServicesError("received wrong command"));
                         }
                     }
                     let rep = Reply::from(&request, 0, data.clone());
                     send_queue.send(rep.into()).await?;
                 }
-                Err(_) => {}
+                Err(_) => {
+                    break;
+                }
             }
         }
+        Ok(())
     }
 }
 
@@ -132,7 +132,7 @@ impl GatewayClient {
             .protocol
             .request(GatewayCommand::GetLastIndex as u8, vec![])
             .await?;
-        let rep: [u8; 4] = rep.try_into().unwrap();
+        let rep: [u8; 4] = rep.try_into().map_err(|_| crate::Error::TryIntoError)?;
         Ok(u32::from_be_bytes(rep))
     }
 }

+ 63 - 20
src/service/reqrep.rs

@@ -1,3 +1,4 @@
+use async_std::sync::Arc;
 use std::io;
 use std::net::SocketAddr;
 
@@ -8,10 +9,16 @@ use bytes::Bytes;
 use futures::FutureExt;
 use rand::Rng;
 use zeromq::*;
+use signal_hook::{iterator::Signals, consts::SIGINT};
+use async_executor::Executor;
+use log::*;
+
+
 
 enum NetEvent {
     Receive(zeromq::ZmqMessage),
     Send(Reply),
+    Stop
 }
 
 pub fn addr_to_string(addr: SocketAddr) -> String {
@@ -19,6 +26,7 @@ pub fn addr_to_string(addr: SocketAddr) -> String {
 }
 
 pub struct RepProtocol {
+    service_name: String,
     addr: SocketAddr,
     socket: zeromq::RepSocket,
     recv_queue: async_channel::Receiver<Reply>,
@@ -30,7 +38,7 @@ pub struct RepProtocol {
 }
 
 impl RepProtocol {
-    pub fn new(addr: SocketAddr) -> RepProtocol {
+    pub fn new(service_name: String, addr: SocketAddr) -> RepProtocol {
         let socket = zeromq::RepSocket::new();
         let (send_queue, recv_channel) = async_channel::unbounded::<Request>();
         let (send_channel, recv_queue) = async_channel::unbounded::<Reply>();
@@ -38,6 +46,7 @@ impl RepProtocol {
         let channels = (send_channel.clone(), recv_channel.clone());
 
         RepProtocol {
+            service_name,
             addr,
             socket,
             recv_queue,
@@ -49,35 +58,59 @@ impl RepProtocol {
     pub async fn start(
         &mut self,
     ) -> Result<(
-        async_channel::Sender<Reply>,
-        async_channel::Receiver<Request>,
+    async_channel::Sender<Reply>,
+    async_channel::Receiver<Request>,
     )> {
         let addr = addr_to_string(self.addr);
         self.socket.bind(addr.as_str()).await?;
+        info!("{} SERVICE: started - bind to {}", self.service_name, addr);
         Ok(self.channels.clone())
     }
 
-    pub async fn run(&mut self) -> Result<()> {
+    pub async fn run(&mut self, executor: Arc<Executor<'_>>) -> Result<()> {
+
+        info!("{} SERVICE: running", self.service_name);
+
+        let (stop_s, stop_r) = async_channel::unbounded::<()>();
+
+        let mut signals = Signals::new(&[SIGINT])?;
+
+        let stop_task = executor.spawn(async move {
+            for _ in signals.forever() {
+                stop_s.send(()).await?;
+                break;
+            }
+            Ok::<(), crate::Error>(())
+        });
+
         loop {
             let event = futures::select! {
                 request = self.socket.recv().fuse() => NetEvent::Receive(request?),
-                reply = self.recv_queue.recv().fuse() => NetEvent::Send(reply?)
+                reply = self.recv_queue.recv().fuse() => NetEvent::Send(reply?),
+                _ = stop_r.recv().fuse() => NetEvent::Stop
             };
 
             match event {
                 NetEvent::Receive(request) => {
-                    let request: &Bytes = request.get(0).unwrap();
-                    let request: Vec<u8> = request.to_vec();
-                    let req: Request = deserialize(&request)?;
-                    self.send_queue.send(req).await?;
+                    if let Some(req) = request.get(0) {
+                        let request: Vec<u8> = req.to_vec();
+                        let req: Request = deserialize(&request)?;
+                        self.send_queue.send(req).await?;
+                    }
                 }
                 NetEvent::Send(reply) => {
                     let reply: Vec<u8> = serialize(&reply);
                     let reply = Bytes::from(reply);
                     self.socket.send(reply.into()).await?;
                 }
+                NetEvent::Stop => {
+                    break
+                }
             }
         }
+        let _ = stop_task.cancel().await;
+        warn!("{} SERVICE: stopped", self.service_name);
+        Ok(())
     }
 }
 
@@ -106,18 +139,21 @@ impl ReqProtocol {
         self.socket.send(req.into()).await?;
 
         let rep: zeromq::ZmqMessage = self.socket.recv().await?;
-        let rep: &Bytes = rep.get(0).unwrap();
-        let rep: Vec<u8> = rep.to_vec();
+        if let Some(rep) = rep.get(0) {
+            let rep: Vec<u8> = rep.to_vec();
 
-        let reply: Reply = deserialize(&rep)?;
+            let reply: Reply = deserialize(&rep)?;
 
-        if reply.has_error() {
-            return Err(crate::Error::ServicesError("response has an error"));
-        }
+            if reply.has_error() {
+                return Err(crate::Error::ServicesError("response has an error"));
+            }
 
-        assert!(reply.get_id() == request.get_id());
+            assert!(reply.get_id() == request.get_id());
 
-        Ok(reply.get_payload())
+            Ok(reply.get_payload())
+        } else {
+            Err(crate::Error::ZMQError("Couldn't parse ZmqMessage".to_string()))
+        }
     }
 }
 
@@ -166,9 +202,16 @@ impl Subscriber {
 
     pub async fn fetch(&mut self) -> Result<Vec<u8>> {
         let data = self.socket.recv().await?;
-        let data: &Bytes = data.get(0).unwrap();
-        let data = data.to_vec();
-        Ok(data)
+        match data.get(0) {
+            Some(d) => {
+                let data = d.to_vec();
+                Ok(data)
+            }
+            None => {
+                Err(crate::Error::ZMQError("Couldn't parse ZmqMessage".to_string()))
+            }
+        }
+
     }
 }