From d262f5242f5385d41578af9a47de68e31b83d5ad Mon Sep 17 00:00:00 2001
From: Daniel LaCosse <3759828+daniellacosse@users.noreply.github.com>
Date: Fri, 21 Feb 2025 17:00:59 -0500
Subject: [PATCH] feat: cache recent prometheus queries in memory (#1643)

* [DO NOT MERGE] prometheus query timer

* add cache layer

* add query result data

* ready for a look

* remove logger
---
 src/shadowbox/server/manager_metrics.ts | 54 +++++++++++++++++++++----
 1 file changed, 46 insertions(+), 8 deletions(-)

diff --git a/src/shadowbox/server/manager_metrics.ts b/src/shadowbox/server/manager_metrics.ts
index 3969ac391..7cbc90ffb 100644
--- a/src/shadowbox/server/manager_metrics.ts
+++ b/src/shadowbox/server/manager_metrics.ts
@@ -16,6 +16,7 @@ import {
   PrometheusClient,
   PrometheusMetric,
   PrometheusValue,
+  QueryResultData,
 } from '../infrastructure/prometheus_scraper';
 import {DataUsageByUser, DataUsageTimeframe} from '../model/metrics';
 
@@ -111,6 +112,8 @@ export class PrometheusManagerMetrics implements ManagerMetrics {
       Math.ceil(now / PROMETHEUS_RANGE_QUERY_STEP_SECONDS) * PROMETHEUS_RANGE_QUERY_STEP_SECONDS;
     const start = end - timeframe.seconds;
 
+    this.prunePrometheusCache();
+
     const [
       bandwidth,
       bandwidthRange,
@@ -121,34 +124,34 @@ export class PrometheusManagerMetrics implements ManagerMetrics {
       dataTransferredByAccessKeyRange,
       tunnelTimeByAccessKeyRange,
     ] = await Promise.all([
-      this.prometheusClient.query(
+      this.cachedPrometheusClient.query(
         `sum(rate(shadowsocks_data_bytes_per_location{dir=~"c<p|p>t"}[${PROMETHEUS_RANGE_QUERY_STEP_SECONDS}s]))`
       ),
-      this.prometheusClient.queryRange(
+      this.cachedPrometheusClient.queryRange(
         `sum(rate(shadowsocks_data_bytes_per_location{dir=~"c<p|p>t"}[${PROMETHEUS_RANGE_QUERY_STEP_SECONDS}s]))`,
         start,
         end,
         `${PROMETHEUS_RANGE_QUERY_STEP_SECONDS}s`
       ),
-      this.prometheusClient.query(
+      this.cachedPrometheusClient.query(
         `sum(increase(shadowsocks_data_bytes_per_location{dir=~"c<p|p>t"}[${timeframe.seconds}s])) by (location, asn, asorg)`
       ),
-      this.prometheusClient.query(
+      this.cachedPrometheusClient.query(
         `sum(increase(shadowsocks_tunnel_time_seconds_per_location[${timeframe.seconds}s])) by (location, asn, asorg)`
       ),
-      this.prometheusClient.query(
+      this.cachedPrometheusClient.query(
         `sum(increase(shadowsocks_data_bytes{dir=~"c<p|p>t"}[${timeframe.seconds}s])) by (access_key)`
       ),
-      this.prometheusClient.query(
+      this.cachedPrometheusClient.query(
         `sum(increase(shadowsocks_tunnel_time_seconds[${timeframe.seconds}s])) by (access_key)`
       ),
-      this.prometheusClient.queryRange(
+      this.cachedPrometheusClient.queryRange(
         `sum(increase(shadowsocks_data_bytes{dir=~"c<p|p>t"}[${PROMETHEUS_RANGE_QUERY_STEP_SECONDS}s])) by (access_key)`,
         start,
         end,
         `${PROMETHEUS_RANGE_QUERY_STEP_SECONDS}s`
       ),
-      this.prometheusClient.queryRange(
+      this.cachedPrometheusClient.queryRange(
         `sum(increase(shadowsocks_tunnel_time_seconds[${PROMETHEUS_RANGE_QUERY_STEP_SECONDS}s])) by (access_key)`,
         start,
         end,
@@ -231,6 +234,41 @@ export class PrometheusManagerMetrics implements ManagerMetrics {
       accessKeys: Array.from(accessKeyMap.values()),
     };
   }
+
+  private prometheusCache = new Map<string, {timestamp: number; result: QueryResultData}>();
+
+  private get cachedPrometheusClient() {
+    return new Proxy(this.prometheusClient, {
+      get: (target, prop) => {
+        if (typeof target[prop] !== 'function') {
+          return target[prop];
+        }
+
+        return async (query, ...args) => {
+          const cacheId = `${String(prop)}: ${query} (args: ${args.join(', ')}))`;
+
+          if (this.prometheusCache.has(cacheId)) {
+            return this.prometheusCache.get(cacheId).result;
+          }
+
+          const result = await (target[prop] as Function)(query, ...args);
+
+          this.prometheusCache.set(cacheId, {timestamp: Date.now(), result});
+
+          return result;
+        };
+      },
+    });
+  }
+
+  private prunePrometheusCache() {
+    const now = Date.now();
+    for (const [key, value] of this.prometheusCache) {
+      if (now - value.timestamp > PROMETHEUS_RANGE_QUERY_STEP_SECONDS * 1000) {
+        this.prometheusCache.delete(key);
+      }
+    }
+  }
 }
 
 function getServerMetricsLocationEntry(