Merge pull request #31106 from grantila/grantila/add-multiget-to-user-settings

Batch consecutive get calls to user-settings

plugins/user-settings/package.json

@@ -69,6 +69,7 @@
     "@material-ui/core": "^4.12.2",
     "@material-ui/icons": "^4.9.1",
     "@material-ui/lab": "4.0.0-alpha.61",
+    "dataloader": "^2.0.0",
     "react-use": "^17.2.4",
     "zen-observable": "^0.10.0"
   },

plugins/user-settings/src/apis/StorageApi/CacheMap.ts

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2025 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * A Map<K, V> with TTL (time-to-live) support, for the Dataloader of user settings.
+ */
+export class CacheMap<K, V> extends Map<K, V> {
+  #ttlMs: number;
+  #timestamps: Map<K, number> = new Map();
+
+  constructor(ttlMs: number) {
+    super();
+    this.#ttlMs = ttlMs;
+  }
+
+  set(key: K, value: V) {
+    const result = super.set(key, value);
+    this.#timestamps.set(key, Date.now());
+    return result;
+  }
+
+  get(key: K) {
+    if (!this.has(key)) {
+      return undefined;
+    }
+    const timestamp = this.#timestamps.get(key)!;
+    if (Date.now() - timestamp > this.#ttlMs) {
+      this.delete(key);
+      return undefined;
+    }
+    return super.get(key);
+  }
+
+  delete(key: K) {
+    this.#timestamps.delete(key);
+    return super.delete(key);
+  }
+
+  clear() {
+    this.#timestamps.clear();
+    return super.clear();
+  }
+}

plugins/user-settings/src/apis/StorageApi/UserSettingsStorage.test.ts

@@ -25,6 +25,7 @@ import {
   mockApis,
   registerMswTestHooks,
 } from '@backstage/test-utils';
+import { createDeferred } from '@backstage/types';
 import { rest } from 'msw';
 import { setupServer } from 'msw/node';
 import { UserSettingsStorage } from './UserSettingsStorage';
