sse support finished

hubertshelley · hubertshelley · commit db63705fab7a · 2023-06-13T14:39:09.000+08:00
diff --git a/examples/custom_response/src/main.rs b/examples/custom_response/src/main.rs
@@ -1,3 +1,4 @@
+use silent::header::CONTENT_TYPE;
 use silent::prelude::*;
 
 fn main() {
@@ -15,7 +16,9 @@ fn main() {
 }
 
 async fn custom_response(_req: Request) -> Result<Response> {
-    let res = Response::empty();
+    let mut res = Response::empty();
+    res.headers_mut()
+        .insert(CONTENT_TYPE, HeaderValue::from_static("text"));
     let html = r#"
     <!DOCTYPE html>
     <html lang="en">
@@ -27,7 +30,7 @@ async fn custom_response(_req: Request) -> Result<Response> {
         <h1>custom response</h1>
     </body>
     </html>"#;
-    let res = res.set_body(full(html));
+    res.set_body(full(html));
     Ok(res)
 }
 
diff --git a/examples/sse-chat/Cargo.toml b/examples/sse-chat/Cargo.toml
@@ -6,7 +6,7 @@ edition = "2021"
 
 [dependencies]
 futures-util = "0.3"
-silent = { path = "../../silent" }
+silent = { path = "../../silent", features = ["sse"] }
 once_cell = "1"
 parking_lot = "0.12"
 tokio = { version = "1", features = ["macros"] }
diff --git a/examples/sse-chat/src/main.rs b/examples/sse-chat/src/main.rs
@@ -4,17 +4,15 @@
 // port from https://github.com/seanmonstar/warp/blob/master/examples/sse_chat.rs
 
 use std::collections::HashMap;
-use std::result::Result as IoResult;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
-use futures_util::{future, FutureExt, Stream, StreamExt, TryFutureExt, TryStreamExt};
+use futures_util::{Stream, StreamExt, TryFutureExt, TryStreamExt};
 use once_cell::sync::Lazy;
 use parking_lot::Mutex;
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::UnboundedReceiverStream;
 
 use silent::prelude::*;
-use silent::sse::{self, keep_alive};
 
 type Users = Mutex<HashMap<usize, mpsc::UnboundedSender<Message>>>;
 
@@ -55,7 +53,7 @@ async fn chat_send(req: Request) -> Result<Response> {
     Ok(Response::empty())
 }
 
-fn get_connected(_req: Request) -> impl Stream<Item = Result<sse::Event>> + Send + 'static {
+async fn user_connected(_req: Request) -> Result<Response> {
     // Use a counter to assign a new unique ID for this user.
     let my_id = NEXT_USER_ID.fetch_add(1, Ordering::Relaxed);
 
@@ -75,47 +73,16 @@ fn get_connected(_req: Request) -> impl Stream<Item = Result<sse::Event>> + Send
 
     // Convert messages into Server-Sent Events and returns resulting stream.
     let stream = rx.map(|msg| match msg {
-        Message::UserId(my_id) => Ok(sse::Event::default().event("user").data(my_id.to_string())),
-        Message::Reply(reply) => Ok(sse::Event::default().data(reply)),
+        Message::UserId(my_id) => {
+            warn!("user {} disconnected", my_id);
+            Ok(SSEEvent::default().event("user").data(my_id.to_string()))
+        }
+        Message::Reply(reply) => {
+            warn!("sse_reply: {}", reply);
+            Ok(SSEEvent::default().data(reply))
+        }
     });
-    stream
-}
-
-async fn user_connected(_req: Request) -> Result<Response> {
-    let stream = get_connected(_req);
-    // SseKeepAlive::new(stream).streaming(res).ok();
-    let mut res = Response::empty()
-        .set_header(
-            HeaderName::from_static("Content-Type"),
-            HeaderValue::from_static("text/event-stream"),
-        )
-        .set_header(
-            HeaderName::from_static("Cache-Control"),
-            HeaderValue::from_static("no-cache"),
-        )
-        .set_header(
-            HeaderName::from_static("Connection"),
-            HeaderValue::from_static("keep-alive"),
-        )
-        .set_header(
-            HeaderName::from_static("Access-Control-Allow-Origin"),
-            HeaderValue::from_static("*"),
-        );
-    res.set_status(StatusCode::PARTIAL_CONTENT);
-    // res.set_body(stream_body(stream));
-    let event_stream = sse::keep_alive().stream(stream);
-    let body_stream = event_stream
-        .map_err(|error| {
-            // FIXME: error logging
-            error!("sse stream error: {}", error);
-            SilentError::BusinessError {
-                code: StatusCode::INTERNAL_SERVER_ERROR,
-                msg: "sse::keep error".to_string(),
-            }
-        })
-        .into_stream()
-        .and_then(|event| future::ready(Ok(event.to_string())));
-    res.set_body(stream_body(body_stream));
+    let res = sse_reply(stream);
     Ok(res)
 }
 
@@ -148,40 +115,38 @@ static INDEX_HTML: &str = r#"
         <title>SSE Chat</title>
     </head>
     <body>
-        <h1>SSE Chat</h1>
+        <h1>SSE chat</h1>
         <div id="chat">
             <p><em>Connecting...</em></p>
         </div>
-        <input type="text" id="msg" />
-        <button type="button" id="submit">Send</button>
-        <script>
-        const chat = document.getElementById('chat');
-        const msg = document.getElementById('msg');
-        const submit = document.getElementById('submit');
-        let sse = new EventSource(`http://${location.host}/chat`);
+        <input type="text" id="text" />
+        <button type="button" id="send">Send</button>
+        <script type="text/javascript">
+        var uri = 'http://' + location.host + '/chat';
+        var sse = new EventSource(uri);
+        function message(data) {
+            var line = document.createElement('p');
+            line.innerText = data;
+            chat.appendChild(line);
+        }
         sse.onopen = function() {
             chat.innerHTML = "<p><em>Connected!</em></p>";
         }
-        var userId;
+        var user_id;
         sse.addEventListener("user", function(msg) {
-            userId = msg.data;
+            user_id = msg.data;
         });
         sse.onmessage = function(msg) {
-            showMessage(msg.data);
+            message(msg.data);
         };
-        document.getElementById('submit').onclick = function() {
+        send.onclick = function() {
             var msg = text.value;
             var xhr = new XMLHttpRequest();
-            xhr.open("POST", `${uri}/${user_id}`, true);
+            xhr.open("POST", uri + '/' + user_id, true);
             xhr.send(msg);
             text.value = '';
-            showMessage('<You>: ' + msg);
+            message('<You>: ' + msg);
         };
-        function showMessage(data) {
-            const line = document.createElement('p');
-            line.innerText = data;
-            chat.appendChild(line);
-        }
         </script>
     </body>
 </html>
diff --git a/silent/Cargo.toml b/silent/Cargo.toml
@@ -19,6 +19,7 @@ default = ["server", "test", "static"]
 full = ["server", "wasi", "static", "test"]
 server = ["tokio/fs", "tokio/net", "tokio/rt-multi-thread", "tokio/signal"]
 ws = []
+sse = []
 static = []
 wasi = ["tokio/sync"]
 test = ["tokio/macros", "tokio/rt"]
diff --git a/silent/src/core/response.rs b/silent/src/core/response.rs
@@ -58,6 +58,16 @@ impl Response {
         self.headers.insert(key, value);
         self
     }
+    #[inline]
+    /// 设置响应header
+    pub fn headers(&self) -> &HeaderMap {
+        &self.headers
+    }
+    #[inline]
+    /// 设置响应header
+    pub fn headers_mut(&mut self) -> &mut HeaderMap {
+        &mut self.headers
+    }
     /// 设置响应header
     pub fn set_typed_header<H>(&mut self, header: H)
     where
diff --git a/silent/src/lib.rs b/silent/src/lib.rs
@@ -12,7 +12,8 @@ mod route;
 mod rt;
 #[cfg(feature = "server")]
 mod service;
-pub mod sse;
+#[cfg(feature = "sse")]
+mod sse;
 #[cfg(feature = "ws")]
 mod ws;
 
@@ -43,6 +44,8 @@ pub mod prelude {
     pub use crate::route::Route;
     #[cfg(feature = "server")]
     pub use crate::service::Server;
+    #[cfg(feature = "sse")]
+    pub use crate::sse::{sse_reply, SSEEvent};
     #[cfg(feature = "ws")]
     pub use crate::ws::{
         FnOnClose, FnOnConnect, FnOnNoneResultFut, FnOnReceive, FnOnSend, FnOnSendFut,
diff --git a/silent/src/service/hyper_service.rs b/silent/src/service/hyper_service.rs
@@ -34,8 +34,7 @@ impl HyperHandler {
         async move {
             match routes.clone().handle(req, remote_addr).await {
                 Ok(res) => {
-                    response.set_status(res.status_code);
-                    response.set_body(res.body);
+                    response = res;
                 }
                 Err((mes, code)) => {
                     tracing::error!("Failed to handle request: {:?}", mes);
diff --git a/silent/src/sse/event.rs b/silent/src/sse/event.rs
@@ -13,59 +13,59 @@ enum DataType {
 
 /// Server-sent event
 #[derive(Default, Debug)]
-pub struct Event {
+pub struct SSEEvent {
     id: Option<String>,
     data: Option<DataType>,
     event: Option<String>,
     comment: Option<String>,
     retry: Option<Duration>,
 }
 
-impl Event {
+impl SSEEvent {
     /// Set Server-sent event data
     /// data field(s) ("data:<content>")
-    pub fn data<T: Into<String>>(mut self, data: T) -> Event {
+    pub fn data<T: Into<String>>(mut self, data: T) -> SSEEvent {
         self.data = Some(DataType::Text(data.into()));
         self
     }
 
     /// Set Server-sent event data
     /// data field(s) ("data:<content>")
-    pub fn json_data<T: Serialize>(mut self, data: T) -> Result<Event, Error> {
+    pub fn json_data<T: Serialize>(mut self, data: T) -> Result<SSEEvent, Error> {
         self.data = Some(DataType::Json(serde_json::to_string(&data)?));
         Ok(self)
     }
 
     /// Set Server-sent event comment
     /// Comment field (":<comment-text>")
-    pub fn comment<T: Into<String>>(mut self, comment: T) -> Event {
+    pub fn comment<T: Into<String>>(mut self, comment: T) -> SSEEvent {
         self.comment = Some(comment.into());
         self
     }
 
     /// Set Server-sent event event
-    /// Event name field ("event:<event-name>")
-    pub fn event<T: Into<String>>(mut self, event: T) -> Event {
+    /// SSEEvent name field ("event:<event-name>")
+    pub fn event<T: Into<String>>(mut self, event: T) -> SSEEvent {
         self.event = Some(event.into());
         self
     }
 
     /// Set Server-sent event retry
     /// Retry timeout field ("retry:<timeout>")
-    pub fn retry(mut self, duration: Duration) -> Event {
+    pub fn retry(mut self, duration: Duration) -> SSEEvent {
         self.retry = Some(duration);
         self
     }
 
     /// Set Server-sent event id
     /// Identifier field ("id:<identifier>")
-    pub fn id<T: Into<String>>(mut self, id: T) -> Event {
+    pub fn id<T: Into<String>>(mut self, id: T) -> SSEEvent {
         self.id = Some(id.into());
         self
     }
 }
 
-impl fmt::Display for Event {
+impl fmt::Display for SSEEvent {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         if let Some(ref comment) = &self.comment {
             ":".fmt(f)?;
diff --git a/silent/src/sse/keep_alive.rs b/silent/src/sse/keep_alive.rs
diff --git a/silent/src/sse/mod.rs b/silent/src/sse/mod.rs
diff --git a/silent/src/sse/reply.rs b/silent/src/sse/reply.rs
