Handle Voice close codes, prevent Songbird spinning WS threads (#1068)

Voice `CloseCode`s now map to a type rather than a collection of constants. Correct close code handling in this way terminates the websocket task when there is no likelihood of resuming, which was causing leftover tasks to spin at the `tokio::select` in some circumstances (i.e., ::leave, which keeps the `Driver` alive).

songbird/src/driver/tasks/ws.rs

@@ -1,13 +1,16 @@
-use super::{error::Result, message::*};
+use super::message::*;
 use crate::{
     events::CoreContext,
     model::{
         payload::{Heartbeat, Speaking},
+        CloseCode as VoiceCloseCode,
         Event as GatewayEvent,
+        FromPrimitive,
         SpeakingState,
     },
     ws::{Error as WsError, ReceiverExt, SenderExt, WsStream},
 };
+use async_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
 use flume::Receiver;
 use rand::random;
 use std::time::Duration;
@@ -52,14 +55,15 @@ impl AuxNetwork {
 
         loop {
             let mut ws_error = false;
+            let mut should_reconnect = false;
 
             let hb = time::delay_until(next_heartbeat);
 
             tokio::select! {
                 _ = hb => {
                     ws_error = match self.send_heartbeat().await {
                         Err(e) => {
-                            error!("Heartbeat send failure {:?}.", e);
+                            should_reconnect = ws_error_is_not_final(e);
                             true
                         },
                         _ => false,
@@ -73,7 +77,7 @@ impl AuxNetwork {
                             false
                         },
                         Err(e) => {
-                            error!("Error processing ws {:?}.", e);
+                            should_reconnect = ws_error_is_not_final(e);
                             true
                         },
                         Ok(Some(msg)) => {
@@ -113,7 +117,7 @@ impl AuxNetwork {
 
                                 ws_error |= match ssu_status {
                                     Err(e) => {
-                                        error!("Issue sending speaking update {:?}.", e);
+                                        should_reconnect = ws_error_is_not_final(e);
                                         true
                                     },
                                     _ => false,
@@ -128,8 +132,13 @@ impl AuxNetwork {
             }
 
             if ws_error {
-                let _ = interconnect.core.send(CoreMessage::Reconnect);
                 self.dont_send = true;
+
+                if should_reconnect {
+                    let _ = interconnect.core.send(CoreMessage::Reconnect);
+                } else {
+                    break;
+                }
             }
         }
     }
@@ -138,7 +147,7 @@ impl AuxNetwork {
         Instant::now() + self.heartbeat_interval
     }
 
-    async fn send_heartbeat(&mut self) -> Result<()> {
+    async fn send_heartbeat(&mut self) -> Result<(), WsError> {
         let nonce = random::<u64>();
         self.last_heartbeat_nonce = Some(nonce);
 
@@ -203,3 +212,21 @@ pub(crate) async fn runner(
     aux.run(&mut interconnect).await;
     info!("WS thread finished.");
 }
+
+fn ws_error_is_not_final(err: WsError) -> bool {
+    match err {
+        WsError::WsClosed(Some(frame)) => match frame.code {
+            CloseCode::Library(l) =>
+                if let Some(code) = VoiceCloseCode::from_u16(l) {
+                    code.should_resume()
+                } else {
+                    true
+                },
+            _ => true,
+        },
+        e => {
+            error!("Error sending/receiving ws {:?}.", e);
+            true
+        },
+    }
+}

voice-model/Cargo.toml

@@ -14,6 +14,7 @@ edition = "2018"
 
 [dependencies]
 bitflags = "1"
+enum_primitive = "0.1"
 serde_repr = "0.1"
 
 [dependencies.serde]

voice-model/src/close_code.rs

@@ -0,0 +1,60 @@
+use enum_primitive::*;
+
+enum_from_primitive! {
+/// Discord Voice Gateway Websocket close codes.
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub enum CloseCode {
+    /// Invalid Voice OP Code.
+    UnknownOpcode = 4001,
+
+    /// Invalid identification payload sent.
+    InvalidPayload = 4002,
+
+    /// A payload was sent prior to identifying.
+    NotAuthenticated = 4003,
+
+    /// The account token sent with the identify payload was incorrect.
+    AuthenticationFailed = 4004,
+
+    /// More than one identify payload was sent.
+    AlreadyAuthenticated = 4005,
+
+    /// The session is no longer valid.
+    SessionInvalid = 4006,
+
+    /// A session timed out.
+    SessionTimeout = 4009,
+
+    /// The server for the last connection attempt could not be found.
+    ServerNotFound = 4011,
+
+    /// Discord did not recognise the voice protocol chosen.
+    UnknownProtocol = 4012,
+
+    /// Disconnected, either due to channel closure/removal
+    /// or kicking.
+    ///
+    /// Should not reconnect.
+    Disconnected = 4014,
+
+    /// Connected voice server crashed.
+    ///
+    /// Should resume.
+    VoiceServerCrash = 4015,
+
+    /// Discord didn't recognise the encryption scheme.
+    UnknownEncryptionMode = 4016,
+}
+}
+
+impl CloseCode {
+    /// Indicates whether a voice client should attempt to reconnect in response to this close code.
+    ///
+    /// Otherwise, the connection should be closed.
+    pub fn should_resume(&self) -> bool {
+        match self {
+            CloseCode::VoiceServerCrash | CloseCode::SessionTimeout => true,
+            _ => false,
+        }
+    }
+}

voice-model/src/constants.rs

@@ -2,43 +2,3 @@
 
 /// Gateway version of the Voice API which this library encodes.
 pub const GATEWAY_VERSION: u8 = 4;
-
-pub mod voice_close_codes {
-    /// Invalid Voice OP Code.
-    pub const UNKNOWN_OPCODE: u16 = 4001;
-
-    /// A payload was sent prior to identifying.
-    pub const NOT_AUTHENTICATED: u16 = 4003;
-
-    /// The account token sent with the identify payload was incorrect.
-    pub const AUTHENTICATION_FAILED: u16 = 4004;
-
-    /// More than one identify payload was sent.
-    pub const ALREADY_AUTHENTICATED: u16 = 4005;
-
-    /// The session is no longer valid.
-    pub const SESSION_INVALID: u16 = 4006;
-
-    /// A session timed out.
-    pub const SESSION_TIMEOUT: u16 = 4009;
-
-    /// The server for the last connection attempt could not be found.
-    pub const SERVER_NOT_FOUND: u16 = 4011;
-
-    /// Discord did not recognise the voice protocol chosen.
-    pub const UNKNOWN_PROTOCOL: u16 = 4012;
-
-    /// Disconnected, either due to channel closure/removal
-    /// or kicking.
-    ///
-    /// Should not reconnect.
-    pub const DISCONNECTED: u16 = 4014;
-
-    /// Connected voice server crashed.
-    ///
-    /// Should resume.
-    pub const VOICE_SERVER_CRASH: u16 = 4015;
-
-    /// Discord didn't recognise the encrytpion scheme.
-    pub const UNKNOWN_ENCRYPTION_MODE: u16 = 4016;
-}

voice-model/src/lib.rs

@@ -1,6 +1,7 @@
 //! Mappings of objects received from Discord's voice gateway API, with implementations
 //! for (de)serialisation.
 
+mod close_code;
 pub mod constants;
 mod event;
 pub mod id;
@@ -11,8 +12,11 @@ mod speaking_state;
 mod util;
 
 pub use self::{
+    close_code::CloseCode,
     event::Event,
     opcode::OpCode,
     protocol_data::ProtocolData,
     speaking_state::SpeakingState,
 };
+
+pub use enum_primitive::FromPrimitive;