Merge pull request #1210 from python-discord/joseph/add-socket-stats

Add a socket stats command

1 files changed, 37 insertions, 5 deletions

diff --git a/bot/exts/utils/eval.py b/bot/exts/utils/internal.py
index 6419b320e..1b4900f42 100644
--- a/bot/exts/utils/eval.py
+++ b/bot/exts/utils/internal.py
@@ -5,6 +5,8 @@ import pprint
 import re
 import textwrap
 import traceback
+from collections import Counter
+from datetime import datetime
 from io import StringIO
 from typing import Any, Optional, Tuple
 
@@ -19,8 +21,8 @@ from bot.utils import find_nth_occurrence, send_to_paste_service
 log = logging.getLogger(__name__)
 
 
-class CodeEval(Cog):
-    """Owner and admin feature that evaluates code and returns the result to the channel."""
+class Internal(Cog):
+    """Administrator and Core Developer commands."""
 
     def __init__(self, bot: Bot):
         self.bot = bot
@@ -30,6 +32,17 @@ class CodeEval(Cog):
 
         self.interpreter = Interpreter(bot)
 
+        self.socket_since = datetime.utcnow()
+        self.socket_event_total = 0
+        self.socket_events = Counter()
+
+    @Cog.listener()
+    async def on_socket_response(self, msg: dict) -> None:
+        """When a websocket event is received, increase our counters."""
+        if event_type := msg.get("t"):
+            self.socket_event_total += 1
+            self.socket_events[event_type] += 1
+
     def _format(self, inp: str, out: Any) -> Tuple[str, Optional[discord.Embed]]:
         """Format the eval output into a string & attempt to format it into an Embed."""
         self._ = out
@@ -198,7 +211,7 @@ async def func():  # (None,) -> Any
         await ctx.send(f"```py\n{out}```", embed=embed)
 
     @group(name='internal', aliases=('int',))
-    @has_any_role(Roles.owners, Roles.admins)
+    @has_any_role(Roles.owners, Roles.admins, Roles.core_developers)
     async def internal_group(self, ctx: Context) -> None:
         """Internal commands. Top secret!"""
         if not ctx.invoked_subcommand:
@@ -220,7 +233,26 @@ async def func():  # (None,) -> Any
 
         await self._eval(ctx, code)
 
+    @internal_group.command(name='socketstats', aliases=('socket', 'stats'))
+    @has_any_role(Roles.admins, Roles.owners, Roles.core_developers)
+    async def socketstats(self, ctx: Context) -> None:
+        """Fetch information on the socket events received from Discord."""
+        running_s = (datetime.utcnow() - self.socket_since).total_seconds()
+
+        per_s = self.socket_event_total / running_s
+
+        stats_embed = discord.Embed(
+            title="WebSocket statistics",
+            description=f"Receiving {per_s:0.2f} event per second.",
+            color=discord.Color.blurple()
+        )
+
+        for event_type, count in self.socket_events.most_common(25):
+            stats_embed.add_field(name=event_type, value=count, inline=False)
+
+        await ctx.send(embed=stats_embed)
+
 
 def setup(bot: Bot) -> None:
-    """Load the CodeEval cog."""
-    bot.add_cog(CodeEval(bot))
+    """Load the Internal cog."""
+    bot.add_cog(Internal(bot))