gateway session resumption

Author: vilgotf

Cargo.toml

@@ -6,7 +6,9 @@ edition = "2024"
 [dependencies]
 anyhow = "1"
 rustls = "0.23"
-tokio = { version = "1", features = ["macros", "rt", "signal"] }
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+tokio = { version = "1", features = ["fs", "macros", "rt", "signal"] }
 tokio-util = { version = "0.7", features = ["rt"] }
 tracing = "0.1"
 tracing-subscriber = "0.3"

src/main.rs

@@ -2,15 +2,26 @@ mod context;
 
 use crate::context::CONTEXT;
 use anyhow::Context;
+use serde::{Deserialize, Serialize};
 use std::{env, error::Error, pin::pin};
-use tokio::signal;
+use tokio::{fs, signal};
 use tokio_util::task::TaskTracker;
-use twilight_gateway::{CloseFrame, Config, Event, EventTypeFlags, Intents, Shard, StreamExt as _};
+use twilight_gateway::{
+    CloseFrame, Config, ConfigBuilder, Event, EventTypeFlags, Intents, Session, Shard, ShardId,
+    StreamExt as _,
+};
 use twilight_http::Client;
 use twilight_model::gateway::payload::incoming::MessageCreate;
 
 const EVENT_TYPES: EventTypeFlags = EventTypeFlags::MESSAGE_CREATE;
 const INTENTS: Intents = Intents::GUILD_MESSAGES.union(Intents::MESSAGE_CONTENT);
+const RESUME_INFO_FILE: &str = "resume-info.json";
+
+#[derive(Debug, Deserialize, Serialize)]
+struct ResumeInfo {
+    resume_url: Option<String>,
+    session: Option<Session>,
+}
 
 #[tokio::main(flavor = "current_thread")]
 async fn main() -> anyhow::Result<()> {
@@ -20,31 +31,73 @@ async fn main() -> anyhow::Result<()> {
 
     let config = Config::new(token.clone(), INTENTS);
     let http = Client::new(token);
-    let shards = twilight_gateway::create_recommended(&http, config, |_, builder| builder.build())
+    let info = async { Ok::<_, anyhow::Error>(http.gateway().authed().await?.model().await?) }
         .await
-        .context("creating shards")?;
+        .context("getting info")?;
+    let resume_info = async {
+        let contents = fs::read(RESUME_INFO_FILE).await?;
+        Ok::<_, anyhow::Error>(serde_json::from_slice::<Vec<ResumeInfo>>(&contents)?)
+    }
+    .await
+    .ok();
+    let shard_ids = (0..info.shards).map(|shard| ShardId::new(shard, info.shards));
+    let shards: Vec<_> = if let Some(mut resume_info) = resume_info
+        && resume_info.len() == info.shards as usize
+    {
+        tracing::info!("resuming previous gateway sessions");
+        let mut resume_info_iter = resume_info.drain(..);
+        shard_ids
+            .map(|shard_id| {
+                let resume_info = resume_info_iter.next().unwrap();
+                let mut builder = ConfigBuilder::from(config.clone());
+
+                if let Some(resume_url) = resume_info.resume_url {
+                    builder = builder.resume_url(resume_url);
+                }
+                if let Some(session) = resume_info.session {
+                    builder = builder.session(session);
+                }
+
+                Shard::with_config(shard_id, builder.build())
+            })
+            .collect()
+    } else {
+        shard_ids
+            .map(|shard_id| Shard::with_config(shard_id, config.clone()))
+            .collect()
+    };
+    _ = fs::remove_file(RESUME_INFO_FILE).await;
     context::initialize(http);
 
-    let tracker = TaskTracker::new();
-    for shard in shards {
-        tracker.spawn(dispatcher(shard));
+    let tasks = shards
+        .into_iter()
+        .map(|shard| tokio::spawn(dispatcher(shard)))
+        .collect::<Vec<_>>();
+    let mut resume_info = Vec::new();
+    for task in tasks {
+        resume_info.push(task.await?);
     }
-    tracker.close();
-    tracker.wait().await;
+
+    async {
+        let contents = serde_json::to_vec(&resume_info)?;
+        fs::write(RESUME_INFO_FILE, contents).await
+    }
+    .await
+    .context("persisting resume info")?;
 
     Ok(())
 }
 
 #[tracing::instrument(fields(shard = %shard.id(), skip_all))]
-async fn dispatcher(mut shard: Shard) {
+async fn dispatcher(mut shard: Shard) -> ResumeInfo {
     let mut is_shutdown = false;
     let tracker = TaskTracker::new();
     let mut shutdown_fut = pin!(signal::ctrl_c());
 
     loop {
         tokio::select! {
             _ = &mut shutdown_fut, if !is_shutdown => {
-                shard.close(CloseFrame::NORMAL);
+                shard.close(CloseFrame::RESUME);
                 is_shutdown = true;
             },
             Some(item) = shard.next_event(EVENT_TYPES) => {
@@ -74,6 +127,11 @@ async fn dispatcher(mut shard: Shard) {
 
     tracker.close();
     tracker.wait().await;
+
+    ResumeInfo {
+        resume_url: shard.resume_url().map(ToOwned::to_owned),
+        session: shard.session().cloned(),
+    }
 }
 
 #[tracing::instrument(fields(id = %event.id), skip_all)]