@@ -9,7 +9,13 @@ use tokio::{
99 net:: TcpStream ,
1010 sync:: mpsc:: { self , Receiver , Sender } ,
1111} ;
12- use tokio_tungstenite:: { tungstenite:: Message , MaybeTlsStream , WebSocketStream } ;
12+ use tokio_tungstenite:: {
13+ tungstenite:: {
14+ protocol:: { frame:: coding:: CloseCode , CloseFrame } ,
15+ Message ,
16+ } ,
17+ MaybeTlsStream , WebSocketStream ,
18+ } ;
1319
1420type Connection = WebSocketStream < MaybeTlsStream < TcpStream > > ;
1521
@@ -19,6 +25,15 @@ pub struct WebApi {
1925 queue : Vec < ClientRequest < ' static > > ,
2026}
2127
28+ impl Drop for WebApi {
29+ fn drop ( & mut self ) {
30+ let req = self . request_tx . clone ( ) ;
31+ tokio:: spawn ( async move {
32+ let _ = req. send ( ClientRequest :: Close ) . await ;
33+ } ) ;
34+ }
35+ }
36+
2237impl Stream for WebApi {
2338 type Item = HostResult ;
2439
@@ -114,6 +129,7 @@ impl WebApi {
114129 res. ok_or_else ( || ClientError :: from ( ErrorKind :: ChannelClosed ) ) ?
115130 }
116131
132+ #[ doc( hidden) ]
117133 pub async fn disconnect ( self , cause : impl Into < Cow < ' static , str > > ) {
118134 let _ = self
119135 . request_tx
@@ -148,6 +164,7 @@ async fn request_handler(
148164 tracing:: debug!( ?error, "request handler error" ) ;
149165 let error = match error {
150166 Error :: ChannelClosed => ErrorKind :: ChannelClosed . into ( ) ,
167+ Error :: ConnectionClosed => ErrorKind :: Disconnect . into ( ) ,
151168 other => ClientError :: from ( format ! ( "{other}" ) ) ,
152169 } ;
153170 let _ = response_tx. send ( Err ( error) ) . await ;
@@ -163,6 +180,17 @@ async fn process_request(
163180 . map_err ( Into :: into)
164181 . map_err ( Error :: OtherError ) ?;
165182 conn. send ( Message :: Binary ( msg. into ( ) ) ) . await ?;
183+ if let ClientRequest :: Disconnect { cause } = req {
184+ conn. close ( cause. map ( |c| CloseFrame {
185+ code : CloseCode :: Normal ,
186+ reason : format ! ( "{c}" ) . into ( ) ,
187+ } ) )
188+ . await ?;
189+ return Err ( Error :: ConnectionClosed ) ;
190+ } else if let ClientRequest :: Close = req {
191+ conn. close ( None ) . await ?;
192+ return Err ( Error :: ConnectionClosed ) ;
193+ }
166194 Ok ( ( ) )
167195}
168196
0 commit comments