diff_yamux

diff --git a/39/kv/Cargo.toml b/39/kv/Cargo.toml
index 4591b6a..fba1f26 100644
--- a/39/kv/Cargo.toml
+++ b/39/kv/Cargo.toml
@@ -24,8 +24,10 @@ sled = "0.34" # sled db
 thiserror = "1" # 错误定义和处理
 tokio = { version = "1", features = ["full" ] } # 异步网络库
 tokio-rustls = "0.22"
+tokio-util = { version = "0.6", features = ["compat"]} # tokio 和 futures 的兼容性库
 tracing = "0.1" # 日志处理
 tracing-subscriber = "0.2" # 日志处理
+yamux = "0.9" # yamux 多路复用支持

 [dev-dependencies]
 async-prost = "0.2.1" # 支持把 protobuf 封装成 TCP frame
diff --git a/39/kv/src/client.rs b/39/kv/src/client.rs
index ef2603c..df6f41f 100644
--- a/39/kv/src/client.rs
+++ b/39/kv/src/client.rs
@@ -1,5 +1,5 @@
 use anyhow::Result;
-use kv5::{CommandRequest, ProstClientStream, TlsClientConnector};
+use kv5::{CommandRequest, ProstClientStream, TlsClientConnector, YamuxCtrl};
 use tokio::net::TcpStream;
 use tracing::info;

