eightball/PluralKit
1
0
Fork 0
mirror of https://github.com/PluralKit/PluralKit.git synced 2026-02-04 04:56:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

feat(gateway): add reconnect timestamp to shard state

Browse source
This commit is contained in:
asleepyskye 2025-08-27 22:31:23 -04:00
parent 2fc5f2a9d9
commit 1378379e14
4 changed files with 18 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

8
crates/gateway/src/discord/gateway.rs
View file

@ -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}");
}

3
crates/gateway/src/discord/shard_state.rs
View file

@ -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?;

11
crates/gateway/src/main.rs
View file

@ -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}")
};
}
}