|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import platform |
| 4 | +from typing import TYPE_CHECKING |
| 5 | + |
| 6 | +import discord |
| 7 | +from discord.ext import commands, tasks |
| 8 | + |
| 9 | +try: |
| 10 | + |
| 11 | + from prometheus_async.aio.web import start_http_server |
| 12 | + from prometheus_client import Counter, Enum, Gauge, Info, Summary |
| 13 | +except ImportError: |
| 14 | + raise RuntimeError( |
| 15 | + "Prometheus libraries are required to be installed. " |
| 16 | + "Either install those libraries or disable Prometheus extension" |
| 17 | + ) |
| 18 | + |
| 19 | +if TYPE_CHECKING: |
| 20 | + from bot.rodhaj import Rodhaj |
| 21 | + |
| 22 | +METRIC_PREFIX = "discord_" |
| 23 | + |
| 24 | + |
| 25 | +class FeatureCollector: |
| 26 | + __slots__ = ( |
| 27 | + "bot", |
| 28 | + "active_tickets", |
| 29 | + "closed_tickets", |
| 30 | + "locked_tickets", |
| 31 | + "blocked_users", |
| 32 | + ) |
| 33 | + |
| 34 | + def __init__(self, bot: Rodhaj): |
| 35 | + self.bot = bot |
| 36 | + self.active_tickets = Gauge( |
| 37 | + f"{METRIC_PREFIX}active_tickets", "Amount of active tickets" |
| 38 | + ) |
| 39 | + self.closed_tickets = Counter( |
| 40 | + f"{METRIC_PREFIX}closed_tickets", "Number of closed tickets in this session" |
| 41 | + ) |
| 42 | + self.locked_tickets = Gauge( |
| 43 | + f"{METRIC_PREFIX}locked_tickets", |
| 44 | + "Number of soft locked tickets in this session", |
| 45 | + ) |
| 46 | + self.blocked_users = Gauge( |
| 47 | + f"{METRIC_PREFIX}blocked_users", "Number of currently blocked users" |
| 48 | + ) |
| 49 | + |
| 50 | + |
| 51 | +# Maybe load all of these from an json file next time |
| 52 | +class Metrics: |
| 53 | + __slots__ = ("bot", "connected", "latency", "commands", "version", "features") |
| 54 | + |
| 55 | + def __init__(self, bot: Rodhaj): |
| 56 | + self.bot = bot |
| 57 | + self.connected = Enum( |
| 58 | + f"{METRIC_PREFIX}connected", |
| 59 | + "Connected to Discord", |
| 60 | + ["shard"], |
| 61 | + states=["connected", "disconnected"], |
| 62 | + ) |
| 63 | + self.latency = Gauge(f"{METRIC_PREFIX}latency", "Latency to Discord", ["shard"]) |
| 64 | + self.commands = Summary(f"{METRIC_PREFIX}commands", "Total commands executed") |
| 65 | + self.version = Info(f"{METRIC_PREFIX}version", "Versions of the bot") |
| 66 | + self.features = FeatureCollector(self.bot) |
| 67 | + |
| 68 | + def get_commands(self) -> int: |
| 69 | + total_commands = 0 |
| 70 | + for _ in self.bot.walk_commands(): |
| 71 | + # As some of the commands are parents, |
| 72 | + # Grouped commands are also counted here |
| 73 | + total_commands += 1 |
| 74 | + |
| 75 | + return total_commands |
| 76 | + |
| 77 | + def fill(self) -> None: |
| 78 | + self.version.info( |
| 79 | + { |
| 80 | + "build_version": self.bot.version, |
| 81 | + "dpy_version": discord.__version__, |
| 82 | + "python_version": platform.python_version(), |
| 83 | + } |
| 84 | + ) |
| 85 | + self.commands.observe(self.get_commands()) |
| 86 | + |
| 87 | + async def start(self, host: str, port: int) -> None: |
| 88 | + await start_http_server(addr=host, port=port) |
| 89 | + |
| 90 | + |
| 91 | +class Prometheus(commands.Cog): |
| 92 | + """Prometheus exporter extension for Rodhaj""" |
| 93 | + |
| 94 | + def __init__(self, bot: Rodhaj): |
| 95 | + self.bot = bot |
| 96 | + self._connected_label = self.bot.metrics.connected.labels(None) |
| 97 | + |
| 98 | + async def cog_load(self) -> None: |
| 99 | + self.latency_loop.start() |
| 100 | + |
| 101 | + async def cog_unload(self) -> None: |
| 102 | + self.latency_loop.stop() |
| 103 | + |
| 104 | + @tasks.loop(seconds=5) |
| 105 | + async def latency_loop(self) -> None: |
| 106 | + self.bot.metrics.latency.labels(None).set(self.bot.latency) |
| 107 | + |
| 108 | + @commands.Cog.listener() |
| 109 | + async def on_connect(self) -> None: |
| 110 | + self._connected_label.state("connected") |
| 111 | + |
| 112 | + @commands.Cog.listener() |
| 113 | + async def on_resumed(self) -> None: |
| 114 | + self._connected_label.state("connected") |
| 115 | + |
| 116 | + @commands.Cog.listener() |
| 117 | + async def on_disconnect(self) -> None: |
| 118 | + self._connected_label.state("disconnected") |
| 119 | + |
| 120 | + |
| 121 | +async def setup(bot: Rodhaj) -> None: |
| 122 | + await bot.add_cog(Prometheus(bot)) |
0 commit comments