From cbb76da108241eb5c40f52611a6a9f39c62ff091 Mon Sep 17 00:00:00 2001
From: yukang <moorekang@gmail.com>
Date: Mon, 29 Jan 2024 16:47:22 +0800
Subject: [PATCH] Fix the websockets terminate issue

---
 rpc/src/server.rs | 20 ++++++--------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/rpc/src/server.rs b/rpc/src/server.rs
index 07a4c20eed..e164ddcdbb 100644
--- a/rpc/src/server.rs
+++ b/rpc/src/server.rs
@@ -88,8 +88,10 @@ impl RpcServer {
     ) -> Result<SocketAddr, AnyError> {
         let stream_config = StreamServerConfig::default()
             .with_keep_alive(true)
-            .with_channel_size(4)
-            .with_pipeline_size(4);
+            .with_pipeline_size(4)
+            .with_shutdown(async move {
+                new_tokio_exit_rx().cancelled().await;
+            });
 
         // HTTP and WS server.
         let post_router = post(handle_jsonrpc::<Option<Session>>);
@@ -100,24 +102,14 @@ impl RpcServer {
         };
         let method_router = post_router.merge(get_router);
 
-        let mut app = Router::new()
+        let app = Router::new()
             .route("/", method_router.clone())
             .route("/*path", method_router)
             .route("/ping", get(ping_handler))
             .layer(Extension(Arc::clone(rpc)))
             .layer(CorsLayer::permissive())
             .layer(TimeoutLayer::new(Duration::from_secs(30)))
-            .layer(Extension(stream_config.clone()));
-
-        if enable_websocket {
-            let ws_config: StreamServerConfig =
-                stream_config
-                    .with_keep_alive(true)
-                    .with_shutdown(async move {
-                        new_tokio_exit_rx().cancelled().await;
-                    });
-            app = app.layer(Extension(ws_config));
-        }
+            .layer(Extension(stream_config));
 
         let (tx_addr, rx_addr) = tokio::sync::oneshot::channel::<SocketAddr>();