@@ -41,14 +42,12 @@ describe('Persistent Storage API', () => {
   const mockIdentityApi = mockApis.identity({ token: 'a-token' });
   const mockIdentityApiFallback = mockApis.identity();
 
-  const createPersistentStorage = (
-    args?: Partial<{
-      fetchApi: FetchApi;
-      discoveryApi: DiscoveryApi;
-      errorApi: ErrorApi;
-      namespace?: string;
-    }>,
-  ): StorageApi => {
+  const createPersistentStorage = (args: {
+    fetchApi?: FetchApi;
+    discoveryApi?: DiscoveryApi;
+    errorApi?: ErrorApi;
+    namespace: string;
+  }): StorageApi => {
     return UserSettingsStorage.create({
       errorApi: mockErrorApi,
       fetchApi: new MockFetchApi(),
@@ -58,14 +57,12 @@ describe('Persistent Storage API', () => {
     });
   };
 
-  const createPersistentStorageFallback = (
-    args?: Partial<{
-      fetchApi: FetchApi;
-      discoveryApi: DiscoveryApi;
-      errorApi: ErrorApi;
-      namespace?: string;
-    }>,
-  ): StorageApi => {
+  const createPersistentStorageFallback = (args: {
+    fetchApi?: FetchApi;
+    discoveryApi?: DiscoveryApi;
+    errorApi?: ErrorApi;
+    namespace: string;
+  }): StorageApi => {
     return UserSettingsStorage.create({
       errorApi: mockErrorApi,
       fetchApi: new MockFetchApi(),
@@ -75,21 +72,17 @@ describe('Persistent Storage API', () => {
     });
   };
 
-  afterEach(() => {
+  afterEach(async () => {
+    // Wait for server callbacks to settle before clearing the handlers.
+    // DataLoader delay is 10ms, so this should be plenty.
+    await new Promise(resolve => setTimeout(resolve, 100));
+
     jest.clearAllMocks();
+    server.resetHandlers();
   });
 
   it('should return undefined for values which are unset', async () => {
-    const storage = createPersistentStorage();
-
-    server.use(
-      rest.get(
-        `${mockBaseUrl}/buckets/:bucket/keys/:key`,
-        async (_req, res, ctx) => {
-          return res(ctx.json({ value: 'a' }));
-        },
-      ),
-    );
+    const storage = createPersistentStorage({ namespace: 'undefined' });
 
     expect(storage.snapshot('myfakekey').value).toBeUndefined();
     expect(storage.snapshot('myfakekey')).toEqual({
@@ -100,7 +93,7 @@ describe('Persistent Storage API', () => {
   });
 
   it('should allow setting of a simple data structure', async () => {
-    const storage = createPersistentStorage();
+    const storage = createPersistentStorage({ namespace: 'simple' });
     const dummyValue = 'a';
 
     server.use(
@@ -121,7 +114,7 @@ describe('Persistent Storage API', () => {
   });
 
   it('should allow setting of a complex data structure', async () => {
-    const storage = createPersistentStorage();
+    const storage = createPersistentStorage({ namespace: 'complex' });
     const dummyValue = {
       some: 'nice data',
       with: { nested: 'values', nice: true },
@@ -144,7 +137,9 @@ describe('Persistent Storage API', () => {
   });
 
   it('should fallback set when user not logged in', async () => {
-    const storage = createPersistentStorageFallback();
+    const storage = createPersistentStorageFallback({
+      namespace: 'not-logged-in',
+    });
 
     const selectedKeyNextHandler = jest.fn();
     const dummyValue = 'my-value';
@@ -171,12 +166,14 @@ describe('Persistent Storage API', () => {
   });
 
   it('should subscribe to key changes when setting a new value', async () => {
-    const storage = createPersistentStorage();
+    const storage = createPersistentStorage({ namespace: 'key-change-set' });
 
     const wrongKeyNextHandler = jest.fn();
     const selectedKeyNextHandler = jest.fn();
     const mockData = { hello: 'im a great new value' };
 
+    const serverCall = createDeferred();
+
     server.use(
       rest.put(
         `${mockBaseUrl}/buckets/:bucket/keys/:key`,
@@ -188,6 +185,10 @@ describe('Persistent Storage API', () => {
           return res(ctx.json(data));
         },
       ),
+      rest.post(`${mockBaseUrl}/multiget`, async (_req, res, ctx) => {
+        serverCall.resolve();
+        return res(ctx.json([]));
+      }),
     );
 
     await new Promise<void>(resolve => {
@@ -212,14 +213,18 @@ describe('Persistent Storage API', () => {
       presence: 'present',
       value: mockData,
     });
+
+    await serverCall;
   });
 
   it('should subscribe to key changes when deleting a value', async () => {
-    const storage = createPersistentStorage();
+    const storage = createPersistentStorage({ namespace: 'key-change-delete' });
 
     const wrongKeyNextHandler = jest.fn();
     const selectedKeyNextHandler = jest.fn();
 
+    const serverCall = createDeferred();
+
     server.use(
       rest.delete(
         `${mockBaseUrl}/buckets/:bucket/keys/:key`,
@@ -227,6 +232,10 @@ describe('Persistent Storage API', () => {
           return res(ctx.status(204));
         },
       ),
+      rest.post(`${mockBaseUrl}/multiget`, async (_req, res, ctx) => {
+        serverCall.resolve();
+        return res(ctx.json([]));
+      }),
     );
 
     await new Promise<void>(resolve => {
@@ -251,10 +260,12 @@ describe('Persistent Storage API', () => {
       presence: 'absent',
       value: undefined,
     });
+
+    await serverCall;
   });
 
   it('should not clash with other namespaces when creating buckets', async () => {
-    const rootStorage = createPersistentStorage();
+    const rootStorage = createPersistentStorage({ namespace: 'clash' });
     const selectedKeyNextHandler = jest.fn();
 
     server.use(
@@ -264,23 +275,21 @@ describe('Persistent Storage API', () => {
           const { bucket, key } = req.params;
           const { value } = await req.json();
 
-          expect(bucket).toEqual('default.profile.something.deep');
+          expect(bucket).toEqual('clash.profile.something.deep');
           expect(key).toEqual('test2');
 
           return res(ctx.json({ value }));
         },
       ),
-      rest.get(
-        `${mockBaseUrl}/buckets/:bucket/keys/:key`,
-        async (req, res, ctx) => {
-          const { bucket, key } = req.params;
+      rest.post(`${mockBaseUrl}/multiget`, async (req, res, ctx) => {
+        const payload = await req.json();
+        const { bucket, key } = payload.items[0];
 
-          expect(bucket).toEqual('default.profile/something');
-          expect(key).toEqual('deep/test2');
+        expect(bucket).toEqual('clash.profile/something');
+        expect(key).toEqual('deep/test2');
 
-          return res(ctx.status(404));
-        },
-      ),
+        return res(ctx.status(404));
+      }),
     );
 
     // when getting key test2 it will translate to default.profile.something.deep/test2
@@ -321,15 +330,13 @@ describe('Persistent Storage API', () => {
     });
 
     server.use(
-      rest.get(
-        `${mockBaseUrl}/buckets/:bucket/keys/:key`,
-        async (req, res, ctx) => {
-          const { bucket, key } = req.params;
-          expect(bucket).toEqual('Test.Mock.Thing');
-          expect(key).toEqual('key');
-          return res(ctx.text('{ invalid: json string }'));
-        },
-      ),
+      rest.post(`${mockBaseUrl}/multiget`, async (req, res, ctx) => {
+        const payload = await req.json();
+        expect(payload).toEqual({
+          items: [{ bucket: 'Test.Mock.Thing', key: 'key' }],
+        });
+        return res(ctx.text('{ invalid: json string }'));
+      }),
     );
 
     await new Promise<void>(resolve => {
@@ -353,17 +360,14 @@ describe('Persistent Storage API', () => {
   });
 
   it('should freeze the snapshot value', async () => {
-    const storage = createPersistentStorage();
+    const storage = createPersistentStorage({ namespace: 'freeze' });
     const selectedKeyNextHandler = jest.fn();
     const data = { foo: 'bar', baz: [{ foo: 'bar' }] };
 
     server.use(
-      rest.get(
-        `${mockBaseUrl}/buckets/:bucket/keys/:key`,
-        async (_req, res, ctx) => {
-          return res(ctx.json({ value: data }));
-        },
-      ),
+      rest.post(`${mockBaseUrl}/multiget`, async (_req, res, ctx) => {
+        return res(ctx.json({ items: [{ value: data }] }));
+      }),
     );
 
     await new Promise<void>(resolve => {
@@ -396,4 +400,112 @@ describe('Persistent Storage API', () => {
       snapshot.value.baz.push({ foo: 'buzz' });
     }).toThrow(/Cannot add property 1, object is not extensible/);
   });
+
+  it('should batch multiple calls into one', async () => {
+    const storage = createPersistentStorage({ namespace: 'multiget' });
+    const selectedKeyNextHandler = jest.fn();
+    const selectedKeyNextHandlerCached = jest.fn();
+    const data1 = { foo: 'bar1', baz: [{ foo: 'bar1' }] };
+    const data2 = { foo: 'bar2', baz: [{ foo: 'bar2' }] };
+
+    let serverCalls = 0;
+
+    server.use(
+      rest.post(`${mockBaseUrl}/multiget`, async (req, res, ctx) => {
+        ++serverCalls;
+        const payload = (await req.json()) as {
+          items: { bucket: string; key: string }[];
+        };
+        const result = payload.items.map(item => {
+          if (item.key === 'key1') {
+            return { value: data1 };
+          } else if (item.key === 'key2') {
+            return { value: data2 };
+          }
+          return null;
+        });
+
+        return res(ctx.json({ items: result }));
+      }),
+    );
+
+    await Promise.all([
+      new Promise<void>(resolve => {
+        storage.observe$('key1').subscribe({
+          next: snapshot => {
+            selectedKeyNextHandler(snapshot);
+            if (snapshot.presence === 'present') {
+              resolve();
+            }
+          },
+        });
+
+        storage.snapshot('key1');
+      }),
+
+      new Promise<void>(resolve => {
+        storage.observe$('missing-key').subscribe({
+          next: snapshot => {
+            selectedKeyNextHandler(snapshot);
+            if (snapshot.presence === 'absent') {
+              resolve();
+            }
+          },
+        });
+
+        storage.snapshot('missing-key');
+      }),
+
+      new Promise<void>(resolve => {
+        storage.observe$('key2').subscribe({
+          next: snapshot => {
+            selectedKeyNextHandler(snapshot);
+            if (snapshot.presence === 'present') {
+              resolve();
+            }
+          },
+        });
+
+        storage.snapshot('key2');
+      }),
+    ]);
+
+    expect(selectedKeyNextHandler).toHaveBeenCalledWith({
+      key: 'key1',
+      presence: 'present',
+      value: data1,
+    });
+    expect(selectedKeyNextHandler).toHaveBeenCalledWith({
+      key: 'key2',
+      presence: 'present',
+      value: data2,
+    });
+    expect(selectedKeyNextHandler).toHaveBeenCalledWith({
+      key: 'missing-key',
+      presence: 'absent',
+    });
+    expect(mockErrorApi.post).not.toHaveBeenCalled();
+    expect(serverCalls).toBe(1);
+
+    // Get key1 again, should use cached value
+    await new Promise<void>(resolve => {
+      storage.observe$('key1').subscribe({
+        next: snapshot => {
+          selectedKeyNextHandlerCached(snapshot);
+          if (snapshot.presence === 'present') {
+            resolve();
+          }
+        },
+      });
+
+      storage.snapshot('key1');
+    });
+
+    expect(selectedKeyNextHandlerCached).toHaveBeenCalledWith({
+      key: 'key1',
+      presence: 'present',
+      value: data1,
+    });
+    expect(serverCalls).toBe(1);
+  });
 });

plugins/user-settings/src/apis/StorageApi/UserSettingsStorage.ts

@@ -27,7 +27,12 @@ import { ResponseError } from '@backstage/errors';
 import { JsonValue, Observable } from '@backstage/types';
 import { SignalApi, SignalSubscriber } from '@backstage/plugin-signals-react';
 import ObservableImpl from 'zen-observable';
-import { UserSettingsSignal } from '@backstage/plugin-user-settings-common';
+import {
+  MultiGetResponse,
+  UserSettingsSignal,
+} from '@backstage/plugin-user-settings-common';
+import DataLoader from 'dataloader';
+import { CacheMap } from './CacheMap';
 
 const JSON_HEADERS = {
   'Content-Type': 'application/json; charset=utf-8',
@@ -36,6 +41,15 @@ const JSON_HEADERS = {
 
 const buckets = new Map<string, UserSettingsStorage>();
 
+const DATALOADER_CACHE_TTL_MS = 2 * 1000; // 2 seconds cache
+const DATALOADER_WINDOW_MS = 10; // 10 ms
+
+type DataLoaderType = DataLoader<
+  { bucket: string; key: string },
+  StorageValueSnapshot<JsonValue>,
+  string
+>;
+
 /**
  * An implementation of the storage API, that uses the user-settings backend to
  * persist the data in the DB.
@@ -59,6 +73,7 @@ export class UserSettingsStorage implements StorageApi {
   private readonly identityApi: IdentityApi;
   private readonly fallback: WebStorage;
   private readonly signalApi?: SignalApi;
+  private readonly userSettingsLoader: DataLoaderType;
 
   private constructor(
     namespace: string,
@@ -68,6 +83,7 @@ export class UserSettingsStorage implements StorageApi {
     identityApi: IdentityApi,
     fallback: WebStorage,
     signalApi?: SignalApi,
+    userSettingsLoader?: DataLoaderType,
   ) {
     this.namespace = namespace;
     this.fetchApi = fetchApi;
@@ -76,6 +92,35 @@ export class UserSettingsStorage implements StorageApi {
     this.identityApi = identityApi;
     this.fallback = fallback;
     this.signalApi = signalApi;
+
+    this.userSettingsLoader =
+      userSettingsLoader ??
+      new DataLoader(
+        async bucketAndKeyList => this.getMulti(bucketAndKeyList),
+        {
+          name: 'UserSettingsStorage.userSettingsLoader',
+          cacheMap: new CacheMap<
+            string,
+            Promise<StorageValueSnapshot<JsonValue>>
+          >(DATALOADER_CACHE_TTL_MS),
+          cacheKeyFn: bucketAndKey => this.stringifyDataLoaderKey(bucketAndKey),
+          maxBatchSize: 100,
+          batchScheduleFn: cb => setTimeout(cb, DATALOADER_WINDOW_MS),
+        },
+      );
+  }
+
+  private stringifyDataLoaderKey({
+    bucket,
+    key,
+  }: {
+    bucket: string;
+    key: string;
+  }) {
+    return `${encodeURIComponent(bucket)}/${encodeURIComponent(key)}`;
+  }
+  private clearCacheKey(key: string) {
+    this.userSettingsLoader.clear({ bucket: this.namespace, key });
   }
 
   static create(options: {
@@ -114,6 +159,8 @@ export class UserSettingsStorage implements StorageApi {
           this.errorApi,
           this.identityApi,
           this.fallback,
+          this.signalApi,
+          this.userSettingsLoader,
         ),
       );
     }
@@ -132,6 +179,8 @@ export class UserSettingsStorage implements StorageApi {
       throw await ResponseError.fromResponse(response);
     }
 
+    this.clearCacheKey(key);
+
     this.notifyChanges({ key, presence: 'absent' });
   }
 
@@ -139,9 +188,12 @@ export class UserSettingsStorage implements StorageApi {
     if (!(await this.isSignedIn())) {
       await this.fallback.set(key, data);
       this.notifyChanges({ key, presence: 'present', value: data });
+      this.clearCacheKey(key);
       return;
     }
 
+    this.clearCacheKey(key);
+
     const fetchUrl = await this.getFetchUrl(key);
 
     const response = await this.fetchApi.fetch(fetchUrl, {
@@ -156,6 +208,15 @@ export class UserSettingsStorage implements StorageApi {
 
     const { value } = await response.json();
 
+    this.userSettingsLoader.prime(
+      { bucket: this.namespace, key },
+      {
+        key,
+        presence: 'present',
+        value,
+      },
+    );
+
     this.notifyChanges({ key, value, presence: 'present' });
   }
 
@@ -171,7 +232,9 @@ export class UserSettingsStorage implements StorageApi {
 
           const updateSnapshot = () => {
             Promise.resolve()
-              .then(() => this.get(key))
+              .then(() =>
+                this.userSettingsLoader.load({ bucket: this.namespace, key }),
+              )
               .then(snapshot => subscriber.next(snapshot))
               .catch(error => this.errorApi.post(error));
           };
@@ -206,38 +269,62 @@ export class UserSettingsStorage implements StorageApi {
     return { key, presence: 'unknown' };
   }
 
-  private async get<T extends JsonValue>(
-    key: string,
-  ): Promise<StorageValueSnapshot<T>> {
+  private async getMulti(
+    bucketAndKeyList: readonly { bucket: string; key: string }[],
+  ): Promise<StorageValueSnapshot<JsonValue>[]> {
+    if (bucketAndKeyList.length === 0) return [];
+
     if (!(await this.isSignedIn())) {
       // This explicitly uses WebStorage, which we know is synchronous and doesn't return presence: unknown
-      return this.fallback.snapshot(key);
-    }
-
-    const fetchUrl = await this.getFetchUrl(key);
-    const response = await this.fetchApi.fetch(fetchUrl);
-
-    if (response.status === 404) {
-      return { key, presence: 'absent' };
-    }
-
-    if (!response.ok) {
-      throw await ResponseError.fromResponse(response);
+      return bucketAndKeyList.map(bucketAndKey =>
+        this.fallback.snapshot(bucketAndKey.key),
+      );
     }
 
     try {
-      const { value: rawValue } = await response.json();
-      const value = JSON.parse(JSON.stringify(rawValue), (_key, val) => {
-        if (typeof val === 'object' && val !== null) {
-          Object.freeze(val);
-        }
-        return val;
+      const baseUrl = await this.discoveryApi.getBaseUrl('user-settings');
+      const response = await this.fetchApi.fetch(`${baseUrl}/multiget`, {
+        method: 'POST',
+        headers: JSON_HEADERS,
+        body: JSON.stringify({ items: bucketAndKeyList }),
       });
 
-      return { key, presence: 'present', value };
-    } catch {
-      // If the value is not valid JSON, we return an unknown presence. This should never happen
-      return { key, presence: 'absent' };
+      if (response.status === 404) {
+        return bucketAndKeyList.map(bucketAndKey => ({
+          key: bucketAndKey.key,
+          presence: 'absent',
+        }));
+      }
+
+      if (!response.ok) {
+        throw await ResponseError.fromResponse(response);
+      }
+
+      const { items: values } = (await response.json()) as MultiGetResponse;
+
+      return bucketAndKeyList.map(
+        ({ key }, i): StorageValueSnapshot<JsonValue> => {
+          if (!values[i]) {
+            return { key, presence: 'absent' };
+          }
+          return {
+            key,
+            presence: 'present',
+            value: JSON.parse(JSON.stringify(values[i].value), (_key, val) => {
+              if (typeof val === 'object' && val !== null) {
+                Object.freeze(val);
+              }
+              return val;
+            }),
+          };
+        },
+      );
+    } catch (e) {
+      this.errorApi.post(new Error(`Failed to fetch user settings, ${e}`));
+      return bucketAndKeyList.map(bucketAndKey => ({
+        key: bucketAndKey.key,
+        presence: 'absent',
+      }));
     }
   }