From 44c5a2d1062e353b1fb73531155a53e8da0f011c Mon Sep 17 00:00:00 2001
From: alyssa <alyssa@mailc.net>
Date: Sun, 20 Apr 2025 19:25:11 +0000
Subject: [PATCH] feat(gateway): get node id from kubernetes

---
 crates/gateway/src/discord/gateway.rs | 7 +++++++
 crates/libpk/src/_config.rs           | 5 +++++
 2 files changed, 12 insertions(+)

diff --git a/crates/gateway/src/discord/gateway.rs b/crates/gateway/src/discord/gateway.rs
index 4c30e6a9..5ca24185 100644
--- a/crates/gateway/src/discord/gateway.rs
+++ b/crates/gateway/src/discord/gateway.rs
@@ -1,3 +1,4 @@
+use anyhow::anyhow;
 use futures::StreamExt;
 use libpk::{_config::ClusterSettings, runtime_config::RuntimeConfig};
 use metrics::counter;
@@ -48,6 +49,12 @@ pub fn create_shards(redis: fred::clients::RedisPool) -> anyhow::Result<Vec<Shar
 
     let (start_shard, end_shard): (u32, u32) = if cluster_settings.total_shards < 16 {
         warn!("we have less than 16 shards, assuming single gateway process");
+        if cluster_settings.node_id != 0 {
+            return Err(anyhow!(
+                "expecting to be node 0 in single-process mode, but we are node {}",
+                cluster_settings.node_id
+            ));
+        }
         (0, (cluster_settings.total_shards - 1).into())
     } else {
         (
diff --git a/crates/libpk/src/_config.rs b/crates/libpk/src/_config.rs
index c0f0b65a..5af40a94 100644
--- a/crates/libpk/src/_config.rs
+++ b/crates/libpk/src/_config.rs
@@ -138,10 +138,15 @@ impl PKConfig {
 lazy_static! {
     #[derive(Debug)]
     pub static ref CONFIG: Arc<PKConfig> = {
+        // hacks
         if let Ok(var) = std::env::var("NOMAD_ALLOC_INDEX")
             && std::env::var("pluralkit__discord__cluster__total_nodes").is_ok() {
             std::env::set_var("pluralkit__discord__cluster__node_id", var);
         }
+        if let Ok(var) = std::env::var("STATEFULSET_NAME_FOR_INDEX")
+            && std::env::var("pluralkit__discord__cluster__total_nodes").is_ok() {
+            std::env::set_var("pluralkit__discord__cluster__node_id", var.split("-").last().unwrap());
+        }
 
         Arc::new(Config::builder()
         .add_source(config::Environment::with_prefix("pluralkit").separator("__"))