perf: reuse existing queries promises to avoid duplicate requests (#701)

In this example, the 'DESCRIBE' query should be fired twice only (for two different tables)
REF: HDX-1527

BEFORE:

<img width="1068" alt="Screenshot 2025-03-21 at 12 00 43 AM" src="https://github.com/user-attachments/assets/ee2d7613-1100-46a3-a5b0-7bfde0118dc0" />

AFTER:

<img width="1048" alt="Screenshot 2025-03-21 at 12 00 11 AM" src="https://github.com/user-attachments/assets/0899a884-1beb-4a18-8140-96c9728ca4f0" />

Author: wrn14897

.changeset/popular-ants-ring.md

@@ -0,0 +1,5 @@
+---
+"@hyperdx/common-utils": patch
+---
+
+perf: reuse existing queries promises to avoid duplicate requests

packages/common-utils/src/__tests__/metadata.test.ts

@@ -1,5 +1,5 @@
 import { ClickhouseClient } from '../clickhouse';
-import { Metadata } from '../metadata';
+import { Metadata, MetadataCache } from '../metadata';
 import * as renderChartConfigModule from '../renderChartConfig';
 import { ChartConfigWithDateRange } from '../types';
 
@@ -20,6 +20,109 @@ jest.mock('../renderChartConfig', () => ({
     .mockResolvedValue({ sql: 'SELECT 1', params: {} }),
 }));
 
+describe('MetadataCache', () => {
+  let metadataCache: MetadataCache;
+
+  beforeEach(() => {
+    metadataCache = new MetadataCache();
+    jest.clearAllMocks();
+  });
+
+  describe('getOrFetch', () => {
+    it('should return cached value if it exists', async () => {
+      const key = 'test-key';
+      const value = { data: 'test-data' };
+
+      // Set a value in the cache
+      metadataCache.set(key, value);
+
+      // Mock query function that should not be called
+      const queryFn = jest.fn().mockResolvedValue('new-value');
+
+      const result = await metadataCache.getOrFetch(key, queryFn);
+
+      expect(result).toBe(value);
+      expect(queryFn).not.toHaveBeenCalled();
+    });
+
+    it('should call query function and store result if no cached value exists', async () => {
+      const key = 'test-key';
+      const expectedValue = { data: 'fetched-data' };
+      const queryFn = jest.fn().mockResolvedValue(expectedValue);
+
+      const result = await metadataCache.getOrFetch(key, queryFn);
+
+      expect(result).toBe(expectedValue);
+      expect(queryFn).toHaveBeenCalledTimes(1);
+      expect(metadataCache.get(key)).toBe(expectedValue);
+    });
+
+    it('should reuse pending promises for the same key', async () => {
+      const key = 'test-key';
+      let resolvePromise: (value: any) => void;
+
+      // Create a promise that we can control when it resolves
+      const pendingPromise = new Promise(resolve => {
+        resolvePromise = resolve;
+      });
+
+      const queryFn = jest.fn().mockReturnValue(pendingPromise);
+
+      // Start two requests for the same key
+      const promise1 = metadataCache.getOrFetch(key, queryFn);
+      const promise2 = metadataCache.getOrFetch(key, queryFn);
+
+      // The query function should only be called once
+      expect(queryFn).toHaveBeenCalledTimes(1);
+
+      // Now resolve the promise
+      resolvePromise!({ data: 'result' });
+
+      // Both promises should resolve to the same value
+      const result1 = await promise1;
+      const result2 = await promise2;
+
+      expect(result1).toEqual({ data: 'result' });
+      expect(result2).toEqual({ data: 'result' });
+      expect(result1).toBe(result2); // Should be the same object reference
+    });
+
+    it('should clean up pending promise after resolution', async () => {
+      const key = 'test-key';
+      const value = { data: 'test-data' };
+      const queryFn = jest.fn().mockResolvedValue(value);
+
+      // Access the private pendingQueries map using any type assertion
+      const pendingQueriesMap = (metadataCache as any).pendingQueries;
+
+      await metadataCache.getOrFetch(key, queryFn);
+
+      // After resolution, the pending query should be removed from the map
+      expect(pendingQueriesMap.has(key)).toBe(false);
+    });
+
+    it('should clean up pending promise after rejection', async () => {
+      const key = 'test-key';
+      const error = new Error('Query failed');
+      const queryFn = jest.fn().mockRejectedValue(error);
+
+      // Access the private pendingQueries map using any type assertion
+      const pendingQueriesMap = (metadataCache as any).pendingQueries;
+
+      try {
+        await metadataCache.getOrFetch(key, queryFn);
+      } catch (e) {
+        // Expected to throw
+      }
+
+      // After rejection, the pending query should be removed from the map
+      expect(pendingQueriesMap.has(key)).toBe(false);
+      // And no value should be stored in the cache
+      expect(metadataCache.get(key)).toBeUndefined();
+    });
+  });
+});
+
 describe('Metadata', () => {
   let metadata: Metadata;
 

packages/common-utils/src/metadata.ts

@@ -15,22 +15,39 @@ const DEFAULT_SAMPLE_SIZE = 1e6;
 
 export class MetadataCache {
   private cache = new Map<string, any>();
+  private pendingQueries = new Map<string, Promise<any>>();
 
   // this should be getOrUpdate... or just query to follow react query
   get<T>(key: string): T | undefined {
     return this.cache.get(key);
   }
 
   async getOrFetch<T>(key: string, query: () => Promise<T>): Promise<T> {
-    const value = this.get(key) as T | undefined;
-    if (value != null) {
-      return value;
+    // Check if value exists in cache
+    const cachedValue = this.cache.get(key) as T | undefined;
+    if (cachedValue != null) {
+      return cachedValue;
     }
 
-    const newValue = await query();
-    this.cache.set(key, newValue);
+    // Check if there is a pending query
+    if (this.pendingQueries.has(key)) {
+      return this.pendingQueries.get(key)!;
+    }
+
+    // If no pending query, initiate the new query
+    const queryPromise = query();
 
-    return newValue;
+    // Store the pending query promise
+    this.pendingQueries.set(key, queryPromise);
+
+    try {
+      const result = await queryPromise;
+      this.cache.set(key, result);
+      return result;
+    } finally {
+      // Clean up the pending query map
+      this.pendingQueries.delete(key);
+    }
   }
 
   set<T>(key: string, value: T) {