From 1ec3663f63f084f66901c7460c54400b43d67398 Mon Sep 17 00:00:00 2001
From: spiral <spiral@spiral.sh>
Date: Mon, 7 Feb 2022 11:05:55 -0500
Subject: [PATCH] feat: add cluster ID to shard state

---
 PluralKit.API/Controllers/PrivateController.cs | 2 ++
 PluralKit.Bot/Services/ShardInfoService.cs     | 6 +++++-
 proto/discord.proto                            | 2 ++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/PluralKit.API/Controllers/PrivateController.cs b/PluralKit.API/Controllers/PrivateController.cs
index 6bbce231..3dba6531 100644
--- a/PluralKit.API/Controllers/PrivateController.cs
+++ b/PluralKit.API/Controllers/PrivateController.cs
@@ -61,6 +61,8 @@ public static class PrivateJsonExt
             s.Add("disconnection_count", shard.DisconnectionCount);
             s.Add("last_heartbeat", shard.LastHeartbeat.ToString());
             s.Add("last_connection", shard.LastConnection.ToString());
+            if (shard.HasClusterId)
+                s.Add("cluster_id", shard.ClusterId);
 
             o.Add(s);
         }
diff --git a/PluralKit.Bot/Services/ShardInfoService.cs b/PluralKit.Bot/Services/ShardInfoService.cs
index 8487ec97..41d1e47f 100644
--- a/PluralKit.Bot/Services/ShardInfoService.cs
+++ b/PluralKit.Bot/Services/ShardInfoService.cs
@@ -16,16 +16,18 @@ namespace PluralKit.Bot;
 
 public class ShardInfoService
 {
+    private readonly int? _clusterId;
     private readonly ILogger _logger;
     private readonly Cluster _client;
     private readonly RedisService _redis;
     private readonly Dictionary<int, ShardInfo> _shardInfo = new();
 
-    public ShardInfoService(ILogger logger, Cluster client, RedisService redis)
+    public ShardInfoService(ILogger logger, Cluster client, RedisService redis, BotConfig config)
     {
         _logger = logger.ForContext<ShardInfoService>();
         _client = client;
         _redis = redis;
+        _clusterId = config.Cluster?.NodeIndex;
     }
 
     public void Init()
@@ -65,6 +67,8 @@ public class ShardInfoService
 
             // latency = 0 because otherwise shard 0 would serialize to an empty array, thanks protobuf
             var state = new ShardState() { ShardId = shard.ShardId, Up = false, Latency = 1 };
+            if (_clusterId != null)
+                state.ClusterId = _clusterId.Value;
 
             // Register listeners for new shard
             shard.Resumed += () => ReadyOrResumed(shard);
diff --git a/proto/discord.proto b/proto/discord.proto
index f3bd765b..448d869b 100644
--- a/proto/discord.proto
+++ b/proto/discord.proto
@@ -11,6 +11,8 @@ message ShardState {
     // unix timestamp
     int32 last_heartbeat = 5;
     int32 last_connection = 6;
+
+    optional int32 cluster_id = 7;
 }
 
 message GatewayEvent {