@@ -16,6 +16,10 @@ async fn main() -> Result<()> {
     let stream = TcpStream::connect(addr).await?;
     let stream = connector.connect(stream).await?;

+    // 打开一个 stream
+    let mut ctrl = YamuxCtrl::new_client(stream, None);
+    let stream = ctrl.open_stream().await?;
+
     let mut client = ProstClientStream::new(stream);

     // 生成一个 HSET 命令
diff --git a/39/kv/src/error.rs b/39/kv/src/error.rs
index a9888f6..b33c647 100644
--- a/39/kv/src/error.rs
+++ b/39/kv/src/error.rs
@@ -26,6 +26,8 @@ pub enum KvError {
     IoError(#[from] std::io::Error),
     #[error("TLS error")]
     TlsError(#[from] tokio_rustls::rustls::TLSError),
+    #[error("Yamux Connection error")]
+    YamuxConnectionError(#[from] yamux::ConnectionError),

     #[error("Internal error: {0}")]
     Internal(String),
diff --git a/39/kv/src/network/mod.rs b/39/kv/src/network/mod.rs
index ca33c4f..1b680ab 100644
--- a/39/kv/src/network/mod.rs
+++ b/39/kv/src/network/mod.rs
@@ -1,9 +1,11 @@
 mod frame;
+mod multiplex;
 mod stream;
 mod tls;

 pub use frame::{read_frame, FrameCoder};
 use futures::{SinkExt, StreamExt};
+pub use multiplex::YamuxCtrl;
 pub use stream::ProstStream;
 pub use tls::{TlsClientConnector, TlsServerAcceptor};

@@ -69,10 +71,12 @@ where

 #[cfg(test)]
 pub mod utils {
+    use anyhow::Result;
     use bytes::{BufMut, BytesMut};
-    use std::task::Poll;
+    use std::{cmp::min, task::Poll};
     use tokio::io::{AsyncRead, AsyncWrite};

+    #[derive(Default)]
     pub struct DummyStream {
         pub buf: BytesMut,
     }
@@ -83,8 +87,9 @@ pub mod utils {
             _cx: &mut std::task::Context<'_>,
             buf: &mut tokio::io::ReadBuf<'_>,
         ) -> Poll<std::io::Result<()>> {
-            let len = buf.capacity();
-            let data = self.get_mut().buf.split_to(len);
+            let this = self.get_mut();
+            let len = min(buf.capacity(), this.buf.len());
+            let data = this.buf.split_to(len);
             buf.put_slice(&data);
             Poll::Ready(Ok(()))
         }
@@ -118,14 +123,13 @@ pub mod utils {

 #[cfg(test)]
 mod tests {
-    use anyhow::Result;
-    use bytes::Bytes;
     use std::net::SocketAddr;
-    use tokio::net::{TcpListener, TcpStream};
-
-    use crate::{assert_res_ok, MemTable, ServiceInner, Value};

     use super::*;
+    use crate::{assert_res_ok, MemTable, ServiceInner, Value};
+    use anyhow::Result;
+    use bytes::Bytes;
+    use tokio::net::{TcpListener, TcpStream};

     #[tokio::test]
     async fn client_server_basic_communication_should_work() -> anyhow::Result<()> {
diff --git a/39/kv/src/network/multiplex.rs b/39/kv/src/network/multiplex.rs
new file mode 100644
index 0000000..b900a63
--- /dev/null
+++ b/39/kv/src/network/multiplex.rs
@@ -0,0 +1,178 @@
+use futures::{future, Future, TryStreamExt};
+use std::marker::PhantomData;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
+use yamux::{Config, Connection, ConnectionError, Control, Mode, WindowUpdateMode};
+
+/// Yamux 控制结构
+pub struct YamuxCtrl<S> {
+    /// yamux control，用于创建新的 stream
+    ctrl: Control,
+    _conn: PhantomData<S>,
+}
+
+impl<S> YamuxCtrl<S>
+where
+    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+{
+    /// 创建 yamux 客户端
+    pub fn new_client(stream: S, config: Option<Config>) -> Self {
+        Self::new(stream, config, true, |_stream| future::ready(Ok(())))
+    }
+
+    /// 创建 yamux 服务端，服务端我们需要具体处理 stream
+    pub fn new_server<F, Fut>(stream: S, config: Option<Config>, f: F) -> Self
+    where
+        F: FnMut(yamux::Stream) -> Fut,
+        F: Send + 'static,
+        Fut: Future<Output = Result<(), ConnectionError>> + Send + 'static,
+    {
+        Self::new(stream, config, false, f)
+    }
+
+    // 创建 YamuxCtrl
+    fn new<F, Fut>(stream: S, config: Option<Config>, is_client: bool, f: F) -> Self
+    where
+        F: FnMut(yamux::Stream) -> Fut,
+        F: Send + 'static,
+        Fut: Future<Output = Result<(), ConnectionError>> + Send + 'static,
+    {
+        let mode = if is_client {
+            Mode::Client
+        } else {
+            Mode::Server
+        };
+
+        // 创建 config
+        let mut config = config.unwrap_or_default();
+        config.set_window_update_mode(WindowUpdateMode::OnRead);
+
+        // 创建 config，yamux::Stream 使用的是 futures 的 trait 所以需要 compat() 到 tokio 的 trait
+        let conn = Connection::new(stream.compat(), config, mode);
+
+        // 创建 yamux ctrl
+        let ctrl = conn.control();
+
+        // pull 所有 stream 下的数据
+        tokio::spawn(yamux::into_stream(conn).try_for_each_concurrent(None, f));
+
+        Self {
+            ctrl,
+            _conn: PhantomData::default(),
+        }
+    }
+
+    /// 打开一个新的 stream
+    pub async fn open_stream(&mut self) -> Result<Compat<yamux::Stream>, ConnectionError> {
+        let stream = self.ctrl.open_stream().await?;
+        Ok(stream.compat())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::net::SocketAddr;
+
+    use super::*;
+    use crate::{
+        assert_res_ok,
+        network::tls::tls_utils::{tls_acceptor, tls_connector},
+        utils::DummyStream,
+        CommandRequest, KvError, MemTable, ProstClientStream, ProstServerStream, Service,
+        ServiceInner, Storage, TlsServerAcceptor,
+    };
+    use anyhow::Result;
+    use tokio::net::{TcpListener, TcpStream};
+    use tokio_rustls::server;
+    use tracing::warn;
+
+    pub async fn start_server_with<Store>(
+        addr: &str,
+        tls: TlsServerAcceptor,
+        store: Store,
+        f: impl Fn(server::TlsStream<TcpStream>, Service) + Send + Sync + 'static,
+    ) -> Result<SocketAddr, KvError>
+    where
+        Store: Storage,
+        Service: From<ServiceInner<Store>>,
+    {
+        let listener = TcpListener::bind(addr).await.unwrap();
+        let addr = listener.local_addr().unwrap();
+        let service: Service = ServiceInner::new(store).into();
+
+        tokio::spawn(async move {
+            loop {
+                match listener.accept().await {
+                    Ok((stream, _addr)) => match tls.accept(stream).await {
+                        Ok(stream) => f(stream, service.clone()),
+                        Err(e) => warn!("Failed to process TLS: {:?}", e),
+                    },
+                    Err(e) => warn!("Failed to process TCP: {:?}", e),
+                }
+            }
+        });
+
+        Ok(addr)
+    }
+
+    /// 创建 ymaux server
+    pub async fn start_yamux_server<Store>(
+        addr: &str,
+        tls: TlsServerAcceptor,
+        store: Store,
+    ) -> Result<SocketAddr, KvError>
+    where
+        Store: Storage,
+        Service: From<ServiceInner<Store>>,
+    {
+        let f = |stream, service: Service| {
+            let svc = service.clone();
+            YamuxCtrl::new_server(stream, None, move |s| {
+                let svc1 = svc.clone();
+                async move {
+                    let stream = ProstServerStream::new(s.compat(), svc1);
+                    stream.process().await.unwrap();
+                    Ok(())
+                }
+            });
+        };
+        start_server_with(addr, tls, store, f).await
+    }
+
+    #[tokio::test]
+    async fn yamux_ctrl_creation_should_work() -> Result<()> {
+        let s = DummyStream::default();
+        let mut ctrl = YamuxCtrl::new_client(s, None);
+        let stream = ctrl.open_stream().await;
+
+        assert!(stream.is_ok());
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn yamux_ctrl_client_server_should_work() -> Result<()> {
+        // 创建使用了 TLS 的 yamux server
+        let acceptor = tls_acceptor(false)?;
+        let addr = start_yamux_server("127.0.0.1:0", acceptor, MemTable::new()).await?;
+
+        let connector = tls_connector(false)?;
+        let stream = TcpStream::connect(addr).await?;
+        let stream = connector.connect(stream).await?;
+        // 创建使用了 TLS 的 yamux client
+        let mut ctrl = YamuxCtrl::new_client(stream, None);
+
+        // 从 client ctrl 中打开一个新的 yamux stream
+        let stream = ctrl.open_stream().await?;
+        // 封装成 ProstClientStream
+        let mut client = ProstClientStream::new(stream);
+
+        let cmd = CommandRequest::new_hset("t1", "k1", "v1".into());
+        client.execute(cmd).await.unwrap();
+
+        let cmd = CommandRequest::new_hget("t1", "k1");
+        let res = client.execute(cmd).await.unwrap();
+        assert_res_ok(res, &["v1".into()], &[]);
+
+        Ok(())
+    }
+}
diff --git a/39/kv/src/network/tls.rs b/39/kv/src/network/tls.rs
index 33ce47b..829d098 100644
--- a/39/kv/src/network/tls.rs
+++ b/39/kv/src/network/tls.rs
@@ -149,16 +149,8 @@ fn load_key(key: &str) -> Result<PrivateKey, KvError> {
 }

 #[cfg(test)]
-mod tests {
-
-    use std::net::SocketAddr;
-
-    use super::*;
-    use anyhow::Result;
-    use tokio::{
-        io::{AsyncReadExt, AsyncWriteExt},
-        net::{TcpListener, TcpStream},
-    };
+pub mod tls_utils {
+    use crate::{KvError, TlsClientConnector, TlsServerAcceptor};

     const CA_CERT: &str = include_str!("../../fixtures/ca.cert");
     const CLIENT_CERT: &str = include_str!("../../fixtures/client.cert");
@@ -166,13 +158,41 @@ mod tests {
     const SERVER_CERT: &str = include_str!("../../fixtures/server.cert");
     const SERVER_KEY: &str = include_str!("../../fixtures/server.key");

-    #[tokio::test]
-    async fn tls_should_work() -> Result<()> {
+    pub fn tls_connector(client_cert: bool) -> Result<TlsClientConnector, KvError> {
+        let ca = Some(CA_CERT);
+        let client_identity = Some((CLIENT_CERT, CLIENT_KEY));
+
+        match client_cert {
+            false => TlsClientConnector::new("kvserver.acme.inc", None, ca),
+            true => TlsClientConnector::new("kvserver.acme.inc", client_identity, ca),
+        }
+    }
+
+    pub fn tls_acceptor(client_cert: bool) -> Result<TlsServerAcceptor, KvError> {
         let ca = Some(CA_CERT);
+        match client_cert {
+            true => TlsServerAcceptor::new(SERVER_CERT, SERVER_KEY, ca),
+            false => TlsServerAcceptor::new(SERVER_CERT, SERVER_KEY, None),
+        }
+    }
+}

-        let addr = start_server(None).await?;
+#[cfg(test)]
+mod tests {
+    use super::tls_utils::tls_acceptor;
+    use crate::network::tls::tls_utils::tls_connector;
+    use anyhow::Result;
+    use std::net::SocketAddr;
+    use std::sync::Arc;
+    use tokio::{
+        io::{AsyncReadExt, AsyncWriteExt},
+        net::{TcpListener, TcpStream},
+    };

-        let connector = TlsClientConnector::new("kvserver.acme.inc", None, ca)?;
+    #[tokio::test]
+    async fn tls_should_work() -> Result<()> {
+        let addr = start_server(false).await?;
+        let connector = tls_connector(false)?;
         let stream = TcpStream::connect(addr).await?;
         let mut stream = connector.connect(stream).await?;
         stream.write_all(b"hello world!").await?;
@@ -185,12 +205,8 @@ mod tests {

     #[tokio::test]
     async fn tls_with_client_cert_should_work() -> Result<()> {
-        let client_identity = Some((CLIENT_CERT, CLIENT_KEY));
-        let ca = Some(CA_CERT);
-
-        let addr = start_server(ca.clone()).await?;
-
-        let connector = TlsClientConnector::new("kvserver.acme.inc", client_identity, ca)?;
+        let addr = start_server(true).await?;
+        let connector = tls_connector(true)?;
         let stream = TcpStream::connect(addr).await?;
         let mut stream = connector.connect(stream).await?;
         stream.write_all(b"hello world!").await?;
@@ -203,9 +219,10 @@ mod tests {

     #[tokio::test]
     async fn tls_with_bad_domain_should_not_work() -> Result<()> {
-        let addr = start_server(None).await?;
+        let addr = start_server(false).await?;

-        let connector = TlsClientConnector::new("kvserver1.acme.inc", None, Some(CA_CERT))?;
+        let mut connector = tls_connector(false)?;
+        connector.domain = Arc::new("kvserver1.acme.inc".into());
         let stream = TcpStream::connect(addr).await?;
         let result = connector.connect(stream).await;

@@ -214,8 +231,8 @@ mod tests {
         Ok(())
     }

-    async fn start_server(ca: Option<&str>) -> Result<SocketAddr> {
-        let acceptor = TlsServerAcceptor::new(SERVER_CERT, SERVER_KEY, ca)?;
+    async fn start_server(client_cert: bool) -> Result<SocketAddr> {
+        let acceptor = tls_acceptor(client_cert)?;

         let echo = TcpListener::bind("127.0.0.1:0").await.unwrap();
         let addr = echo.local_addr().unwrap();
diff --git a/39/kv/src/server.rs b/39/kv/src/server.rs
index cc175e2..2c9546b 100644
--- a/39/kv/src/server.rs
+++ b/39/kv/src/server.rs
@@ -1,6 +1,7 @@
 use anyhow::Result;
-use kv5::{MemTable, ProstServerStream, Service, ServiceInner, TlsServerAcceptor};
+use kv5::{MemTable, ProstServerStream, Service, ServiceInner, TlsServerAcceptor, YamuxCtrl};
 use tokio::net::TcpListener;
+use tokio_util::compat::FuturesAsyncReadCompatExt;
 use tracing::info;

 #[tokio::main]
@@ -13,6 +14,7 @@ async fn main() -> Result<()> {
     let server_key = include_str!("../fixtures/server.key");

     let acceptor = TlsServerAcceptor::new(server_cert, server_key, None)?;
+
     let service: Service = ServiceInner::new(MemTable::new()).into();
     let listener = TcpListener::bind(addr).await?;
     info!("Start listening on {}", addr);
@@ -20,8 +22,18 @@ async fn main() -> Result<()> {
         let tls = acceptor.clone();
         let (stream, addr) = listener.accept().await?;
         info!("Client {:?} connected", addr);
-        let stream = tls.accept(stream).await?;
-        let stream = ProstServerStream::new(stream, service.clone());
-        tokio::spawn(async move { stream.process().await });
+
+        let svc = service.clone();
+        tokio::spawn(async move {
+            let stream = tls.accept(stream).await.unwrap();
+            YamuxCtrl::new_server(stream, None, move |stream| {
+                let svc1 = svc.clone();
+                async move {
+                    let stream = ProstServerStream::new(stream.compat(), svc1.clone());
+                    stream.process().await.unwrap();
+                    Ok(())
+                }
+            });
+        });
     }
 }