@powersync/service-core 0.18.0 → 1.7.0

package/src/sync/sync.ts

@@ -1,5 +1,5 @@
 import { JSONBig, JsonContainer } from '@powersync/service-jsonbig';
-import { RequestParameters } from '@powersync/service-sync-rules';
+import { BucketDescription, BucketPriority, RequestParameters, SqlSyncRules } from '@powersync/service-sync-rules';
 import { Semaphore, withTimeout } from 'async-mutex';
 
 import { AbortError } from 'ix/aborterror.js';
@@ -9,9 +9,10 @@ import * as storage from '../storage/storage-index.js';
 import * as util from '../util/util-index.js';
 
 import { logger } from '@powersync/lib-services-framework';
+import { BucketChecksumState } from './BucketChecksumState.js';
 import { mergeAsyncIterables } from './merge.js';
 import { RequestTracker } from './RequestTracker.js';
-import { TokenStreamOptions, tokenStream } from './util.js';
+import { acquireSemaphoreAbortable, settledPromise, tokenStream, TokenStreamOptions } from './util.js';
 
 /**
  * Maximum number of connections actively fetching data.
@@ -32,11 +33,11 @@ const syncSemaphore = withTimeout(
 );
 
 export interface SyncStreamParameters {
-  storage: storage.BucketStorageFactory;
+  bucketStorage: storage.SyncRulesBucketStorage;
+  syncRules: SqlSyncRules;
   params: util.StreamingSyncRequest;
   syncParams: RequestParameters;
   token: auth.JwtPayload;
-  parseOptions: storage.ParseSyncRulesOptions;
   /**
    * If this signal is aborted, the stream response ends as soon as possible, without error.
    */
