ready for a look

Author: daniellacosse

src/shadowbox/server/manager_metrics.ts

@@ -19,6 +19,7 @@ import {
   QueryResultData,
 } from '../infrastructure/prometheus_scraper';
 import {DataUsageByUser, DataUsageTimeframe} from '../model/metrics';
+import * as logging from '../infrastructure/logging';
 
 const PROMETHEUS_RANGE_QUERY_STEP_SECONDS = 5 * 60;
 
@@ -99,61 +100,6 @@ export class PrometheusManagerMetrics implements ManagerMetrics {
     return {bytesTransferredByUserId: usage};
   }
 
-  private queryCache = new Map<string, {timestamp: number; result: QueryResultData}>();
-
-  private async promethusClientTimedQuery(query: string) {
-    const cacheId = query;
-
-    if (this.queryCache.has(cacheId)) {
-      const cached = this.queryCache.get(cacheId);
-
-      if (cached && cached.timestamp + PROMETHEUS_RANGE_QUERY_STEP_SECONDS * 1000 > Date.now()) {
-        console.log(cacheId, 'cache hit');
-
-        return cached.result;
-      }
-    }
-
-    console.time(query);
-
-    const result = await this.prometheusClient.query(query);
-
-    this.queryCache.set(cacheId, {timestamp: Date.now(), result});
-
-    console.timeEnd(query);
-
-    return result;
-  }
-
-  private async promethusClientTimedRangeQuery(
-    query: string,
-    start: number,
-    end: number,
-    step: string
-  ) {
-    const cacheId = `${query}-${start}-${end}-${step}`;
-
-    if (this.queryCache.has(cacheId)) {
-      const cached = this.queryCache.get(cacheId);
-
-      if (cached && cached.timestamp + PROMETHEUS_RANGE_QUERY_STEP_SECONDS * 1000 > Date.now()) {
-        console.log(cacheId, 'cache hit');
-
-        return cached.result;
-      }
-    }
-
-    console.time(cacheId);
-
-    const result = await this.prometheusClient.queryRange(query, start, end, step);
-
-    this.queryCache.set(cacheId, {timestamp: Date.now(), result});
-
-    console.timeEnd(cacheId);
-
-    return result;
-  }
-
   async getServerMetrics(timeframe: Duration): Promise<ServerMetrics> {
     const now = new Date().getTime() / 1000;
     // We need to calculate consistent start and end times for Prometheus range
@@ -167,6 +113,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,
@@ -177,34 +125,34 @@ export class PrometheusManagerMetrics implements ManagerMetrics {
       dataTransferredByAccessKeyRange,
       tunnelTimeByAccessKeyRange,
     ] = await Promise.all([
-      this.promethusClientTimedQuery(
+      this.cachedPrometheusClient.query(
         `sum(rate(shadowsocks_data_bytes_per_location{dir=~"c<p|p>t"}[${PROMETHEUS_RANGE_QUERY_STEP_SECONDS}s]))`
       ),
-      this.promethusClientTimedRangeQuery(
+      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.promethusClientTimedQuery(
+      this.cachedPrometheusClient.query(
         `sum(increase(shadowsocks_data_bytes_per_location{dir=~"c<p|p>t"}[${timeframe.seconds}s])) by (location, asn, asorg)`
       ),
-      this.promethusClientTimedQuery(
+      this.cachedPrometheusClient.query(
         `sum(increase(shadowsocks_tunnel_time_seconds_per_location[${timeframe.seconds}s])) by (location, asn, asorg)`
       ),
-      this.promethusClientTimedQuery(
+      this.cachedPrometheusClient.query(
         `sum(increase(shadowsocks_data_bytes{dir=~"c<p|p>t"}[${timeframe.seconds}s])) by (access_key)`
       ),
-      this.promethusClientTimedQuery(
+      this.cachedPrometheusClient.query(
         `sum(increase(shadowsocks_tunnel_time_seconds[${timeframe.seconds}s])) by (access_key)`
       ),
-      this.promethusClientTimedRangeQuery(
+      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.promethusClientTimedRangeQuery(
+      this.cachedPrometheusClient.queryRange(
         `sum(increase(shadowsocks_tunnel_time_seconds[${PROMETHEUS_RANGE_QUERY_STEP_SECONDS}s])) by (access_key)`,
         start,
         end,
@@ -287,6 +235,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(