From ad27025fda087274dcf27542eb7d844d5a409411 Mon Sep 17 00:00:00 2001
From: Nico Flaig <nflaig@protonmail.com>
Date: Wed, 3 Jan 2024 10:59:43 +0100
Subject: [PATCH] refactor: type safe metric labels

---
 packages/discv5/src/metrics.ts | 41 +++++++++++++++++-----------------
 1 file changed, 20 insertions(+), 21 deletions(-)

diff --git a/packages/discv5/src/metrics.ts b/packages/discv5/src/metrics.ts
index f3d82373..bd143142 100644
--- a/packages/discv5/src/metrics.ts
+++ b/packages/discv5/src/metrics.ts
@@ -1,33 +1,32 @@
 export interface MetricsRegister {
-  gauge<T extends string>(config: GaugeConfig<T>): IGauge<T>;
+  gauge<Labels extends LabelsGeneric = NoLabels>(config: GaugeConfig<Labels>): Gauge<Labels>;
 }
 
-type GaugeConfig<T extends string> = {
+type GaugeConfig<Labels extends LabelsGeneric> = {
   name: string;
   help: string;
-  labelNames?: T[] | readonly T[];
-};
+} & (NoLabels extends Labels ? { labelNames?: never } : { labelNames: [LabelKeys<Labels>, ...LabelKeys<Labels>[]] });
 
-type Labels<T extends string> = Partial<Record<T, string | number>>;
-interface IGauge<T extends string = string> {
-  inc(value?: number): void;
-  inc(labels: Labels<T>, value?: number): void;
-  set(value: number): void;
-  set(labels: Labels<T>, value: number): void;
-  collect(): void;
+type NoLabels = Record<string, never>;
+type LabelsGeneric = Record<string, string | number>;
+type LabelKeys<Labels extends LabelsGeneric> = Extract<keyof Labels, string>;
+interface Gauge<Labels extends LabelsGeneric = NoLabels> {
+  inc: NoLabels extends Labels ? (value?: number) => void : (labels: Labels, value?: number) => void;
+  set: NoLabels extends Labels ? (value: number) => void : (labels: Labels, value: number) => void;
+  collect?(): void;
 }
 
 export interface IDiscv5Metrics {
-  kadTableSize: IGauge;
-  activeSessionCount: IGauge;
-  connectedPeerCount: IGauge;
-  lookupCount: IGauge;
+  kadTableSize: Gauge;
+  activeSessionCount: Gauge;
+  connectedPeerCount: Gauge;
+  lookupCount: Gauge;
 
-  sentMessageCount: IGauge<"type">;
-  rcvdMessageCount: IGauge<"type">;
+  sentMessageCount: Gauge<{ type: string }>;
+  rcvdMessageCount: Gauge<{ type: string }>;
 
-  rateLimitHitIP: IGauge;
-  rateLimitHitTotal: IGauge;
+  rateLimitHitIP: Gauge;
+  rateLimitHitTotal: Gauge;
 }
 
 export type Metrics = ReturnType<typeof createDiscv5Metrics>;
@@ -59,13 +58,13 @@ export function createDiscv5Metrics(register: MetricsRegister) {
       help: "Count of the discv5 connected peers",
     }),
     /** Total number messages sent by message type */
-    sentMessageCount: register.gauge<"type">({
+    sentMessageCount: register.gauge<{ type: string }>({
       name: "discv5_sent_message_count",
       help: "Count of the discv5 messages sent by message type",
       labelNames: ["type"],
     }),
     /** Total number messages received by message type */
-    rcvdMessageCount: register.gauge<"type">({
+    rcvdMessageCount: register.gauge<{ type: string }>({
       name: "discv5_rcvd_message_count",
       help: "Count of the discv5 messages received by message type",
       labelNames: ["type"],