From f138e88e5b3667966a5a1639df57da0c38d7ba10 Mon Sep 17 00:00:00 2001
From: gmega <giuliano.mega@gmail.com>
Date: Fri, 18 Oct 2024 18:39:09 -0300
Subject: [PATCH] add simple metrics

---
 deluge/core/alertmanager.py |  1 +
 deluge/core/core.py         |  5 +++
 deluge/core/daemon_entry.py |  3 ++
 deluge/core/metrics.py      | 63 +++++++++++++++++++++++++++++++++++++
 requirements.txt            |  1 +
 5 files changed, 73 insertions(+)
 create mode 100644 deluge/core/metrics.py

diff --git a/deluge/core/alertmanager.py b/deluge/core/alertmanager.py
index 24f47fddd..1323c6b37 100644
--- a/deluge/core/alertmanager.py
+++ b/deluge/core/alertmanager.py
@@ -53,6 +53,7 @@ class AlertManager(component.Component):
             | lt.alert.category_t.ip_block_notification
             | lt.alert.category_t.performance_warning
             | lt.alert.category_t.file_progress_notification
+            | lt.alert.category_t.piece_progress_notification
         )
 
         self.session.apply_settings({'alert_mask': alert_mask})
diff --git a/deluge/core/core.py b/deluge/core/core.py
index e2130f595..8db085b12 100644
--- a/deluge/core/core.py
+++ b/deluge/core/core.py
@@ -34,6 +34,7 @@ from deluge.core.authmanager import (
 )
 from deluge.core.eventmanager import EventManager
 from deluge.core.filtermanager import FilterManager
+from deluge.core.metrics import Metrics
 from deluge.core.pluginmanager import PluginManager
 from deluge.core.preferencesmanager import PreferencesManager
 from deluge.core.rpcserver import export
@@ -138,6 +139,7 @@ class Core(component.Component):
         self.torrentmanager = TorrentManager()
         self.filtermanager = FilterManager(self)
         self.authmanager = AuthManager()
+        self.metrics = Metrics(self)
 
         # New release check information
         self.new_release = None
@@ -200,6 +202,9 @@ class Core(component.Component):
         self.alertmanager.register_handler(
             'session_stats', self._on_alert_session_stats
         )
+        self.alertmanager.register_handler(
+            'piece_finished', self.metrics.handle_alert
+        )
         self.session_rates_timer_interval = 2
         self.session_rates_timer = task.LoopingCall(self._update_session_rates)
 
diff --git a/deluge/core/daemon_entry.py b/deluge/core/daemon_entry.py
index c49fd2a35..0180ae150 100644
--- a/deluge/core/daemon_entry.py
+++ b/deluge/core/daemon_entry.py
@@ -138,3 +138,6 @@ def start_daemon(skip_start=False):
     return run_profiled(
         run_daemon, options, output_file=options.profile, do_profile=options.profile
     )
+
+if __name__ == '__main__':
+    start_daemon()
\ No newline at end of file
diff --git a/deluge/core/metrics.py b/deluge/core/metrics.py
new file mode 100644
index 000000000..a09c25b46
--- /dev/null
+++ b/deluge/core/metrics.py
@@ -0,0 +1,63 @@
+import logging
+import re
+from csv import DictWriter
+from datetime import datetime, timezone
+from io import StringIO
+from typing import Optional
+
+from prometheus_client import Counter
+
+logger = logging.getLogger('M')
+
+torrent_pieces_downloaded = Counter(
+    name='deluge_torrent_pieces_downloaded',
+    documentation='Number of pieces downloaded for a torrent',
+    labelnames=['peer_id', 'torrent_name']
+)
+
+NID = re.compile(r'\(([a-zA-Z0-9]+)\)')
+
+
+class Metrics:
+    def __init__(self, core: 'deluge.core.Core'):
+        self.peer_id: Optional[str] = None
+        self._buffer = StringIO()
+        self.log_writer = DictWriter(
+            f=self._buffer,
+            fieldnames=['metric', 'timestamp', 'labels', 'value'],
+        )
+        self.log_writer.writeheader()
+        core.session.post_dht_stats()
+
+    def handle_alert(self, alert):
+        alert_type = alert.what()
+        handler = getattr(self, f'_{alert_type}', lambda _: None)
+        handler(alert)
+
+    def _dht_stats(self, alert):
+        # Since the node id is not exposed in the alert by libtorrent's Python binding,
+        # we need to extract the digest from the string representation.
+        result = NID.search(alert.message())
+        if result is None:
+            raise Exception('Could not extract node id from DHT stats alert')
+        self.peer_id = result.group(1)
+
+    def _piece_finished(self, alert):
+        labels = {
+            'peer_id': self.peer_id if self.peer_id else '<unknown>',
+            'torrent_name': alert.torrent_name,
+        }
+
+        torrent_pieces_downloaded.labels(**labels).inc(1)
+        self.log_writer.writerow(
+            dict({
+                'timestamp': datetime.now(timezone.utc).isoformat(),
+                'metric': 'deluge_torrent_pieces_downloaded',
+                'labels': ','.join(labels.values()),
+                'value': alert.piece_index,
+            })
+        )
+
+        logger.debug(self._buffer.getvalue())
+        self._buffer.truncate(0)
+        self._buffer.seek(0)
diff --git a/requirements.txt b/requirements.txt
index 4aee26d2d..f5b186317 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,3 +17,4 @@ zope.interface>=4.4.2
 distro; 'linux' in sys_platform or 'bsd' in sys_platform
 pygeoip
 ifaddr>=0.2.0
+prometheus-client>=0.21.0