feat(thrift): add fallback address option for shmipc

Author: Joshuahoky

volo-thrift/src/client/mod.rs

@@ -17,6 +17,8 @@ use motore::{
 };
 use pilota::thrift::TMessageType;
 use tokio::time::Duration;
+#[cfg(feature = "shmipc")]
+use volo::net::shmipc_fallback::ShmipcMakeTransportWithFallback;
 use volo::{
     FastStr,
     client::WithOptService,
@@ -495,6 +497,37 @@ impl<IL, OL, C, Req, Resp, MkT, MkC, LB> ClientBuilder<IL, OL, C, Req, Resp, MkT
         }
     }
 
+    #[cfg(feature = "shmipc")]
+    /// Set the address for the client with shmipc fallback.
+    pub fn address_with_fallback<A1: Into<Address>, A2: Into<Address>>(
+        self,
+        shmipc_addr: A1,
+        fallback_addr: A2,
+    ) -> ClientBuilder<IL, OL, C, Req, Resp, ShmipcMakeTransportWithFallback, MkC, LB> {
+        ClientBuilder {
+            address: Some(shmipc_addr.into()),
+            make_transport: ShmipcMakeTransportWithFallback::new(
+                DefaultMakeTransport::default(),
+                DefaultMakeTransport::default(),
+                fallback_addr.into(),
+            ),
+            config: self.config,
+            pool: self.pool,
+            caller_name: self.caller_name,
+            callee_name: self.callee_name,
+            inner_layer: self.inner_layer,
+            outer_layer: self.outer_layer,
+            mk_client: self.mk_client,
+            _marker: PhantomData,
+            make_codec: self.make_codec,
+            mk_lb: self.mk_lb,
+            disable_timeout_layer: self.disable_timeout_layer,
+            enable_biz_error: self.enable_biz_error,
+            #[cfg(feature = "multiplex")]
+            multiplex: self.multiplex,
+        }
+    }
+
     #[doc(hidden)]
     pub fn get_callee_name(&self) -> &FastStr {
         &self.callee_name

volo-thrift/src/server/mod.rs

@@ -16,6 +16,8 @@ use tokio::{
     sync::Notify,
 };
 use tracing::{info, trace};
+#[cfg(feature = "shmipc")]
+use volo::net::shmipc_fallback::ShmipcAddressWithFallback;
 use volo::{
     net::{
         Address,
@@ -408,6 +410,38 @@ impl<S, L, Req, MkC, SP> Server<S, L, Req, MkC, SP> {
         Ok(())
     }
 
+    #[cfg(feature = "shmipc")]
+    /// Run the server with shmipc and TCP fallback support.
+    ///
+    /// The server will listen on both the shmipc address and the fallback TCP address.
+    /// Clients can connect via either transport.
+    pub async fn run_with_fallback<A1, A2>(
+        self,
+        shmipc_addr: A1,
+        fallback_addr: A2,
+    ) -> Result<(), BoxError>
+    where
+        A1: Into<Address>,
+        A2: volo::net::incoming::MakeIncoming + Send,
+        A2::Incoming: Send,
+        L: Layer<BoxService<ServerContext, Req, S::Response, crate::ServerError>>,
+        MkC: MakeCodec<OwnedReadHalf, OwnedWriteHalf>,
+        L::Service: Service<ServerContext, Req, Response = S::Response, Error = crate::ServerError>
+            + Send
+            + 'static
+            + Sync,
+        S: Service<ServerContext, Req, Error = crate::ServerError> + Send + 'static + Sync,
+        S::Response: EntryMessage + Send + 'static + Sync,
+        Req: EntryMessage + Send + 'static,
+        SP: SpanProvider,
+    {
+        self.run(ShmipcAddressWithFallback {
+            shmipc_addr: shmipc_addr.into(),
+            default_mi: fallback_addr,
+        })
+        .await
+    }
+
     #[cfg(feature = "multiplex")]
     /// Use multiplexing to handle multiple requests in one connection.
     ///

volo/src/net/mod.rs

@@ -4,6 +4,8 @@ pub mod ext;
 pub mod incoming;
 #[cfg(feature = "shmipc")]
 pub mod shmipc;
+#[cfg(feature = "shmipc")]
+pub mod shmipc_fallback;
 #[cfg(feature = "__tls")]
 #[cfg_attr(docsrs, doc(cfg(any(feature = "rustls", feature = "native-tls"))))]
 pub mod tls;

volo/src/net/shmipc_fallback.rs

@@ -0,0 +1,123 @@
+use std::io;
+
+use futures::{FutureExt, future::Either};
+
+use super::{
+    Address, DefaultIncoming, MakeIncoming,
+    conn::{Conn, OwnedReadHalf, OwnedWriteHalf},
+    dial::{DefaultMakeTransport, MakeTransport},
+    incoming::Incoming,
+};
+
+pub struct ShmipcAddressWithFallback<MI> {
+    pub shmipc_addr: Address,
+    pub default_mi: MI,
+}
+
+impl<MI, I> MakeIncoming for ShmipcAddressWithFallback<MI>
+where
+    MI: MakeIncoming<Incoming = I> + Send,
+    I: Incoming + Send,
+{
+    type Incoming = ShmipcIncoming<I>;
+
+    async fn make_incoming(self) -> io::Result<Self::Incoming> {
+        Ok(ShmipcIncoming {
+            shmipc_listener: self.shmipc_addr.make_incoming().await?,
+            default_incoming: self.default_mi.make_incoming().await?,
+        })
+    }
+}
+
+#[derive(Debug)]
+pub struct ShmipcIncoming<I> {
+    shmipc_listener: DefaultIncoming,
+    default_incoming: I,
+}
+
+impl<I> Incoming for ShmipcIncoming<I>
+where
+    I: Incoming,
+{
+    async fn accept(&mut self) -> io::Result<Option<Conn>> {
+        self.try_next().await
+    }
+}
+
+impl<I> ShmipcIncoming<I>
+where
+    I: Incoming,
+{
+    async fn try_next(&mut self) -> io::Result<Option<Conn>> {
+        let shmipc_conn = self.shmipc_listener.accept().fuse();
+        let default_conn = self.default_incoming.accept().fuse();
+        futures::pin_mut!(shmipc_conn, default_conn);
+        match futures::future::select(shmipc_conn, default_conn).await {
+            Either::Left((conn, _)) => {
+                tracing::trace!("recv a conn from shmipc");
+                conn
+            }
+            Either::Right((conn, _)) => {
+                tracing::trace!("recv a conn from default");
+                conn
+            }
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct ShmipcMakeTransportWithFallback {
+    pub shmipc_mkt: DefaultMakeTransport,
+    pub default_mkt: DefaultMakeTransport,
+    pub fallback_addr: Address,
+}
+
+impl ShmipcMakeTransportWithFallback {
+    pub fn new(
+        shmipc: DefaultMakeTransport,
+        default_mkt: DefaultMakeTransport,
+        fallback_addr: Address,
+    ) -> Self {
+        Self {
+            shmipc_mkt: shmipc,
+            default_mkt,
+            fallback_addr,
+        }
+    }
+}
+
+impl MakeTransport for ShmipcMakeTransportWithFallback {
+    type ReadHalf = OwnedReadHalf;
+    type WriteHalf = OwnedWriteHalf;
+
+    async fn make_transport(
+        &self,
+        mut addr: Address,
+    ) -> io::Result<(Self::ReadHalf, Self::WriteHalf)> {
+        if addr.is_shmipc() {
+            match self.shmipc_mkt.make_transport(addr).await {
+                Ok(ret) => return Ok(ret),
+                Err(e) => {
+                    tracing::info!(
+                        "failed to connect to shmipc target: {e}, fallback to default target"
+                    );
+                    addr = self.fallback_addr.clone();
+                }
+            }
+        }
+
+        self.default_mkt.make_transport(addr).await
+    }
+
+    fn set_connect_timeout(&mut self, timeout: Option<std::time::Duration>) {
+        self.default_mkt.set_connect_timeout(timeout);
+    }
+
+    fn set_read_timeout(&mut self, timeout: Option<std::time::Duration>) {
+        self.default_mkt.set_read_timeout(timeout);
+    }
+
+    fn set_write_timeout(&mut self, timeout: Option<std::time::Duration>) {
+        self.default_mkt.set_write_timeout(timeout);
+    }
+}