ttrpc: support connection recovery with fdstore

Cargo.toml

@@ -17,10 +17,12 @@ nix = "0.23.0"
 log = "0.4"
 byteorder = "1.3.2"
 thiserror = "1.0"
+uuid = { version = "1.1.2", features = ["v4"] }
 
 async-trait = { version = "0.1.31", optional = true }
-tokio = { version = "1", features = ["rt", "sync", "io-util", "macros", "time"], optional = true }
+tokio = { version = "1", features = ["rt", "sync", "io-util", "macros", "time", "fs"], optional = true }
 futures = { version = "0.3", optional = true }
+libsystemd = { version = "0.7.0", optional = true }
 
 [target.'cfg(target_os = "linux")'.dependencies]
 tokio-vsock = { version = "0.3.1", optional = true }
@@ -32,6 +34,7 @@ protobuf-codegen = "3.1.0"
 default = ["sync"]
 async = ["async-trait", "tokio", "futures", "tokio-vsock"]
 sync = []
+fdstore = ["libsystemd"]
 
 [package.metadata.docs.rs]
 all-features = true

src/asynchronous/client.rs

@@ -54,9 +54,16 @@ impl Client {
             rx: Some(rx),
             streams: req_map.clone(),
         };
-
-        let conn = Connection::new(stream, delegate);
-        tokio::spawn(async move { conn.run().await });
+        #[cfg(not(feature = "fdstore"))]
+        {
+            let conn = Connection::new(stream, delegate);
+            tokio::spawn(async move { conn.run().await });
+        }
+        #[cfg(feature = "fdstore")]
+        {
+            let conn = Connection::new(stream, delegate, "".to_string());
+            tokio::spawn(async move { conn.run().await });
+        }
 
         Client {
             req_tx,
@@ -254,7 +261,7 @@ impl ReaderDelegate for ClientReader {
 
     async fn exit(&self) {}
 
-    async fn handle_msg(&self, msg: GenMessage) {
+    async fn handle_msg(&self, _id: u64, msg: GenMessage) {
         let req_map = self.streams.clone();
         tokio::spawn(async move {
             let resp_tx = match msg.header.type_ {

src/asynchronous/connection.rs

@@ -17,6 +17,8 @@ use crate::error::Error;
 use crate::proto::GenMessage;
 
 use super::stream::SendingMessage;
+#[cfg(feature = "fdstore")]
+use crate::r#async::fdstore::MessageStore;
 
 pub trait Builder {
     type Reader;
@@ -37,10 +39,13 @@ pub trait ReaderDelegate {
     async fn wait_shutdown(&self);
     async fn disconnect(&self, e: Error, task: &mut task::JoinHandle<()>);
     async fn exit(&self);
-    async fn handle_msg(&self, msg: GenMessage);
+    // handle message with id, the id is only for message store.
+    async fn handle_msg(&self, id: u64, msg: GenMessage);
 }
 
 pub struct Connection<S, B: Builder> {
+    #[cfg(feature = "fdstore")]
+    name: String,
     reader: ReadHalf<S>,
     writer_task: task::JoinHandle<()>,
     reader_delegate: B::Reader,
@@ -53,7 +58,7 @@ where
     B::Reader: ReaderDelegate + Send + Sync + 'static,
     B::Writer: WriterDelegate + Send + Sync + 'static,
 {
-    pub fn new(conn: S, mut builder: B) -> Self {
+    pub fn new(conn: S, mut builder: B, #[cfg(feature = "fdstore")] name: String) -> Self {
         let (reader, mut writer) = split(conn);
 
         let (reader_delegate, mut writer_delegate) = builder.build();
@@ -73,6 +78,8 @@ where
         });
 
         Self {
+            #[cfg(feature = "fdstore")]
+            name,
             reader,
             writer_task,
             reader_delegate,
@@ -81,6 +88,8 @@ where
 
     pub async fn run(self) -> std::io::Result<()> {
         let Connection {
+            #[cfg(feature = "fdstore")]
+                name: _,
             mut reader,
             mut writer_task,
             reader_delegate,
@@ -91,7 +100,7 @@ where
                     match res {
                         Ok(msg) => {
                             trace!("Got Message {:?}", msg);
-                            reader_delegate.handle_msg(msg).await;
+                            reader_delegate.handle_msg(0, msg).await;
                         }
                         Err(e) => {
                             trace!("Read msg err: {:?}", e);
@@ -111,4 +120,63 @@ where
 
         Ok(())
     }
+
+    #[cfg(feature = "fdstore")]
+    pub async fn run_with_message_store(self, message_store: MessageStore) -> std::io::Result<()> {
+        let Connection {
+            name,
+            mut reader,
+            mut writer_task,
+            reader_delegate,
+        } = self;
+
+        let messages = message_store.get_messages(&name).await;
+        // the next message id should be larger than the message id before restart.
+        let mut id = 0u64;
+        for m in messages {
+            if m.id >= id {
+                id = m.id + 1;
+            }
+            // handle the stored request
+            reader_delegate.handle_msg(m.id, m.message).await;
+        }
+        loop {
+            select! {
+                res = GenMessage::read_from(&mut reader) => {
+                    match res {
+                        Ok(msg) => {
+                            trace!("Got Message {:?}", msg);
+                            message_store.insert(name.clone(), id, msg.clone()).await;
+                            reader_delegate.handle_msg(id, msg).await;
+                            id += 1;
+                        }
+                        Err(e) => {
+                            trace!("Read msg err: {:?}", e);
+                            reader_delegate.disconnect(e, &mut writer_task).await;
+                            break;
+                        }
+                    }
+                }
+                _v = reader_delegate.wait_shutdown() => {
+                    trace!("Receive shutdown.");
+                    break;
+                }
+            }
+        }
+        reader_delegate.exit().await;
+        #[cfg(feature = "fdstore")]
+        if let Err(e) = libsystemd::daemon::notify_with_fds(
+            false,
+            &[
+                libsystemd::daemon::NotifyState::Fdname(name.to_string()),
+                libsystemd::daemon::NotifyState::FdstoreRemove,
+            ],
+            &[],
+        ) {
+            warn!("failed to notify systemd to remove the fd {}: {}", name, e);
+        }
+        trace!("Reader task exit.");
+
+        Ok(())
+    }
 }

src/asynchronous/fdstore.rs

@@ -0,0 +1,157 @@
+use crate::{error::Result, proto::GenMessage, Error};
+use std::{collections::HashMap, io::ErrorKind, ops::DerefMut, sync::Arc};
+use tokio::{
+    fs::File,
+    io::{AsyncSeekExt, AsyncWriteExt},
+    sync::Mutex,
+};
+
+const SOCK_NAME_LEN: usize = 36;
+
+#[derive(Clone)]
+pub struct MessageStore {
+    file: Arc<Mutex<File>>,
+    cache: Arc<Mutex<HashMap<String, Vec<SockMessage>>>>,
+}
+
+#[derive(Default, Debug, Clone, PartialEq, Eq)]
+pub struct SockMessage {
+    pub(crate) sock_name: String,
+    pub(crate) id: u64,
+    pub(crate) message: GenMessage,
+}
+
+impl SockMessage {
+    pub async fn write_to(&self, mut writer: impl tokio::io::AsyncWriteExt + Unpin) -> Result<()> {
+        writer
+            .write_all(self.sock_name.as_bytes())
+            .await
+            .map_err(|e| Error::Others(e.to_string()))?;
+        writer
+            .write_u64(self.id)
+            .await
+            .map_err(|e| Error::Others(e.to_string()))?;
+        self.message.write_to(writer).await?;
+        Ok(())
+    }
+
+    pub async fn read_from(mut reader: impl tokio::io::AsyncReadExt + Unpin) -> Result<Self> {
+        let mut sock_name_buf = vec![0u8; SOCK_NAME_LEN];
+        let len = reader.read_exact(&mut sock_name_buf).await.map_err(|e| {
+            if e.kind() == ErrorKind::UnexpectedEof {
+                Error::Eof
+            } else {
+                Error::Others(format!("failed to read messages from memfd {}", e))
+            }
+        })?;
+        if len < SOCK_NAME_LEN {
+            return Err(Error::Others(format!("read {} bytes for socket name", len)));
+        }
+        let sock_name =
+            String::from_utf8(sock_name_buf).map_err(|e| Error::Others(e.to_string()))?;
+        let id = reader
+            .read_u64()
+            .await
+            .map_err(|e| Error::Others(e.to_string()))?;
+        let message = GenMessage::read_from(reader).await?;
+        Ok(Self {
+            sock_name,
+            id,
+            message,
+        })
+    }
+}
+
+impl MessageStore {
+    pub async fn load(mut f: File) -> Result<Self> {
+        f.seek(std::io::SeekFrom::Start(0))
+            .await
+            .map_err(|e| Error::Others(e.to_string()))?;
+        let s = Self {
+            file: Arc::new(Mutex::new(f)),
+            cache: Arc::new(Mutex::new(HashMap::new())),
+        };
+        let file = s.file.clone();
+        let mut f = file.lock().await;
+        loop {
+            let a = SockMessage::read_from(f.deref_mut()).await;
+            match a {
+                Ok(m) => {
+                    trace!("load a message from {}, with id {}", m.sock_name, m.id);
+                    s.insert_sock_message(m).await;
+                }
+                Err(e) => match e {
+                    Error::Eof => {
+                        break;
+                    }
+                    _ => {
+                        return Err(Error::Others(format!(
+                            "failed to read message from memfd in fdstore: {}",
+                            e
+                        )));
+                    }
+                },
+            }
+        }
+        return Ok(s);
+    }
+
+    pub async fn insert(&self, sock_name: String, id: u64, m: GenMessage) {
+        assert_eq!(SOCK_NAME_LEN, sock_name.len());
+        self.insert_sock_message(SockMessage {
+            sock_name,
+            id,
+            message: m,
+        })
+        .await;
+        self.dump().await;
+    }
+
+    pub async fn dump(&self) {
+        let mut file = self.file.lock().await;
+        file.set_len(0).await.unwrap_or_default();
+        file.rewind().await.unwrap_or_default();
+        let cache = self.cache.lock().await;
+        for v in cache.values() {
+            for m in v {
+                m.write_to(file.deref_mut()).await.unwrap_or_default();
+            }
+        }
+        file.flush().await.unwrap_or_default();
+    }
+
+    pub async fn get_messages(&self, key: &str) -> Vec<SockMessage> {
+        let mut res = vec![];
+        let cache = self.cache.lock().await;
+        if let Some(l) = cache.get(key) {
+            for m in l {
+                res.push(m.clone());
+            }
+        }
+        res
+    }
+
+    pub async fn remove(&self, sock_name: String, id: u64) {
+        self.remove_sock_message(sock_name, id).await;
+        self.dump().await;
+    }
+
+    async fn remove_sock_message(&self, sock_name: String, id: u64) {
+        let mut cache = self.cache.lock().await;
+        if let Some(v) = cache.get_mut(&sock_name) {
+            v.retain(|x| x.id != id);
+        }
+    }
+
+    async fn insert_sock_message(&self, m: SockMessage) {
+        let mut cache = self.cache.lock().await;
+        let sock_name = m.sock_name.clone();
+        if let Some(v) = cache.get_mut(&sock_name) {
+            v.push(m);
+        } else {
+            let mut l = Vec::new();
+            l.push(m);
+            cache.insert(sock_name, l);
+        }
+    }
+}

src/asynchronous/mod.rs

@@ -12,6 +12,8 @@ mod stream;
 #[doc(hidden)]
 mod utils;
 mod connection;
+#[cfg(feature = "fdstore")]
+mod fdstore;
 pub mod shutdown;
 mod unix_incoming;