@@ -49,7 +50,7 @@ export interface SyncStreamParameters {
 export async function* streamResponse(
   options: SyncStreamParameters
 ): AsyncIterable<util.StreamingSyncLine | string | null> {
-  const { storage, params, syncParams, token, tokenStreamOptions, tracker, signal, parseOptions } = options;
+  const { bucketStorage, syncRules, params, syncParams, token, tokenStreamOptions, tracker, signal } = options;
   // We also need to be able to abort, so we create our own controller.
   const controller = new AbortController();
   if (signal) {
@@ -65,7 +66,7 @@ export async function* streamResponse(
     }
   }
   const ki = tokenStream(token, controller.signal, tokenStreamOptions);
-  const stream = streamResponseInner(storage, params, syncParams, tracker, parseOptions, controller.signal);
+  const stream = streamResponseInner(bucketStorage, syncRules, params, syncParams, tracker, controller.signal);
   // Merge the two streams, and abort as soon as one of the streams end.
   const merged = mergeAsyncIterables([stream, ki], controller.signal);
 
@@ -84,157 +85,155 @@ export async function* streamResponse(
   }
 }
 
+export type BucketSyncState = {
+  description?: BucketDescription; // Undefined if the bucket has not yet been resolved by us.
+  start_op_id: string;
+};
+
 async function* streamResponseInner(
-  storage: storage.BucketStorageFactory,
+  bucketStorage: storage.SyncRulesBucketStorage,
+  syncRules: SqlSyncRules,
   params: util.StreamingSyncRequest,
   syncParams: RequestParameters,
   tracker: RequestTracker,
-  parseOptions: storage.ParseSyncRulesOptions,
   signal: AbortSignal
 ): AsyncGenerator<util.StreamingSyncLine | string | null> {
-  // Bucket state of bucket id -> op_id.
-  // This starts with the state from the client. May contain buckets that the user do not have access to (anymore).
-  let dataBuckets = new Map<string, string>();
-
-  let lastChecksums: util.ChecksumMap | null = null;
-  let lastWriteCheckpoint: bigint | null = null;
-
   const { raw_data, binary_data } = params;
 
-  if (params.buckets) {
-    for (let { name, after: start } of params.buckets) {
-      dataBuckets.set(name, start);
-    }
-  }
-
   const checkpointUserId = util.checkpointUserId(syncParams.token_parameters.user_id as string, params.client_id);
-  const stream = storage.watchWriteCheckpoint(checkpointUserId, signal);
-  for await (const next of stream) {
-    const { base, writeCheckpoint } = next;
-    const checkpoint = base.checkpoint;
-
-    const storage = await base.getBucketStorage();
-    if (storage == null) {
-      // Sync rules deleted in the meantime - try again with the next checkpoint.
-      continue;
-    }
-    const syncRules = storage.getParsedSyncRules(parseOptions);
 
-    const allBuckets = await syncRules.queryBucketIds({
-      getParameterSets(lookups) {
-        return storage.getParameterSets(checkpoint, lookups);
-      },
-      parameters: syncParams
-    });
-
-    if (allBuckets.length > 1000) {
-      logger.error(`Too many buckets`, {
-        checkpoint,
-        user_id: syncParams.user_id,
-        buckets: allBuckets.length
-      });
-      // TODO: Limit number of buckets even before we get to this point
-      throw new Error(`Too many buckets: ${allBuckets.length}`);
-    }
+  const checksumState = new BucketChecksumState({
+    bucketStorage,
+    syncRules,
+    syncParams,
+    initialBucketPositions: params.buckets
+  });
+  const stream = bucketStorage.watchWriteCheckpoint({
+    user_id: checkpointUserId,
+    signal
+  });
+  const newCheckpoints = stream[Symbol.asyncIterator]();
 
-    let dataBucketsNew = new Map<string, string>();
-    for (let bucket of allBuckets) {
-      dataBucketsNew.set(bucket, dataBuckets.get(bucket) ?? '0');
-    }
-    dataBuckets = dataBucketsNew;
-
-    const bucketList = [...dataBuckets.keys()];
-    const checksumMap = await storage.getChecksums(checkpoint, bucketList);
-    // Subset of buckets for which there may be new data in this batch.
-    let bucketsToFetch: string[];
-
-    if (lastChecksums) {
-      const diff = util.checksumsDiff(lastChecksums, checksumMap);
-
-      if (
-        lastWriteCheckpoint == writeCheckpoint &&
-        diff.removedBuckets.length == 0 &&
-        diff.updatedBuckets.length == 0
-      ) {
-        // No changes - don't send anything to the client
+  try {
+    let nextCheckpointPromise:
+      | Promise<PromiseSettledResult<IteratorResult<storage.StorageCheckpointUpdate>>>
+      | undefined;
+
+    do {
+      if (!nextCheckpointPromise) {
+        // Wrap in a settledPromise, so that abort errors after the parent stopped iterating
+        // does not result in uncaught errors.
+        nextCheckpointPromise = settledPromise(newCheckpoints.next());
+      }
+      const next = await nextCheckpointPromise;
+      nextCheckpointPromise = undefined;
+      if (next.status == 'rejected') {
+        throw next.reason;
+      }
+      if (next.value.done) {
+        break;
+      }
+      const line = await checksumState.buildNextCheckpointLine(next.value.value);
+      if (line == null) {
+        // No update to sync
         continue;
       }
-      bucketsToFetch = diff.updatedBuckets.map((c) => c.bucket);
-
-      let message = `Updated checkpoint: ${checkpoint} | `;
-      message += `write: ${writeCheckpoint} | `;
-      message += `buckets: ${allBuckets.length} | `;
-      message += `updated: ${limitedBuckets(diff.updatedBuckets, 20)} | `;
-      message += `removed: ${limitedBuckets(diff.removedBuckets, 20)}`;
-      logger.info(message, {
-        checkpoint,
-        user_id: syncParams.user_id,
-        buckets: allBuckets.length,
-        updated: diff.updatedBuckets.length,
-        removed: diff.removedBuckets.length
-      });
 
-      const checksum_line: util.StreamingSyncCheckpointDiff = {
-        checkpoint_diff: {
-          last_op_id: checkpoint,
-          write_checkpoint: writeCheckpoint ? String(writeCheckpoint) : undefined,
-          removed_buckets: diff.removedBuckets,
-          updated_buckets: diff.updatedBuckets
+      const { checkpointLine, bucketsToFetch } = line;
+
+      yield checkpointLine;
+      // Start syncing data for buckets up to the checkpoint. As soon as we have completed at least one priority and
+      // at least 1000 operations, we also start listening for new checkpoints concurrently. When a new checkpoint comes
+      // in while we're still busy syncing data for lower priorities, interrupt the current operation and start syncing
+      // the new checkpoint.
+      const abortCheckpointController = new AbortController();
+      let syncedOperations = 0;
+
+      const abortCheckpointSignal = AbortSignal.any([abortCheckpointController.signal, signal]);
+
+      const bucketsByPriority = [...Map.groupBy(bucketsToFetch, (bucket) => bucket.priority).entries()];
+      bucketsByPriority.sort((a, b) => a[0] - b[0]); // Sort from high to lower priorities
+      const lowestPriority = bucketsByPriority.at(-1)?.[0];
+
+      function maybeRaceForNewCheckpoint() {
+        if (syncedOperations >= 1000 && nextCheckpointPromise === undefined) {
+          nextCheckpointPromise = (async () => {
+            const next = await settledPromise(newCheckpoints.next());
+            if (next.status == 'rejected') {
+              abortCheckpointController.abort();
+            } else if (!next.value.done) {
+              // Stop the running bucketDataInBatches() iterations, making the main flow reach the new checkpoint.
+              abortCheckpointController.abort();
+            }
+
+            return next;
+          })();
         }
-      };
+      }
 
-      yield checksum_line;
-    } else {
-      let message = `New checkpoint: ${checkpoint} | write: ${writeCheckpoint} | `;
-      message += `buckets: ${allBuckets.length} ${limitedBuckets(allBuckets, 20)}`;
-      logger.info(message, { checkpoint, user_id: syncParams.user_id, buckets: allBuckets.length });
-      bucketsToFetch = allBuckets;
-      const checksum_line: util.StreamingSyncCheckpoint = {
-        checkpoint: {
-          last_op_id: checkpoint,
-          write_checkpoint: writeCheckpoint ? String(writeCheckpoint) : undefined,
-          buckets: [...checksumMap.values()]
+      function markOperationsSent(operations: number) {
+        syncedOperations += operations;
+        tracker.addOperationsSynced(operations);
+        maybeRaceForNewCheckpoint();
+      }
+
+      // This incrementally updates dataBuckets with each individual bucket position.
+      // At the end of this, we can be sure that all buckets have data up to the checkpoint.
+      for (const [priority, buckets] of bucketsByPriority) {
+        const isLast = priority === lowestPriority;
+        if (abortCheckpointSignal.aborted) {
+          break;
         }
-      };
-      yield checksum_line;
-    }
-    lastChecksums = checksumMap;
-    lastWriteCheckpoint = writeCheckpoint;
-
-    // This incrementally updates dataBuckets with each individual bucket position.
-    // At the end of this, we can be sure that all buckets have data up to the checkpoint.
-    yield* bucketDataInBatches({
-      storage,
-      checkpoint,
-      bucketsToFetch,
-      dataBuckets,
-      raw_data,
-      binary_data,
-      signal,
-      tracker,
-      user_id: syncParams.user_id
-    });
-
-    await new Promise((resolve) => setTimeout(resolve, 10));
+
+        yield* bucketDataInBatches({
+          bucketStorage: bucketStorage,
+          checkpoint: next.value.value.base.checkpoint,
+          bucketsToFetch: buckets,
+          checksumState,
+          raw_data,
+          binary_data,
+          onRowsSent: markOperationsSent,
+          abort_connection: signal,
+          abort_batch: abortCheckpointSignal,
+          user_id: syncParams.user_id,
+          // Passing undefined will emit a full sync complete message at the end. If we pass a priority, we'll emit a partial
+          // sync complete message.
+          forPriority: !isLast ? priority : undefined
+        });
+      }
+
+      if (!abortCheckpointSignal.aborted) {
+        await new Promise((resolve) => setTimeout(resolve, 10));
+      }
+    } while (!signal.aborted);
+  } finally {
+    await newCheckpoints.return?.();
   }
 }
 
 interface BucketDataRequest {
-  storage: storage.SyncRulesBucketStorage;
+  bucketStorage: storage.SyncRulesBucketStorage;
   checkpoint: string;
-  bucketsToFetch: string[];
-  /** Bucket data position, modified by the request.  */
-  dataBuckets: Map<string, string>;
+  bucketsToFetch: BucketDescription[];
+  /** Contains current bucket state. Modified by the request as data is sent.  */
+  checksumState: BucketChecksumState;
   raw_data: boolean | undefined;
   binary_data: boolean | undefined;
-  tracker: RequestTracker;
-  signal: AbortSignal;
+  /** Signals that the connection was aborted and that streaming should stop ASAP. */
+  abort_connection: AbortSignal;
+  /**
+   * Signals that higher-priority batches are available. The current batch can stop at a sensible point.
+   * This signal also fires when abort_connection fires.
+   */
+  abort_batch: AbortSignal;
   user_id?: string;
+  forPriority?: BucketPriority;
+  onRowsSent: (amount: number) => void;
 }
 
 async function* bucketDataInBatches(request: BucketDataRequest) {
   let isDone = false;
-  while (!request.signal.aborted && !isDone) {
+  while (!request.abort_batch.aborted && !isDone) {
     // The code below is functionally the same as this for-await loop below.
     // However, the for-await loop appears to have a memory leak, so we avoid it.
     // for await (const { done, data } of bucketDataBatch(storage, checkpoint, dataBuckets, raw_data, signal)) {
@@ -273,7 +272,17 @@ interface BucketDataBatchResult {
  * Extracted as a separate internal function just to avoid memory leaks.
  */
 async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<BucketDataBatchResult, void> {
-  const { storage, checkpoint, bucketsToFetch, dataBuckets, raw_data, binary_data, tracker, signal } = request;
+  const {
+    bucketStorage: storage,
+    checkpoint,
+    bucketsToFetch,
+    checksumState,
+    raw_data,
+    binary_data,
+    abort_connection,
+    abort_batch,
+    onRowsSent
+  } = request;
 
   const checkpointOp = BigInt(checkpoint);
   let checkpointInvalidated = false;
@@ -281,7 +290,12 @@ async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<Buck
   if (syncSemaphore.isLocked()) {
     logger.info('Sync concurrency limit reached, waiting for lock', { user_id: request.user_id });
   }
-  const [value, release] = await syncSemaphore.acquire();
+  const acquired = await acquireSemaphoreAbortable(syncSemaphore, AbortSignal.any([abort_batch]));
+  if (acquired === 'aborted') {
+    return;
+  }
+
+  const [value, release] = acquired;
   try {
     if (value <= 3) {
       // This can be noisy, so we only log when we get close to the
@@ -293,13 +307,14 @@ async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<Buck
     }
     // Optimization: Only fetch buckets for which the checksums have changed since the last checkpoint
     // For the first batch, this will be all buckets.
-    const filteredBuckets = new Map(bucketsToFetch.map((bucket) => [bucket, dataBuckets.get(bucket)!]));
-    const data = storage.getBucketDataBatch(checkpoint, filteredBuckets);
+    const filteredBuckets = checksumState.getFilteredBucketPositions(bucketsToFetch);
+    const dataBatches = storage.getBucketDataBatch(checkpoint, filteredBuckets);
 
     let has_more = false;
 
-    for await (let { batch: r, targetOp } of data) {
-      if (signal.aborted) {
+    for await (let { batch: r, targetOp } of dataBatches) {
+      // Abort in current batch if the connection is closed
+      if (abort_connection.aborted) {
         return;
       }
       if (r.has_more) {
@@ -339,9 +354,15 @@ async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<Buck
         // iterator memory in case if large data sent.
         yield { data: null, done: false };
       }
-      tracker.addOperationsSynced(r.data.length);
+      onRowsSent(r.data.length);
+
+      checksumState.updateBucketPosition({ bucket: r.bucket, nextAfter: r.next_after, hasMore: r.has_more });
 
-      dataBuckets.set(r.bucket, r.next_after);
+      // Check if syncing bucket data is supposed to stop before fetching more data
+      // from storage.
+      if (abort_batch.aborted) {
+        return;
+      }
     }
 
     if (!has_more) {
@@ -351,12 +372,22 @@ async function* bucketDataBatch(request: BucketDataRequest): AsyncGenerator<Buck
         // More data should be available immediately for a new checkpoint.
         yield { data: null, done: true };
       } else {
-        const line: util.StreamingSyncCheckpointComplete = {
-          checkpoint_complete: {
-            last_op_id: checkpoint
-          }
-        };
-        yield { data: line, done: true };
+        if (request.forPriority !== undefined) {
+          const line: util.StreamingSyncCheckpointPartiallyComplete = {
+            partial_checkpoint_complete: {
+              last_op_id: checkpoint,
+              priority: request.forPriority
+            }
+          };
+          yield { data: line, done: true };
+        } else {
+          const line: util.StreamingSyncCheckpointComplete = {
+            checkpoint_complete: {
+              last_op_id: checkpoint
+            }
+          };
+          yield { data: line, done: true };
+        }
       }
     }
   } finally {
@@ -383,18 +414,3 @@ function transformLegacyResponse(bucketData: util.SyncBucketData): any {
     })
   };
 }
-
-function limitedBuckets(buckets: string[] | util.BucketChecksum[], limit: number) {
-  buckets = buckets.map((b) => {
-    if (typeof b != 'string') {
-      return b.bucket;
-    } else {
-      return b;
-    }
-  });
-  if (buckets.length <= limit) {
-    return JSON.stringify(buckets);
-  }
-  const limited = buckets.slice(0, limit);
-  return `${JSON.stringify(limited)}...`;
-}

package/src/sync/util.ts

@@ -2,6 +2,7 @@ import * as timers from 'timers/promises';
 
 import * as util from '../util/util-index.js';
 import { RequestTracker } from './RequestTracker.js';
+import { SemaphoreInterface } from 'async-mutex';
 
 export type TokenStreamOptions = {
   /**
@@ -99,3 +100,56 @@ export async function* transformToBytesTracked(
     yield encoded;
   }
 }
+
+export function acquireSemaphoreAbortable(
+  semaphone: SemaphoreInterface,
+  abort: AbortSignal
+): Promise<[number, SemaphoreInterface.Releaser] | 'aborted'> {
+  return new Promise((resolve, reject) => {
+    let aborted = false;
+    let hasSemaphore = false;
+
+    const listener = () => {
+      if (!hasSemaphore) {
+        aborted = true;
+        abort.removeEventListener('abort', listener);
+        resolve('aborted');
+      }
+    };
+    abort.addEventListener('abort', listener);
+
+    semaphone.acquire().then((acquired) => {
+      hasSemaphore = true;
+      if (aborted) {
+        // Release semaphore, already aborted
+        acquired[1]();
+      } else {
+        abort.removeEventListener('abort', listener);
+        resolve(acquired);
+      }
+    }, reject);
+  });
+}
+
+/**
+ * Wrap a promise in the style of Promise.allSettled.
+ *
+ * This is specifically useful if rejections should not be treated as uncaught rejections
+ * if it is not specifically handled.
+ */
+export function settledPromise<T>(promise: Promise<T>): Promise<PromiseSettledResult<T>> {
+  return promise.then(
+    (result) => {
+      return {
+        status: 'fulfilled',
+        value: result
+      };
+    },
+    (error) => {
+      return {
+        status: 'rejected',
+        reason: error
+      };
+    }
+  );
+}

package/src/util/checkpointing.ts

@@ -0,0 +1,41 @@
+import { ErrorCode, ServiceError } from '@powersync/lib-services-framework';
+import { RouteAPI } from '../api/RouteAPI.js';
+import { BucketStorageFactory } from '../storage/storage-index.js';
+
+export interface CreateWriteCheckpointOptions {
+  userId: string | undefined;
+  clientId: string | undefined;
+  api: RouteAPI;
+  storage: BucketStorageFactory;
+}
+export async function createWriteCheckpoint(options: CreateWriteCheckpointOptions) {
+  const full_user_id = checkpointUserId(options.userId, options.clientId);
+
+  const syncBucketStorage = await options.storage.getActiveStorage();
+  if (!syncBucketStorage) {
+    throw new ServiceError(ErrorCode.PSYNC_S2302, `Cannot create Write Checkpoint since no sync rules are active.`);
+  }
+
+  const { writeCheckpoint, currentCheckpoint } = await options.api.createReplicationHead(async (currentCheckpoint) => {
+    const writeCheckpoint = await syncBucketStorage.createManagedWriteCheckpoint({
+      user_id: full_user_id,
+      heads: { '1': currentCheckpoint }
+    });
+    return { writeCheckpoint, currentCheckpoint };
+  });
+
+  return {
+    writeCheckpoint: String(writeCheckpoint),
+    replicationHead: currentCheckpoint
+  };
+}
+
+export function checkpointUserId(user_id: string | undefined, client_id: string | undefined) {
+  if (user_id == null) {
+    throw new Error('user_id is required');
+  }
+  if (client_id == null) {
+    return user_id;
+  }
+  return `${user_id}/${client_id}`;
+}

package/src/util/protocol-types.ts

@@ -1,5 +1,5 @@
 import * as t from 'ts-codec';
-import { SqliteJsonValue } from '@powersync/service-sync-rules';
+import { BucketDescription, BucketPriority, SqliteJsonValue } from '@powersync/service-sync-rules';
 
 export const BucketRequest = t.object({
   name: t.string,
@@ -59,7 +59,7 @@ export interface StreamingSyncCheckpointDiff {
   checkpoint_diff: {
     last_op_id: OpId;
     write_checkpoint?: OpId;
-    updated_buckets: BucketChecksum[];
+    updated_buckets: BucketChecksumWithDescription[];
     removed_buckets: string[];
   };
 }
@@ -74,13 +74,23 @@ export interface StreamingSyncCheckpointComplete {
   };
 }
 
-export interface StreamingSyncKeepalive {}
+export interface StreamingSyncCheckpointPartiallyComplete {
+  partial_checkpoint_complete: {
+    last_op_id: OpId;
+    priority: BucketPriority;
+  };
+}
+
+export interface StreamingSyncKeepalive {
+  token_expires_in: number;
+}
 
 export type StreamingSyncLine =
   | StreamingSyncData
   | StreamingSyncCheckpoint
   | StreamingSyncCheckpointDiff
   | StreamingSyncCheckpointComplete
+  | StreamingSyncCheckpointPartiallyComplete
   | StreamingSyncKeepalive;
 
 /**
@@ -91,7 +101,7 @@ export type OpId = string;
 export interface Checkpoint {
   last_op_id: OpId;
   write_checkpoint?: OpId;
-  buckets: BucketChecksum[];
+  buckets: BucketChecksumWithDescription[];
 }
 
 export interface BucketState {
@@ -142,3 +152,5 @@ export interface BucketChecksum {
    */
   count: number;
 }
+
+export interface BucketChecksumWithDescription extends BucketChecksum, BucketDescription {}

package/src/util/util-index.ts

@@ -5,6 +5,7 @@ export * from './Mutex.js';
 export * from './protocol-types.js';
 export * from './secs.js';
 export * from './utils.js';
+export * from './checkpointing.js';
 
 export * from './config.js';
 export * from './config/compound-config-collector.js';

package/src/util/utils.ts

@@ -145,16 +145,6 @@ export function isCompleteRow(storeData: boolean, row: sync_rules.ToastableSqlit
   return !hasToastedValues(row);
 }
 
-export function checkpointUserId(user_id: string | undefined, client_id: string | undefined) {
-  if (user_id == null) {
-    throw new Error('user_id is required');
-  }
-  if (client_id == null) {
-    return user_id;
-  }
-  return `${user_id}/${client_id}`;
-}
-
 /**
  * Reduce a bucket to the final state as stored on the client.
  *

package/test/src/auth.test.ts

@@ -274,7 +274,7 @@ describe('JWT Auth', () => {
     ).rejects.toThrow('Token must expire in a maximum of');
   });
 
-  test('http', async () => {
+  test('http', { timeout: 20_000 }, async () => {
     // Not ideal to rely on an external endpoint for tests, but it is good to test that this
     // one actually works.
     const remote = new RemoteJWKSCollector(
@@ -290,9 +290,9 @@ describe('JWT Auth', () => {
         reject_ip_ranges: ['local']
       }
     });
-    expect(invalid.getKeys()).rejects.toThrow('IPs in this range are not supported');
+    await expect(invalid.getKeys()).rejects.toThrow('IPs in this range are not supported');
 
-    // IPS throw an error immediately
+    // IPs throw an error immediately
     expect(
       () =>
         new RemoteJWKSCollector('https://127.0.0.1/.well-known/jwks.json', {
@@ -315,11 +315,11 @@ describe('JWT Auth', () => {
 
     const invalid = new RemoteJWKSCollector('https://127.0.0.1/.well-known/jwks.json');
     // Should try and fetch
-    expect(invalid.getKeys()).rejects.toThrow();
+    await expect(invalid.getKeys()).rejects.toThrow();
 
     const invalid2 = new RemoteJWKSCollector('https://localhost/.well-known/jwks.json');
     // Should try and fetch
-    expect(invalid2.getKeys()).rejects.toThrow();
+    await expect(invalid2.getKeys()).rejects.toThrow();
   });
 
   test('caching', async () => {