I think it's there?

Jigsaw-Code · Nov 12, 2024 · d4957f4 · d4957f4
1 parent e06ab96
commit d4957f4
Show file tree

Hide file tree

Showing 3 changed files with 255 additions and 41 deletions.
diff --git a/src/shadowbox/server/manager_metrics.spec.ts b/src/shadowbox/server/manager_metrics.spec.ts
@@ -13,12 +13,133 @@
 // limitations under the License.
 
 import {PrometheusManagerMetrics} from './manager_metrics';
-import {FakePrometheusClient} from './mocks/mocks';
+import {FakeAccessKeyPrometheusClient} from './mocks/mocks';
 
 describe('PrometheusManagerMetrics', () => {
+  it('getServerMetrics', async (done) => {
+    const managerMetrics = new PrometheusManagerMetrics(
+      new FakeAccessKeyPrometheusClient([
+        {
+          accessKeyId: 0,
+          location: 'US',
+          asn: 49490,
+          asOrg: null,
+          dataTransferred: {
+            bytes: 50000,
+          },
+          tunnelTime: {
+            seconds: 10000,
+          },
+        },
+        {
+          accessKeyId: 1,
+          location: 'US',
+          asn: 49490,
+          asOrg: null,
+          dataTransferred: {
+            bytes: 50000,
+          },
+          tunnelTime: {
+            seconds: 5000,
+          },
+        },
+        {
+          accessKeyId: 2,
+          location: 'CA',
+          asn: null,
+          asOrg: null,
+          dataTransferred: {
+            bytes: 40000,
+          },
+          tunnelTime: {
+            seconds: 7500,
+          },
+        },
+      ])
+    );
+
+    const serverMetrics = await managerMetrics.getServerMetrics({hours: 0});
+
+    expect(JSON.stringify(serverMetrics, null, 2)).toEqual(`{
+  "server": [
+    {
+      "location": "US",
+      "asn": 49490,
+      "asOrg": "null",
+      "dataTransferred": {
+        "bytes": 100000
+      },
+      "tunnelTime": {
+        "seconds": 15000
+      }
+    },
+    {
+      "location": "CA",
+      "asn": null,
+      "asOrg": "null",
+      "dataTransferred": {
+        "bytes": 40000
+      },
+      "tunnelTime": {
+        "seconds": 7500
+      }
+    }
+  ],
+  "accessKeys": [
+    {
+      "accessKeyId": 0,
+      "dataTransferred": {
+        "bytes": 50000
+      },
+      "tunnelTime": {
+        "seconds": 10000
+      }
+    },
+    {
+      "accessKeyId": 1,
+      "dataTransferred": {
+        "bytes": 50000
+      },
+      "tunnelTime": {
+        "seconds": 5000
+      }
+    },
+    {
+      "accessKeyId": 2,
+      "dataTransferred": {
+        "bytes": 40000
+      },
+      "tunnelTime": {
+        "seconds": 7500
+      }
+    }
+  ]
+}`);
+    done();
+  });
+
   it('getOutboundByteTransfer', async (done) => {
     const managerMetrics = new PrometheusManagerMetrics(
-      new FakePrometheusClient({'access-key-1': 1000, 'access-key-2': 10000})
+      new FakeAccessKeyPrometheusClient([
+        {
+          accessKeyId: 'access-key-1',
+          asn: null,
+          asOrg: null,
+          location: null,
+          dataTransferred: {
+            bytes: 1000,
+          },
+        },
+        {
+          accessKeyId: 'access-key-2',
+          asn: null,
+          asOrg: null,
+          location: null,
+          dataTransferred: {
+            bytes: 10000,
+          },
+        },
+      ])
     );
     const dataUsage = await managerMetrics.getOutboundByteTransfer({hours: 0});
     const bytesTransferredByUserId = dataUsage.bytesTransferredByUserId;

diff --git a/src/shadowbox/server/manager_metrics.ts b/src/shadowbox/server/manager_metrics.ts
@@ -19,27 +19,31 @@ interface ServerMetricsTimeframe {
   hours: number;
 }
 
+interface ServerMetricsServerEntry {
+  location: string;
+  asn: number;
+  asOrg: string;
+  tunnelTime?: {
+    seconds: number;
+  };
+  dataTransferred: {
+    bytes: number;
+  };
+}
+
+interface ServerMetricsAccessKeyEntry {
+  accessKeyId: number;
+  tunnelTime?: {
+    seconds: number;
+  };
+  dataTransferred: {
+    bytes: number;
+  };
+}
+
 interface ServerMetrics {
-  server: {
-    location: string;
-    asn: number;
-    asOrg: string;
-    tunnelTime: {
-      seconds: number;
-    };
-    dataTransferred: {
-      bytes: number;
-    };
-  }[];
-  accessKeys: {
-    accessKeyId: number;
-    tunnelTime: {
-      seconds: number;
-    };
-    dataTransferred: {
-      bytes: number;
-    };
-  }[];
+  server: ServerMetricsServerEntry[];
+  accessKeys: ServerMetricsAccessKeyEntry[];
 }
 
 export interface ManagerMetrics {
@@ -71,7 +75,7 @@ export class PrometheusManagerMetrics implements ManagerMetrics {
 
   async getServerMetrics({hours}: ServerMetricsTimeframe): Promise<ServerMetrics> {
     const dataTransferredByLocation = await this.prometheusClient.query(
-      `sum(increase(shadowsocks_data_bytes{dir=~"c<p|p>t"}[${hours}h])) by (location, asn, asorg)`
+      `sum(increase(shadowsocks_data_bytes_per_location{dir=~"c<p|p>t"}[${hours}h])) by (location, asn, asorg)`
     );
     const tunnelTimeByLocation = await this.prometheusClient.query(
       `sum(increase(shadowsocks_tunnel_time_seconds_per_location[${hours}h])) by (location, asn, asorg)`
@@ -85,38 +89,52 @@ export class PrometheusManagerMetrics implements ManagerMetrics {
 
     const server = [];
     for (const entry of dataTransferredByLocation.result) {
-      server.push({
+      const result: ServerMetricsServerEntry = {
         location: entry.metric['location'],
         asn: parseInt(entry.metric['asn']),
         asOrg: entry.metric['asorg'],
-        tunnelTime: {
-          seconds: tunnelTimeByLocation.result.find((target) => {
-            return (
-              entry.metric['location'] === target.metric['location'] &&
-              entry.metric['asn'] === target.metric['asn'] &&
-              entry.metric['asorg'] === target.metric['asorg']
-            );
-          }),
-        },
         dataTransferred: {
-          bytes: Math.round(parseFloat(entry.value[1])),
+          bytes: parseFloat(entry.value[1]),
         },
+      };
+
+      const matchingTunnelTimeResult = tunnelTimeByLocation.result.find((target) => {
+        return (
+          entry.metric['location'] === target.metric['location'] &&
+          entry.metric['asn'] === target.metric['asn'] &&
+          entry.metric['asorg'] === target.metric['asorg']
+        );
       });
+
+      if (matchingTunnelTimeResult) {
+        result.tunnelTime = {
+          seconds: parseFloat(matchingTunnelTimeResult.value[1]),
+        };
+      }
+
+      server.push(result);
     }
 
     const accessKeys = [];
     for (const entry of dataTransferredByAccessKey.result) {
-      accessKeys.push({
+      const result: ServerMetricsAccessKeyEntry = {
         accessKeyId: parseInt(entry.metric['access_key']),
-        tunnelTime: {
-          seconds: tunnelTimeByAccessKey.result.find((target) => {
-            return entry.metric['access_key'] === target.metric['access_key'];
-          }),
-        },
         dataTransferred: {
-          bytes: Math.round(parseFloat(entry.value[1])),
+          bytes: parseFloat(entry.value[1]),
         },
+      };
+
+      const matchingTunnelTimeResult = tunnelTimeByAccessKey.result.find((target) => {
+        return entry.metric['access_key'] === target.metric['access_key'];
       });
+
+      if (matchingTunnelTimeResult) {
+        result.tunnelTime = {
+          seconds: parseFloat(matchingTunnelTimeResult.value[1]),
+        };
+      }
+
+      accessKeys.push(result);
     }
 
     return {

diff --git a/src/shadowbox/server/mocks/mocks.ts b/src/shadowbox/server/mocks/mocks.ts
@@ -65,3 +65,78 @@ export class FakePrometheusClient extends PrometheusClient {
     return queryResultData;
   }
 }
+
+interface AccessKeyPrometheusClientMetric {
+  accessKeyId: number | string;
+  location?: string;
+  asn?: number;
+  asOrg?: string;
+  tunnelTime?: {
+    seconds: number;
+  };
+  dataTransferred: {
+    bytes: number;
+  };
+}
+
+export class FakeAccessKeyPrometheusClient extends PrometheusClient {
+  constructor(public rawAccessKeyMetrics: AccessKeyPrometheusClientMetric[]) {
+    super('');
+  }
+
+  async query(_query: string): Promise<QueryResultData> {
+    const queryResultData = {result: []} as QueryResultData;
+
+    if (_query.startsWith('sum(increase(shadowsocks_data_bytes_per_location')) {
+      const locations = {};
+
+      for (const {location, asn, asOrg, dataTransferred} of this.rawAccessKeyMetrics) {
+        const locationKey = `${location},${asn},${asOrg}`;
+
+        locations[locationKey] ??= 0;
+        locations[locationKey] += dataTransferred.bytes;
+      }
+
+      for (const [locationKey, bytes] of Object.entries(locations)) {
+        const [location, asn, asorg] = locationKey.split(',');
+        queryResultData.result.push({
+          metric: {location, asn, asorg},
+          value: [Date.now() / 1000, `${bytes}`],
+        });
+      }
+    } else if (_query.startsWith('sum(increase(shadowsocks_tunnel_time_seconds_per_location')) {
+      const locations = {};
+
+      for (const {location, asn, asOrg, tunnelTime} of this.rawAccessKeyMetrics) {
+        const locationKey = `${location},${asn},${asOrg}`;
+
+        locations[locationKey] ??= 0;
+        locations[locationKey] += tunnelTime.seconds;
+      }
+
+      for (const [locationKey, seconds] of Object.entries(locations)) {
+        const [location, asn, asorg] = locationKey.split(',');
+        queryResultData.result.push({
+          metric: {location, asn, asorg},
+          value: [Date.now() / 1000, `${seconds}`],
+        });
+      }
+    } else if (_query.startsWith('sum(increase(shadowsocks_data_bytes')) {
+      for (const {accessKeyId, dataTransferred} of this.rawAccessKeyMetrics) {
+        queryResultData.result.push({
+          metric: {access_key: `${accessKeyId}`},
+          value: [Date.now() / 1000, `${dataTransferred.bytes}`],
+        });
+      }
+    } else if (_query.startsWith('sum(increase(shadowsocks_tunnel_time_seconds')) {
+      for (const {accessKeyId, tunnelTime} of this.rawAccessKeyMetrics) {
+        queryResultData.result.push({
+          metric: {access_key: `${accessKeyId}`},
+          value: [Date.now() / 1000, `${tunnelTime.seconds}`],
+        });
+      }
+    }
+
+    return queryResultData;
+  }
+}