fix(core): dispose tenant cache on demand in order to hot reload oidc provider (#4641)

* fix(core): dispose tenant cache on demand

* chore: add comments for the tenant cache invalidation mechanism

* refactor(core): refactor dispose tenant cache implementation per review comments

* refactor(core): change `invalidateCache` to a class member method

* test(core): add test cases for Tenant class

packages/core/src/caches/well-known.test.ts

@@ -28,11 +28,14 @@ describe('Well-known cache basics', () => {
       pick(mockConnector0, 'connectorId', 'id', 'metadata'),
     ]);
 
+    await cache.set('tenant-cache-expires-at', WellKnownCache.defaultKey, 123);
+    expect(await cache.get('tenant-cache-expires-at', WellKnownCache.defaultKey)).toBe(123);
+
     await cache.delete('sie', WellKnownCache.defaultKey);
     expect(await cache.get('sie', WellKnownCache.defaultKey)).toBe(undefined);
   });
 
-  it('should be able to set the value with wrong structure', async () => {
+  it('should NOT be able to set the value with wrong structure', async () => {
     const cache = new WellKnownCache(tenantId, cacheStore);
 
     // @ts-expect-error
@@ -40,7 +43,7 @@ describe('Well-known cache basics', () => {
     expect(await cache.get('custom-phrases-tags', WellKnownCache.defaultKey)).toBe(undefined);
   });
 
-  it('should be able to set and get when cache type is wrong', async () => {
+  it('should NOT be able to set and get when cache type is wrong', async () => {
     const cache = new WellKnownCache(tenantId, cacheStore);
 
     // @ts-expect-error

packages/core/src/caches/well-known.ts

@@ -12,6 +12,7 @@ type WellKnownMap = {
   'connectors-well-known': ConnectorWellKnown[];
   'custom-phrases': Record<string, unknown>;
   'custom-phrases-tags': string[];
+  'tenant-cache-expires-at': number;
 };
 
 type WellKnownCacheType = keyof WellKnownMap;
@@ -52,6 +53,9 @@ function getValueGuard(type: WellKnownCacheType): ZodType<WellKnownMap[typeof ty
     case 'custom-phrases': {
       return z.record(z.unknown());
     }
+    case 'tenant-cache-expires-at': {
+      return z.number();
+    }
     default: {
       throw new Error(`No proper value guard found for cache key "${String(type)}".`);
     }

packages/core/src/routes/logto-config.ts

@@ -65,7 +65,7 @@ const getRedactedOidcKeyResponse = async (
   );
 
 export default function logtoConfigRoutes<T extends AuthedRouter>(
-  ...[router, { queries, logtoConfigs, envSet }]: RouterInitArgs<T>
+  ...[router, { queries, logtoConfigs, invalidateCache }]: RouterInitArgs<T>
 ) {
   const { getAdminConsoleConfig, updateAdminConsoleConfig, updateOidcConfigsByKey } =
     queries.logtoConfigs;
@@ -153,9 +153,7 @@ export default function logtoConfigRoutes<T extends AuthedRouter>(
       const updatedKeys = existingKeys.filter(({ id }) => id !== keyId);
 
       await updateOidcConfigsByKey(configKey, updatedKeys);
-
-      // Reload OIDC configs in envSet in order to apply the changes immediately
-      await envSet.load();
+      void invalidateCache();
 
       ctx.status = 204;
 
@@ -198,9 +196,7 @@ export default function logtoConfigRoutes<T extends AuthedRouter>(
       const updatedKeys = [newPrivateKey, ...existingKeys].slice(0, 2);
 
       await updateOidcConfigsByKey(configKey, updatedKeys);
-
-      // Reload OIDC configs in envSet in order to apply the changes immediately
-      await envSet.load();
+      void invalidateCache();
 
       // Remove actual values of the private keys from response
       ctx.body = await getRedactedOidcKeyResponse(configKey, updatedKeys);

packages/core/src/tenants/Tenant.test.ts

@@ -1,7 +1,9 @@
 import { adminTenantId, defaultTenantId } from '@logto/schemas';
 import { createMockUtils, pickDefault } from '@logto/shared/esm';
+import Sinon from 'sinon';
 
 import { RedisCache } from '#src/caches/index.js';
+import { WellKnownCache } from '#src/caches/well-known.js';
 import { createMockProvider } from '#src/test-utils/oidc-provider.js';
 import { emptyMiddleware } from '#src/utils/test-utils.js';
 
@@ -82,3 +84,31 @@ describe('Tenant `.run()`', () => {
     expect(typeof tenant.run).toBe('function');
   });
 });
+
+describe('Tenant cache health check', () => {
+  it('should set expiration timestamp in redis', async () => {
+    const redisCache = new RedisCache();
+    const tenant = await Tenant.create(defaultTenantId, redisCache);
+    expect(typeof tenant.invalidateCache).toBe('function');
+
+    Sinon.stub(tenant.wellKnownCache, 'set').value(jest.fn());
+
+    await tenant.invalidateCache();
+    expect(tenant.wellKnownCache.set).toBeCalledWith(
+      'tenant-cache-expires-at',
+      WellKnownCache.defaultKey,
+      expect.any(Number)
+    );
+  });
+
+  it('should be able to check the health of tenant cache', async () => {
+    const tenant = await Tenant.create(defaultTenantId, new RedisCache());
+    expect(typeof tenant.checkHealth).toBe('function');
+    expect(await tenant.checkHealth()).toBe(true);
+
+    // Stub the `wellKnownCache.get()` to set current timestamp as the tenant expiration timestamp
+    Sinon.stub(tenant.wellKnownCache, 'get').value(jest.fn(async () => Date.now()));
+
+    expect(await tenant.checkHealth()).toBe(false);
+  });
+});

packages/core/src/tenants/Tenant.ts

@@ -46,6 +46,7 @@ export default class Tenant implements TenantContext {
 
   private readonly app: Koa;
 
+  #createdAt = Date.now();
   #requestCount = 0;
   #onRequestEmpty?: () => Promise<void>;
 
@@ -95,6 +96,7 @@ export default class Tenant implements TenantContext {
       libraries,
       envSet,
       sentinel,
+      invalidateCache: this.invalidateCache.bind(this),
     };
 
     // Mount APIs
@@ -192,4 +194,33 @@ export default class Tenant implements TenantContext {
       };
     });
   }
+
+  /**
+   * Set a expiration timestamp in redis cache, and check it before returning the tenant LRU cache. This helps
+   * determine when to invalidate the cached tenant and force a in-place rolling reload of the OIDC provider.
+   */
+  public async invalidateCache() {
+    await this.wellKnownCache.set('tenant-cache-expires-at', WellKnownCache.defaultKey, Date.now());
+  }
+
+  /**
+   * Check if the tenant cache is healthy by comparing its creation timestamp with the global expiration timestamp.
+   *
+   * The global tenant expiration timestamp is stored in redis and shared across all server cluster instances. It
+   * can be set by calling `invalidateCache()` method on any tenant instance.
+   *
+   * @returns Resolves `true` if the tenant cache is healthy, `false` if it should be invalidated.
+   */
+  public async checkHealth() {
+    // `tenant-cache-expires-at` is a timestamp set in redis, which indicates all existing tenant instances in LRU
+    // cache should be invalidated after this timestamp, effective for the entire server cluster.
+    const tenantCacheExpiresAt = await this.wellKnownCache.get(
+      'tenant-cache-expires-at',
+      WellKnownCache.defaultKey
+    );
+
+    // Healthy if there's no expiration timestamp, or the current LRU cached tenant instance is created after the
+    // expiration timestamp.
+    return !tenantCacheExpiresAt || tenantCacheExpiresAt < this.#createdAt;
+  }
 }

packages/core/src/tenants/TenantContext.ts

@@ -19,4 +19,5 @@ export default abstract class TenantContext {
   public abstract readonly connectors: ConnectorLibrary;
   public abstract readonly libraries: Libraries;
   public abstract readonly sentinel: Sentinel;
+  public abstract invalidateCache(): Promise<void>;
 }

packages/core/src/tenants/index.ts

@@ -16,17 +16,22 @@ export class TenantPool {
   });
 
   async get(tenantId: string): Promise<Tenant> {
-    const tenant = this.cache.get(tenantId);
+    const tenantPromise = this.cache.get(tenantId);
 
-    if (tenant) {
-      return tenant;
+    if (tenantPromise) {
+      const tenant = await tenantPromise;
+      // If the current LRU cached tenant instance is still healthy, return it
+      if (await tenant.checkHealth()) {
+        return tenantPromise;
+      }
+      // Otherwise, create a new tenant instance and store in LRU cache, using the code below.
     }
 
     consoleLog.info('Init tenant:', tenantId);
-    const newTenant = Tenant.create(tenantId, redisCache);
-    this.cache.set(tenantId, newTenant);
+    const newTenantPromise = Tenant.create(tenantId, redisCache);
+    this.cache.set(tenantId, newTenantPromise);
 
-    return newTenant;
+    return newTenantPromise;
   }
 
   async endAll(): Promise<void> {

packages/core/src/test-utils/tenant.ts

@@ -89,6 +89,10 @@ export class MockTenant implements TenantContext {
     this.sentinel = new MockSentinel();
   }
 
+  public async invalidateCache() {
+    // Do nothing
+  }
+
   setPartialKey<Type extends 'queries' | 'libraries', Key extends keyof this[Type]>(
     type: Type,
     key: Key,

packages/integration-tests/src/tests/api/logto-config.test.ts

@@ -65,30 +65,32 @@ describe('admin console sign-in experience', () => {
   });
 
   it('should rotate OIDC keys successfully', async () => {
-    const privateKeys = await rotateOidcKeys('private-keys', SupportedSigningKeyAlgorithm.RSA);
+    const existingPrivateKeys = await getOidcKeys('private-keys');
+    const newPrivateKeys = await rotateOidcKeys('private-keys', SupportedSigningKeyAlgorithm.RSA);
 
-    expect(privateKeys).toHaveLength(2);
-    expect(privateKeys).toMatchObject([
+    expect(newPrivateKeys).toHaveLength(2);
+    expect(newPrivateKeys).toMatchObject([
       // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
       { id: expect.any(String), signingKeyAlgorithm: 'RSA', createdAt: expect.any(Number) },
       // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
       { id: expect.any(String), signingKeyAlgorithm: 'EC', createdAt: expect.any(Number) },
     ]);
+    expect(newPrivateKeys[1]?.id).toBe(existingPrivateKeys[0]?.id);
 
-    const cookieKeys = await rotateOidcKeys('cookie-keys');
+    const existingCookieKeys = await getOidcKeys('cookie-keys');
+    const newCookieKeys = await rotateOidcKeys('cookie-keys');
 
-    expect(cookieKeys).toHaveLength(2);
-    expect(cookieKeys).toMatchObject([
+    expect(newCookieKeys).toHaveLength(2);
+    expect(newCookieKeys).toMatchObject([
       // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
       { id: expect.any(String), createdAt: expect.any(Number) },
       // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
       { id: expect.any(String), createdAt: expect.any(Number) },
     ]);
+    expect(newCookieKeys[1]?.id).toBe(existingCookieKeys[0]?.id);
   });
 
   it('should only keep 2 recent OIDC keys', async () => {
-    await rotateOidcKeys('private-keys', SupportedSigningKeyAlgorithm.RSA);
-    await rotateOidcKeys('private-keys', SupportedSigningKeyAlgorithm.RSA);
     const privateKeys = await rotateOidcKeys('private-keys'); // Defaults to 'EC' algorithm
 
     expect(privateKeys).toHaveLength(2);
@@ -98,5 +100,16 @@ describe('admin console sign-in experience', () => {
       // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
       { id: expect.any(String), signingKeyAlgorithm: 'RSA', createdAt: expect.any(Number) },
     ]);
+
+    const privateKeys2 = await rotateOidcKeys('private-keys', SupportedSigningKeyAlgorithm.RSA);
+
+    expect(privateKeys2).toHaveLength(2);
+    expect(privateKeys2).toMatchObject([
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+      { id: expect.any(String), signingKeyAlgorithm: 'RSA', createdAt: expect.any(Number) },
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+      { id: expect.any(String), signingKeyAlgorithm: 'EC', createdAt: expect.any(Number) },
+    ]);
+    expect(privateKeys2[1]?.id).toBe(privateKeys[0]?.id);
   });
 });