diff --git a/crates/gateway/src/discord/gateway.rs b/crates/gateway/src/discord/gateway.rs index 6d16bfd1..2f657138 100644 --- a/crates/gateway/src/discord/gateway.rs +++ b/crates/gateway/src/discord/gateway.rs @@ -6,7 +6,7 @@ use std::sync::Arc; use tokio::sync::mpsc::Sender; use tracing::{error, info, warn}; use twilight_gateway::{ - create_iterator, ConfigBuilder, Event, EventTypeFlags, Message, Shard, ShardId, + create_iterator, ConfigBuilder, Event, EventTypeFlags, Message, Shard, ShardId, CloseFrame }; use twilight_model::gateway::{ payload::outgoing::update_presence::UpdatePresencePayload, @@ -116,7 +116,11 @@ pub async fn runner( let raw_event = match item { Ok(evt) => match evt { Message::Close(frame) => { + let mut state_event = ShardStateEvent::Closed; let close_code = if let Some(close) = frame { + if close == CloseFrame::RESUME { + state_event = ShardStateEvent::Reconnect; + } close.code.to_string() } else { "unknown".to_string() @@ -132,7 +136,7 @@ pub async fn runner( .increment(1); if let Err(error) = - tx_state.try_send((shard.id(), ShardStateEvent::Closed, None, None)) + tx_state.try_send((shard.id(), state_event, None, None)) { error!("failed to update shard state for socket closure: {error}"); } diff --git a/crates/gateway/src/discord/shard_state.rs b/crates/gateway/src/discord/shard_state.rs index c85e02c8..fe047b50 100644 --- a/crates/gateway/src/discord/shard_state.rs +++ b/crates/gateway/src/discord/shard_state.rs @@ -86,7 +86,7 @@ impl ShardStateManager { Ok(()) } - pub async fn socket_closed(&self, shard_id: u32) -> anyhow::Result<()> { + pub async fn socket_closed(&self, shard_id: u32, reconnect: bool) -> anyhow::Result<()> { gauge!("pluralkit_gateway_shard_up").decrement(1); let mut info = self @@ -97,6 +97,7 @@ impl ShardStateManager { info.shard_id = shard_id as i32; info.cluster_id = Some(cluster_config().node_id as i32); info.up = false; + info.last_reconnect = chrono::offset::Utc::now().timestamp() as i32; info.disconnection_count += 1; self.save_shard(shard_id, info).await?; diff --git a/crates/gateway/src/main.rs b/crates/gateway/src/main.rs index e61c3445..a4bf7d2a 100644 --- a/crates/gateway/src/main.rs +++ b/crates/gateway/src/main.rs @@ -109,8 +109,13 @@ async fn main() -> anyhow::Result<()> { }; } ShardStateEvent::Closed => { - if let Err(error) = shard_state.socket_closed(shard_id.number()).await { - error!("failed to update shard state for heartbeat: {error}") + if let Err(error) = shard_state.socket_closed(shard_id.number(), false).await { + error!("failed to update shard state for closed: {error}") + }; + } + ShardStateEvent::Reconnect => { + if let Err(error) = shard_state.socket_closed(shard_id.number(), true).await { + error!("failed to update shard state for reconnect: {error}") }; } ShardStateEvent::Other => { @@ -121,7 +126,7 @@ async fn main() -> anyhow::Result<()> { ) .await { - error!("failed to update shard state for heartbeat: {error}") + error!("failed to update shard state for other evt: {error}") }; } } diff --git a/crates/libpk/src/state.rs b/crates/libpk/src/state.rs index df44ea1d..3ea6342a 100644 --- a/crates/libpk/src/state.rs +++ b/crates/libpk/src/state.rs @@ -8,11 +8,13 @@ pub struct ShardState { /// unix timestamp pub last_heartbeat: i32, pub last_connection: i32, + pub last_reconnect: i32, pub cluster_id: Option, } pub enum ShardStateEvent { Closed, Heartbeat, + Reconnect, Other, }