From 4ee2f3940819e408d45808bc5c1b44c1ec34dcc0 Mon Sep 17 00:00:00 2001
From: Sunli <scott_s829@163.com>
Date: Wed, 2 Aug 2023 09:43:46 +0800
Subject: [PATCH] a

---
 src/http/mod.rs       |  2 ++
 src/http/timer.rs     | 26 ++++++++++++++++++++++
 src/http/websocket.rs | 51 +++++++++++++------------------------------
 3 files changed, 43 insertions(+), 36 deletions(-)
 create mode 100644 src/http/timer.rs

diff --git a/src/http/mod.rs b/src/http/mod.rs
index ff3bc3ee3..20839c3cf 100644
--- a/src/http/mod.rs
+++ b/src/http/mod.rs
@@ -7,6 +7,7 @@ mod graphiql_v2_source;
 mod multipart;
 #[cfg(feature = "playground")]
 mod playground_source;
+mod timer;
 mod websocket;
 
 use std::io::ErrorKind;
@@ -21,6 +22,7 @@ pub use multipart::MultipartOptions;
 #[cfg(feature = "playground")]
 pub use playground_source::{playground_source, GraphQLPlaygroundConfig};
 use serde::Deserialize;
+pub use timer::Timer;
 pub use websocket::{
     ClientMessage, Protocols as WebSocketProtocols, WebSocket, WsMessage, ALL_WEBSOCKET_PROTOCOLS,
 };
diff --git a/src/http/timer.rs b/src/http/timer.rs
new file mode 100644
index 000000000..fd1ea1dad
--- /dev/null
+++ b/src/http/timer.rs
@@ -0,0 +1,26 @@
+use std::{
+    task::{Context, Poll},
+    time::Duration,
+};
+
+pub trait Timer {
+    fn new(interval: Duration) -> Self;
+
+    fn poll_tick(&mut self, cx: &mut Context<'_>) -> Poll<()>;
+
+    fn reset(&mut self);
+}
+
+pub struct DefaultTimer;
+
+impl Timer for DefaultTimer {
+    fn new(_interval: Duration) -> Self {
+        Self
+    }
+
+    fn poll_tick(&mut self, cx: &mut Context<'_>) -> Poll<()> {
+        Poll::Pending
+    }
+
+    fn reset(&mut self) {}
+}
diff --git a/src/http/websocket.rs b/src/http/websocket.rs
index 27604868d..ce569ac6e 100644
--- a/src/http/websocket.rs
+++ b/src/http/websocket.rs
@@ -6,6 +6,7 @@ use std::{
     pin::Pin,
     sync::Arc,
     task::{Context, Poll},
+    time::Duration,
 };
 
 use futures_util::{
@@ -16,51 +17,25 @@ use futures_util::{
 use pin_project_lite::pin_project;
 use serde::{Deserialize, Serialize};
 
-use crate::{Data, Error, Executor, Request, Response, Result};
+use crate::{http::Timer, Data, Error, Executor, Request, Response, Result};
 
 /// All known protocols based on WebSocket.
 pub const ALL_WEBSOCKET_PROTOCOLS: [&str; 2] = ["graphql-transport-ws", "graphql-ws"];
 
 /// An enum representing the various forms of a WebSocket message.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum WsMessage {
     /// A text WebSocket message
     Text(String),
 
-    /// A close message with the close frame.
+    /// A close message with the close frame
     Close(u16, String),
-}
 
-impl WsMessage {
-    /// Returns the contained [WsMessage::Text] value, consuming the `self`
-    /// value.
-    ///
-    /// Because this function may panic, its use is generally discouraged.
-    ///
-    /// # Panics
-    ///
-    /// Panics if the self value not equals [WsMessage::Text].
-    pub fn unwrap_text(self) -> String {
-        match self {
-            Self::Text(text) => text,
-            Self::Close(_, _) => panic!("Not a text message"),
-        }
-    }
+    /// A ping message with the specified payload
+    Ping(Vec<u8>),
 
-    /// Returns the contained [WsMessage::Close] value, consuming the `self`
-    /// value.
-    ///
-    /// Because this function may panic, its use is generally discouraged.
-    ///
-    /// # Panics
-    ///
-    /// Panics if the self value not equals [WsMessage::Close].
-    pub fn unwrap_close(self) -> (u16, String) {
-        match self {
-            Self::Close(code, msg) => (code, msg),
-            Self::Text(_) => panic!("Not a close message"),
-        }
-    }
+    /// A pong message with the specified payload
+    Pong(Vec<u8>),
 }
 
 pin_project! {
@@ -165,6 +140,13 @@ where
             protocol: self.protocol,
         }
     }
+
+    pub fn with_keep_alive<T>(self, timeout: Duration)
+    where
+        T: Timer,
+    {
+        let timer = Timer::new(timeout);
+    }
 }
 
 impl<S, E, OnInit, InitFut> Stream for WebSocket<S, E, OnInit>
@@ -439,7 +421,4 @@ enum ServerMessage<'a> {
         #[serde(skip_serializing_if = "Option::is_none")]
         payload: Option<serde_json::Value>,
     },
-    // Not used by this library
-    // #[serde(rename = "ka")]
-    // KeepAlive
 }