Export postgres and redis query timing metrics as histograms

src/databases/Postgres.ts

@@ -5,6 +5,9 @@ import { Client, Pool, QueryResult, types } from "pg";
 import fs from "fs";
 import { CustomPostgresReadOnlyConfig, CustomWritePostgresConfig } from "../types/config.model";
 import { timeoutPomise, PromiseWithState, savePromiseState, nextFulfilment } from "../utils/promise";
+import { EmitMetricsOptions, Histogram } from "../utils/histogram";
+
+const HISTOGRAM_MAX_BUCKET = 4096; // in ms, the max value of the last histogram bucket, must be a power of 2
 
 // return numeric (pg_type oid=1700) as float
 types.setTypeParser(1700, function(val) {
@@ -51,6 +54,11 @@ export class Postgres implements IDatabase {
     failedResponseTime: number[] = [];
     maxStoredTimes = 200;
 
+    readSuccessHistogram: Histogram = new Histogram({ maxBucket: HISTOGRAM_MAX_BUCKET });
+    writeSuccessHistogram: Histogram = new Histogram({ maxBucket: HISTOGRAM_MAX_BUCKET });
+    readFailureHistogram: Histogram = new Histogram({ maxBucket: HISTOGRAM_MAX_BUCKET });
+    writeFailureHistogram: Histogram = new Histogram({ maxBucket: HISTOGRAM_MAX_BUCKET });
+
     constructor(private config: DatabaseConfig) {}
 
     async init(): Promise<void> {
@@ -145,7 +153,7 @@ export class Postgres implements IDatabase {
                 else if (this.config.postgres.timeout) currentPromises.push(savePromiseState(timeoutPomise(this.config.postgres.timeout)));
                 const queryResult = await nextFulfilment(currentPromises);
 
-                this.updateResponseTime(type, start);
+                this.updateResponseTime(type, start, false);
 
                 this.activePostgresRequests--;
                 switch (type) {
@@ -175,7 +183,7 @@ export class Postgres implements IDatabase {
                     }
                 }
 
-                this.updateResponseTime(type, start, this.failedResponseTime);
+                this.updateResponseTime(type, start, true);
                 this.activePostgresRequests--;
                 Logger.error(`prepare (postgres) try ${tries}: ${err}`);
             }
@@ -256,14 +264,41 @@ export class Postgres implements IDatabase {
         return result;
     }
 
-    private updateResponseTime(type: string, start: number, customArray?: number[]): void {
+    private updateResponseTime(type: string, start: number, isFailure: boolean): void {
         const responseTime = Date.now() - start;
 
-        const array = customArray ?? (this.isReadQuery(type) ?
-            this.readResponseTime : this.writeResponseTime);
+        // the "legacy" rolling average calculation
+        const array = isFailure
+            ? this.failedResponseTime
+            : (this.isReadQuery(type)
+                ? this.readResponseTime
+                : this.writeResponseTime);
 
         array.push(responseTime);
         if (array.length > this.maxStoredTimes) array.shift();
+
+        // histograms
+        const histogram = this.isReadQuery(type)
+            ? isFailure
+                ? this.readFailureHistogram
+                : this.readSuccessHistogram
+            : isFailure
+                ? this.writeFailureHistogram
+                : this.writeSuccessHistogram;
+
+        histogram.observe(responseTime);
+    }
+
+    emitHistograms({
+        baseName,
+        labels = {},
+    }: EmitMetricsOptions): string[] {
+        return [
+            ...this.readSuccessHistogram.emitMetrics({ baseName, labels: { ...labels, type: "read", failure: "no" } }),
+            ...this.writeSuccessHistogram.emitMetrics({ baseName, labels: { ...labels, type: "write", failure: "no" } }),
+            ...this.readFailureHistogram.emitMetrics({ baseName, labels: { ...labels, type: "read", failure: "yes" } }),
+            ...this.writeFailureHistogram.emitMetrics({ baseName, labels: { ...labels, type: "write", failure: "yes" } }),
+        ];
     }
 
     getStats(): PostgresStats {

src/routes/getMetrics.ts

@@ -1,12 +1,13 @@
 import { db, privateDB } from "../databases/databases";
 import { Request, Response } from "express";
 import os from "os";
-import redis, { getRedisStats } from "../utils/redis";
+import redis, { getRedisHistograms, getRedisStats } from "../utils/redis";
 import { Postgres } from "../databases/Postgres";
 import { Server } from "http";
 
 export async function getMetrics(req: Request, res: Response, server: Server): Promise<Response> {
     const redisStats = getRedisStats();
+    const redisHistograms = getRedisHistograms();
 
     return res.type("text").send([
         `# HELP sb_uptime Uptime of this instance`,
@@ -102,5 +103,13 @@ export async function getMetrics(req: Request, res: Response, server: Server): P
         `# HELP sb_redis_last_invalidation_message The time of the last invalidation message in redis`,
         `# TYPE sb_redis_last_invalidation_message gauge`,
         `sb_redis_last_invalidation_message ${redisStats.lastInvalidationMessage}`,
+        `# HELP sb_postgres_timings Query timings for postgres databases`,
+        `# TYPE sb_postgres_timings histogram`,
+        ...(db as Postgres)?.emitHistograms?.({ baseName: "sb_postgres_timings", labels: { db: "public" } }) ?? [],
+        ...(privateDB as Postgres)?.emitHistograms?.({ baseName: "sb_postgres_timings", labels: { db: "private" } }) ?? [],
+        `# HELP sb_redis_timings Query timings for redis`,
+        `# TYPE sb_redis_timings histogram`,
+        ...redisHistograms.read.emitMetrics({ baseName: "sb_redis_timings", labels: { type: "read" } }),
+        ...redisHistograms.write.emitMetrics({ baseName: "sb_redis_timings", labels: { type: "write" } }),
     ].join("\n"));
-}
+}

src/utils/histogram.ts

@@ -0,0 +1,125 @@
+export interface HistogramOptions {
+    /**
+     * The threshold for the last bucket, required
+     */
+    maxBucket: number;
+    /**
+     * The threshold for the first bucket, defaults to 1
+     *
+     * Must be >0
+     */
+    minBucket?: number;
+}
+
+export interface EmitMetricsOptions {
+    /**
+     * The basename of the histogram metric
+     */
+    baseName: string;
+    /**
+     * Extra labels to attach to the metrics
+     */
+    labels?: Record<string, string>;
+}
+
+export class Histogram {
+    private minBucket: number;
+    private storage: number[];
+    /**
+     * An object that maps bucket lower bounds to values in private histogram storage
+     *
+     * Keys are powers of 2 + "+Inf" for the "infinity bucket"/"total count submetric"
+     */
+    public readonly buckets: Record<string, number>;
+    private sum: number;
+
+    constructor({
+        maxBucket,
+        minBucket = 1,
+    }: HistogramOptions) {
+        this.minBucket = minBucket;
+        const numBuckets = Math.ceil(Math.log2(maxBucket / minBucket)) + 2; // log2 + <1ms bucket + <Infinity bucket
+
+        // initialize the storage array with zeros
+        this.storage = Array.from({ length: numBuckets }, () => 0);
+        this.sum = 0;
+
+        // generate the "buckets" access object
+        const buckets: Record<string, number> = {};
+
+        let max = minBucket;
+        for (let i = 0; i < numBuckets-1; i++) {
+            // rebind to a new const, otherwise all props will return the last bucket
+            const storageIndex = i;
+            Object.defineProperty(buckets, `${max}`, {
+                enumerable: true,
+                get: () => this.storage[storageIndex],
+            });
+            max *= 2;
+        }
+        Object.defineProperty(buckets, "+Inf", {
+            enumerable: true,
+            get: () => this.storage[numBuckets-1],
+        });
+
+        this.buckets = buckets;
+    }
+
+    public getSum(): number {
+        return this.sum;
+    }
+
+    public getCount(): number {
+        return this.buckets["+Inf"];
+    }
+
+    /**
+     * Log an observation
+     *
+     * Increments appropriate buckets and adds this value to the sum.
+     * @param observation the observed value
+     */
+    public observe(observation: number) {
+        this.sum += observation;
+
+        // find the appropriate bucket
+        let max = this.minBucket;
+        let smallestBucket;
+        for (smallestBucket = 0; smallestBucket < this.storage.length-1; smallestBucket++) {
+            if (observation <= max) {
+                break;
+            }
+            max *= 2;
+        }
+        // fun fact: if the break branch is not reached, i will be equal to this.storage.length-1 here
+        //           which corresponds to the +Inf bucket - everything works out
+        //
+        // increment every bucket after the smallest matching found
+        for (let i = smallestBucket; i < this.storage.length; i++) {
+            this.storage[i]++;
+        }
+    }
+
+    public emitMetrics({
+        baseName,
+        labels = {},
+    }: EmitMetricsOptions): string[] {
+        return [
+            `${baseName}_sum${serializeLabels(labels)} ${this.getSum()}`,
+            `${baseName}_count${serializeLabels(labels)} ${this.getCount()}`,
+            ...Object.entries(this.buckets).map(([bucket, value]) => `${baseName}_bucket${serializeLabels({ ...labels, le: bucket })} ${value}`),
+        ];
+    }
+}
+
+/**
+ * Serializes a set of labels into the format for prometheus.
+ *
+ * Makes no attempt to escape the names/values. Do not use quotation marks in names or values.
+ * @param labels map of label name -> value to serialize
+ */
+export function serializeLabels(labels: Record<string, string>): string {
+    const labelEntries = Object.entries(labels);
+    if (labelEntries.length === 0) return "";
+    return `{${labelEntries.map(([key, value]) => `${key}="${value}"`).join(", ")}}`;
+}

src/utils/redis.ts

@@ -10,6 +10,7 @@ import { compress, uncompress } from "lz4-napi";
 import { LRUCache } from "lru-cache";
 import { shouldClientCacheKey } from "./redisKeys";
 import { ZMember } from "@redis/client/dist/lib/commands/generic-transformers";
+import { Histogram } from "./histogram";
 
 export interface RedisStats {
     activeRequests: number;
@@ -24,6 +25,11 @@ export interface RedisStats {
     lastInvalidationMessage: number;
 }
 
+export interface RedisHistograms {
+    read: Histogram,
+    write: Histogram,
+}
+
 interface RedisSB {
     get(key: RedisCommandArgument, useClientCache?: boolean): Promise<string>;
     getWithCache(key: RedisCommandArgument): Promise<string>;
@@ -74,6 +80,9 @@ const writeResponseTime: number[] = [];
 let lastResponseTimeLimit = 0;
 const maxStoredTimes = 200;
 
+const readHistogram = new Histogram({ maxBucket: 4096 });
+const writeHistogram = new Histogram({ maxBucket: 4096 });
+
 const activeRequestPromises: Record<string, Promise<string>> = {};
 // Used to handle race conditions
 const resetKeys: Set<RedisCommandArgument> = new Set();
@@ -257,6 +266,7 @@ if (config.redis?.enabled) {
 
             const responseTime = Date.now() - start;
             readResponseTime.push(responseTime);
+            readHistogram.observe(responseTime);
             if (readResponseTime.length > maxStoredTimes) readResponseTime.shift();
             if (config.redis.stopWritingAfterResponseTime
                     && responseTime > config.redis.stopWritingAfterResponseTime) {
@@ -293,7 +303,9 @@ if (config.redis?.enabled) {
                 writeRequests--;
                 resolve(reply);
 
-                writeResponseTime.push(Date.now() - start);
+                const time = Date.now() - start;
+                writeResponseTime.push(time);
+                writeHistogram.observe(time);
                 if (writeResponseTime.length > maxStoredTimes) writeResponseTime.shift();
             }).catch((err) => {
                 activeRequests--;
@@ -412,6 +424,13 @@ export function getRedisStats(): RedisStats {
     };
 }
 
+export function getRedisHistograms(): RedisHistograms {
+    return {
+        read: readHistogram,
+        write: writeHistogram,
+    };
+}
+
 async function setupCacheClientListener(cacheClient: RedisClientType,
     cache: LRUCache<RedisCommandArgument, string>) {