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

@@ -1,11 +1,15 @@
 mod context;
+mod resume;
 
 use crate::context::CONTEXT;
-use anyhow::Context;
-use std::{env, error::Error, pin::pin};
+use anyhow::Context as _;
+use std::{env, error::Error, pin::pin, time::Duration};
 use tokio::signal;
 use tokio_util::task::TaskTracker;
-use twilight_gateway::{CloseFrame, Config, Event, EventTypeFlags, Intents, Shard, StreamExt as _};
+use twilight_gateway::{
+    CloseFrame, ConfigBuilder, Event, EventTypeFlags, Intents, Shard, StreamExt as _,
+    queue::InMemoryQueue,
+};
 use twilight_http::Client;
 use twilight_model::gateway::payload::incoming::MessageCreate;
 
@@ -18,13 +22,24 @@ async fn main() -> anyhow::Result<()> {
 
     let token = env::var("TOKEN").context("reading `TOKEN`")?;
 
-    let config = Config::new(token.clone(), INTENTS);
-    let http = Client::new(token);
-    let shards = twilight_gateway::create_recommended(&http, config, |_, builder| builder.build())
+    let http = Client::new(token.clone());
+    let info = async { Ok::<_, anyhow::Error>(http.gateway().authed().await?.model().await?) }
         .await
-        .context("creating shards")?;
+        .context("getting info")?;
     context::initialize(http);
 
+    // The queue defaults are static and may be incorrect for large or newly
+    // restarted bots.
+    let queue = InMemoryQueue::new(
+        info.session_start_limit.max_concurrency,
+        info.session_start_limit.remaining,
+        Duration::from_millis(info.session_start_limit.reset_after),
+        info.session_start_limit.total,
+    );
+    let config = ConfigBuilder::new(token, INTENTS).queue(queue).build();
+
+    let shards = resume::restore(config, info.shards).await;
+
     let tasks = shards
         .into_iter()
         .map(|shard| tokio::spawn(dispatcher(shard)))
@@ -34,29 +49,35 @@ async fn main() -> anyhow::Result<()> {
     tracing::info!("shutting down; press CTRL-C to abort");
 
     let join_all_tasks = async {
+        let mut resume_info = Vec::new();
         for task in tasks {
-            task.await?;
+            resume_info.push(task.await?);
         }
-        Ok::<_, anyhow::Error>(())
+        Ok::<_, anyhow::Error>(resume_info)
     };
-    tokio::select! {
-        _ = signal::ctrl_c() => {},
-        _ = join_all_tasks => {},
+    let resume_info = tokio::select! {
+        _ = signal::ctrl_c() => Vec::new(),
+        resume_info = join_all_tasks => resume_info?,
     };
 
+    // Save shard information to be restored.
+    resume::save(&resume_info)
+        .await
+        .context("saving resume info")?;
+
     Ok(())
 }
 
 #[tracing::instrument(fields(shard = %shard.id(), skip_all))]
-async fn dispatcher(mut shard: Shard) {
+async fn dispatcher(mut shard: Shard) -> resume::Info {
     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) => {
@@ -86,6 +107,8 @@ async fn dispatcher(mut shard: Shard) {
 
     tracker.close();
     tracker.wait().await;
+
+    resume::Info::from(&shard)
 }
 
 #[tracing::instrument(fields(id = %event.id), skip_all)]

src/resume.rs

@@ -0,0 +1,81 @@
+use serde::{Deserialize, Serialize};
+use tokio::fs;
+use twilight_gateway::{Config, ConfigBuilder, Session, Shard, ShardId};
+
+const INFO_FILE: &str = "resume-info.json";
+
+/// [`Shard`] session resumption information.
+#[derive(Debug, Deserialize, Serialize)]
+pub struct Info {
+    resume_url: Option<String>,
+    session: Option<Session>,
+}
+
+impl Info {
+    fn is_none(&self) -> bool {
+        self.resume_url.is_none() && self.session.is_none()
+    }
+}
+
+impl From<&Shard> for Info {
+    fn from(value: &Shard) -> Self {
+        Self {
+            resume_url: value.resume_url().map(ToOwned::to_owned),
+            session: value.session().cloned(),
+        }
+    }
+}
+
+/// Saves shard resumption information to the file system.
+pub async fn save(info: &[Info]) -> anyhow::Result<()> {
+    if !info.iter().all(Info::is_none) {
+        let contents = serde_json::to_vec(&info)?;
+        fs::write(INFO_FILE, contents).await?;
+    }
+
+    Ok(())
+}
+
+/// Restores shard resumption information from the file system.
+pub async fn restore(config: Config, shards: u32) -> Vec<Shard> {
+    let resume_info = async {
+        let contents = fs::read(INFO_FILE).await?;
+        Ok::<_, anyhow::Error>(serde_json::from_slice::<Vec<Info>>(&contents)?)
+    }
+    .await;
+
+    let shard_ids = (0..shards).map(|shard| ShardId::new(shard, shards));
+
+    // A session may only successfully be resumed if it retains its shard ID,
+    // but Discord may have recommend a different shard count (producing
+    // different shard IDs).
+    let shards: Vec<_> = if let Ok(resume_info) = resume_info
+        && resume_info.len() == shards as usize
+    {
+        tracing::info!("resuming previous gateway sessions");
+        shard_ids
+            .zip(resume_info)
+            .map(|(shard_id, resume_info)| {
+                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()
+    };
+
+    // Resumed or not, the saved resume info is now stale.
+    _ = fs::remove_file(INFO_FILE).await;
+
+    shards
+}