refactor(frontend): simplify UpdateService further

2 files changed, 173 insertions, 325 deletions

diff --git a/subprojects/frontend/src/xtext/UpdateService.ts b/subprojects/frontend/src/xtext/UpdateService.ts
index 94e01ca2..d8782d90 100644
--- a/subprojects/frontend/src/xtext/UpdateService.ts
+++ b/subprojects/frontend/src/xtext/UpdateService.ts
@@ -6,11 +6,7 @@ import { nanoid } from 'nanoid';
 import type EditorStore from '../editor/EditorStore';
 import getLogger from '../utils/getLogger';
 
-import UpdateStateTracker, {
-  type LockedState,
-  type PendingUpdate,
-} from './UpdateStateTracker';
-import type { StateUpdateResult, Delta } from './UpdateStateTracker';
+import UpdateStateTracker from './UpdateStateTracker';
 import type XtextWebSocketClient from './XtextWebSocketClient';
 import {
   type ContentAssistEntry,
@@ -86,10 +82,10 @@ export default class UpdateService {
   }
 
   private idleUpdate(): void {
-    if (!this.webSocketClient.isOpen || !this.tracker.hasDirtyChanges) {
+    if (!this.webSocketClient.isOpen || !this.tracker.needsUpdate) {
       return;
     }
-    if (!this.tracker.locekdForUpdate) {
+    if (!this.tracker.lockedForUpdate) {
       this.updateOrThrow().catch((error) => {
         if (error === E_CANCELED || error === E_TIMEOUT) {
           log.debug('Idle update cancelled');
@@ -111,88 +107,64 @@ export default class UpdateService {
    * @returns a promise resolving when the update is completed
    */
   private async updateOrThrow(): Promise<void> {
-    // We may check here for the delta to avoid locking,
-    // but we'll need to recompute the delta in the critical section,
-    // because it may have changed by the time we can acquire the lock.
-    if (
-      !this.tracker.hasDirtyChanges &&
-      this.tracker.xtextStateId !== undefined
-    ) {
+    if (!this.tracker.needsUpdate) {
       return;
     }
-    await this.tracker.runExclusive((lockedState) =>
-      this.updateExclusive(lockedState),
-    );
+    await this.tracker.runExclusive(() => this.updateExclusive());
   }
 
-  private async updateExclusive(lockedState: LockedState): Promise<void> {
+  private async updateExclusive(): Promise<void> {
     if (this.xtextStateId === undefined) {
-      await this.updateFullTextExclusive(lockedState);
+      await this.updateFullTextExclusive();
     }
-    if (!this.tracker.hasDirtyChanges) {
+    const delta = this.tracker.prepareDeltaUpdateExclusive();
+    if (delta === undefined) {
       return;
     }
-    await lockedState.updateExclusive(async (pendingUpdate) => {
-      const delta = pendingUpdate.prepareDeltaUpdateExclusive();
-      if (delta === undefined) {
-        return undefined;
-      }
-      log.trace('Editor delta', delta);
-      const result = await this.webSocketClient.send({
-        resource: this.resourceName,
-        serviceType: 'update',
-        requiredStateId: this.xtextStateId,
-        ...delta,
-      });
-      const parsedDocumentStateResult = DocumentStateResult.safeParse(result);
-      if (parsedDocumentStateResult.success) {
-        return parsedDocumentStateResult.data.stateId;
-      }
-      if (isConflictResult(result, 'invalidStateId')) {
-        return this.doUpdateFullTextExclusive(pendingUpdate);
-      }
-      throw parsedDocumentStateResult.error;
+    log.trace('Editor delta', delta);
+    const result = await this.webSocketClient.send({
+      resource: this.resourceName,
+      serviceType: 'update',
+      requiredStateId: this.xtextStateId,
+      ...delta,
     });
+    const parsedDocumentStateResult = DocumentStateResult.safeParse(result);
+    if (parsedDocumentStateResult.success) {
+      this.tracker.setStateIdExclusive(parsedDocumentStateResult.data.stateId);
+      return;
+    }
+    if (isConflictResult(result, 'invalidStateId')) {
+      await this.updateFullTextExclusive();
+    }
+    throw parsedDocumentStateResult.error;
   }
 
   private updateFullTextOrThrow(): Promise<void> {
-    return this.tracker.runExclusive((lockedState) =>
-      this.updateFullTextExclusive(lockedState),
-    );
+    return this.tracker.runExclusive(() => this.updateFullTextExclusive());
   }
 
-  private async updateFullTextExclusive(
-    lockedState: LockedState,
-  ): Promise<void> {
-    await lockedState.updateExclusive((pendingUpdate) =>
-      this.doUpdateFullTextExclusive(pendingUpdate),
-    );
-  }
-
-  private async doUpdateFullTextExclusive(
-    pendingUpdate: PendingUpdate,
-  ): Promise<string> {
+  private async updateFullTextExclusive(): Promise<void> {
     log.debug('Performing full text update');
-    pendingUpdate.extendPendingUpdateExclusive();
+    this.tracker.prepareFullTextUpdateExclusive();
     const result = await this.webSocketClient.send({
       resource: this.resourceName,
       serviceType: 'update',
       fullText: this.store.state.doc.sliceString(0),
     });
     const { stateId } = DocumentStateResult.parse(result);
-    return stateId;
+    this.tracker.setStateIdExclusive(stateId);
   }
 
   async fetchContentAssist(
     params: ContentAssistParams,
     signal: AbortSignal,
   ): Promise<ContentAssistEntry[]> {
-    if (this.tracker.upToDate && this.xtextStateId !== undefined) {
+    if (!this.tracker.hasPendingChanges && this.xtextStateId !== undefined) {
       return this.fetchContentAssistFetchOnly(params, this.xtextStateId);
     }
     try {
-      return await this.tracker.runExclusiveHighPriority((lockedState) =>
-        this.fetchContentAssistExclusive(params, lockedState, signal),
+      return await this.tracker.runExclusiveHighPriority(() =>
+        this.fetchContentAssistExclusive(params, signal),
       );
     } catch (error) {
       if ((error === E_CANCELED || error === E_TIMEOUT) && signal.aborted) {
@@ -204,37 +176,29 @@ export default class UpdateService {
 
   private async fetchContentAssistExclusive(
     params: ContentAssistParams,
-    lockedState: LockedState,
     signal: AbortSignal,
   ): Promise<ContentAssistEntry[]> {
     if (this.xtextStateId === undefined) {
-      await this.updateFullTextExclusive(lockedState);
+      await this.updateFullTextExclusive();
+      if (this.xtextStateId === undefined) {
+        throw new Error('failed to obtain Xtext state id');
+      }
     }
     if (signal.aborted) {
       return [];
     }
-    if (this.tracker.hasDirtyChanges) {
-      // Try to fetch while also performing a delta update.
-      const fetchUpdateEntries = await lockedState.withUpdateExclusive(
-        async (pendingUpdate) => {
-          const delta = pendingUpdate.prepareDeltaUpdateExclusive();
-          if (delta === undefined) {
-            return { newStateId: undefined, data: undefined };
-          }
-          log.trace('Editor delta', delta);
-          return this.doFetchContentAssistWithDeltaExclusive(
-            params,
-            pendingUpdate,
-            delta,
-          );
-        },
+    let entries: ContentAssistEntry[] | undefined;
+    if (this.tracker.needsUpdate) {
+      entries = await this.fetchContentAssistWithDeltaExclusive(
+        params,
+        this.xtextStateId,
       );
-      if (fetchUpdateEntries !== undefined) {
-        return fetchUpdateEntries;
-      }
-      if (signal.aborted) {
-        return [];
-      }
+    }
+    if (entries !== undefined) {
+      return entries;
+    }
+    if (signal.aborted) {
+      return [];
     }
     if (this.xtextStateId === undefined) {
       throw new Error('failed to obtain Xtext state id');
@@ -242,32 +206,35 @@ export default class UpdateService {
     return this.fetchContentAssistFetchOnly(params, this.xtextStateId);
   }
 
-  private async doFetchContentAssistWithDeltaExclusive(
+  private async fetchContentAssistWithDeltaExclusive(
     params: ContentAssistParams,
-    pendingUpdate: PendingUpdate,
-    delta: Delta,
-  ): Promise<StateUpdateResult<ContentAssistEntry[] | undefined>> {
+    requiredStateId: string,
+  ): Promise<ContentAssistEntry[] | undefined> {
+    const delta = this.tracker.prepareDeltaUpdateExclusive();
+    if (delta === undefined) {
+      return undefined;
+    }
+    log.trace('Editor delta for content assist', delta);
     const fetchUpdateResult = await this.webSocketClient.send({
       ...params,
       resource: this.resourceName,
       serviceType: 'assist',
-      requiredStateId: this.xtextStateId,
+      requiredStateId,
       ...delta,
     });
     const parsedContentAssistResult =
       ContentAssistResult.safeParse(fetchUpdateResult);
     if (parsedContentAssistResult.success) {
-      const { stateId, entries: resultEntries } =
-        parsedContentAssistResult.data;
-      return { newStateId: stateId, data: resultEntries };
+      const {
+        data: { stateId, entries },
+      } = parsedContentAssistResult;
+      this.tracker.setStateIdExclusive(stateId);
+      return entries;
     }
     if (isConflictResult(fetchUpdateResult, 'invalidStateId')) {
       log.warn('Server state invalid during content assist');
-      const newStateId = await this.doUpdateFullTextExclusive(pendingUpdate);
-      // We must finish this state update transaction to prepare for any push events
-      // before querying for content assist, so we just return `undefined` and will query
-      // the content assist service later.
-      return { newStateId, data: undefined };
+      await this.updateFullTextExclusive();
+      return undefined;
     }
     throw parsedContentAssistResult.error;
   }
@@ -294,33 +261,30 @@ export default class UpdateService {
   }
 
   formatText(): Promise<void> {
-    return this.tracker.runExclusiveWithRetries((lockedState) =>
-      this.formatTextExclusive(lockedState),
+    return this.tracker.runExclusiveWithRetries(() =>
+      this.formatTextExclusive(),
     );
   }
 
-  private async formatTextExclusive(lockedState: LockedState): Promise<void> {
-    await this.updateExclusive(lockedState);
+  private async formatTextExclusive(): Promise<void> {
+    await this.updateExclusive();
     let { from, to } = this.store.state.selection.main;
     if (to <= from) {
       from = 0;
       to = this.store.state.doc.length;
     }
     log.debug('Formatting from', from, 'to', to);
-    await lockedState.updateExclusive(async (pendingUpdate) => {
-      const result = await this.webSocketClient.send({
-        resource: this.resourceName,
-        serviceType: 'format',
-        selectionStart: from,
-        selectionEnd: to,
-      });
-      const { stateId, formattedText } = FormattingResult.parse(result);
-      pendingUpdate.applyBeforeDirtyChangesExclusive({
-        from,
-        to,
-        insert: formattedText,
-      });
-      return stateId;
+    const result = await this.webSocketClient.send({
+      resource: this.resourceName,
+      serviceType: 'format',
+      selectionStart: from,
+      selectionEnd: to,
+    });
+    const { stateId, formattedText } = FormattingResult.parse(result);
+    this.tracker.setStateIdExclusive(stateId, {
+      from,
+      to,
+      insert: formattedText,
     });
   }
 
@@ -335,7 +299,8 @@ export default class UpdateService {
       }
       throw error;
     }
-    if (!this.tracker.upToDate) {
+    const expectedStateId = this.xtextStateId;
+    if (expectedStateId === undefined || this.tracker.hasPendingChanges) {
       // Just give up if another update is in progress.
       return { cancelled: true };
     }
@@ -343,11 +308,6 @@ export default class UpdateService {
     if (caretOffsetResult.cancelled) {
       return { cancelled: true };
     }
-    const expectedStateId = this.xtextStateId;
-    if (expectedStateId === undefined) {
-      // If there is no state on the server, don't bother with finding occurrences.
-      return { cancelled: true };
-    }
     const data = await this.webSocketClient.send({
       resource: this.resourceName,
       serviceType: 'occurrences',
diff --git a/subprojects/frontend/src/xtext/UpdateStateTracker.ts b/subprojects/frontend/src/xtext/UpdateStateTracker.ts
index 04359060..a529f9a0 100644
--- a/subprojects/frontend/src/xtext/UpdateStateTracker.ts
+++ b/subprojects/frontend/src/xtext/UpdateStateTracker.ts
@@ -1,18 +1,3 @@
-/**
- * @file State tracker for pushing updates to the Xtext server.
- *
- * This file implements complex logic to avoid missing or overwriting state updates
- * and to avoid sending conflicting updates to the Xtext server.
- *
- * The `LockedState` and `PendingUpdate` objects are used as capabilities to
- * signify whether the socket to the Xtext server is locked for updates and
- * whether an update is in progress, respectively.
- * Always use these objects only received as an argument of a lambda expression
- * or method and never leak them into class field or global variables.
- * The presence of such an objects in the scope should always imply that
- * the corresponding condition holds.
- */
-
 import {
   type ChangeDesc,
   ChangeSet,
@@ -23,12 +8,9 @@ import {
 import { E_CANCELED, Mutex, withTimeout } from 'async-mutex';
 
 import type EditorStore from '../editor/EditorStore';
-import getLogger from '../utils/getLogger';
 
 const WAIT_FOR_UPDATE_TIMEOUT_MS = 1000;
 
-const log = getLogger('xtext.UpdateStateTracker');
-
 /**
  * State effect used to override the dirty changes after a transaction.
  *
@@ -40,57 +22,6 @@ const log = getLogger('xtext.UpdateStateTracker');
  */
 const setDirtyChanges = StateEffect.define<ChangeSet>();
 
-export interface StateUpdateResult<T> {
-  /** The new state ID on the server or `undefined` if no update was performed. */
-  newStateId: string | undefined;
-
-  /** Optional data payload received during the update. */
-  data: T;
-}
-
-/**
- * Signifies a capability that the Xtext server state is locked for update.
- */
-export interface LockedState {
-  /**
-   *
-   * @param callback the asynchronous callback that updates the server state
-   * @returns a promise resolving after the update
-   */
-  updateExclusive(
-    callback: (pendingUpdate: PendingUpdate) => Promise<string | undefined>,
-  ): Promise<void>;
-
-  /**
-   * Executes an asynchronous callback that updates the state on the server.
-   *
-   * If the callback returns `undefined` as the `newStateId`,
-   * the update is assumed to be aborted and any pending changes will be marked as dirt again.
-   * Any exceptions thrown in `callback` will also cause the update to be aborted.
-   *
-   * Ensures that updates happen sequentially and manages `pendingUpdate`
-   * and `dirtyChanges` to reflect changes being synchronized to the server
-   * and not yet synchronized to the server, respectively.
-   *
-   * Optionally, `callback` may return a second value that is retured by this function.
-   *
-   * Once the remote procedure call to update the server state finishes
-   * and returns the new `stateId`, `callback` must return _immediately_
-   * to ensure that the local `stateId` is updated likewise to be able to handle
-   * push messages referring to the new `stateId` from the server.
-   * If additional asynchronous work is needed to compute the second value in some cases,
-   * use `T | undefined` instead of `T` as a return type and signal the need for additional
-   * computations by returning `undefined`. Thus additional computations can be performed
-   * outside of the critical section.
-   *
-   * @param callback the asynchronous callback that updates the server state
-   * @returns a promise resolving to the second value returned by `callback`
-   */
-  withUpdateExclusive<T>(
-    callback: (pendingUpdate: PendingUpdate) => Promise<StateUpdateResult<T>>,
-  ): Promise<T>;
-}
-
 export interface Delta {
   deltaOffset: number;
 
@@ -99,32 +30,19 @@ export interface Delta {
   deltaText: string;
 }
 
-/**
- * Signifies a capability that dirty changes are being marked for uploading.
- */
-export interface PendingUpdate {
-  prepareDeltaUpdateExclusive(): Delta | undefined;
-
-  extendPendingUpdateExclusive(): void;
-
-  applyBeforeDirtyChangesExclusive(changeSpec: ChangeSpec): void;
-}
-
 export default class UpdateStateTracker {
   xtextStateId: string | undefined;
 
   /**
-   * The changes being synchronized to the server if a full or delta text update is running
-   * withing a `withUpdateExclusive` block, `undefined` otherwise.
+   * The changes marked for synchronization to the server if a full or delta text update
+   * is running, `undefined` otherwise.
    *
-   * Must be `undefined` before and after entering the critical section of `mutex`
-   * and may only be changes in the critical section of `mutex`.
+   * Must be `undefined` upon entering the critical section of `mutex`,
+   * may only be changed in the critical section of `mutex`,
+   * and will be set to `undefined` (marking any changes as dirty again) when leaving it.
    *
    * Methods named with an `Exclusive` suffix in this class assume that the mutex is held
-   * and may call `updateExclusive` or `withUpdateExclusive` to mutate this field.
-   *
-   * Methods named with a `do` suffix assume that they are called in a `withUpdateExclusive`
-   * block and require this field to be non-`undefined`.
+   * and may mutate this field.
    */
   private pendingChanges: ChangeSet | undefined;
 
@@ -142,24 +60,24 @@ export default class UpdateStateTracker {
     this.dirtyChanges = this.newEmptyChangeSet();
   }
 
-  get locekdForUpdate(): boolean {
-    return this.mutex.isLocked();
+  private get hasDirtyChanges(): boolean {
+    return !this.dirtyChanges.empty;
   }
 
-  get hasDirtyChanges(): boolean {
-    return !this.dirtyChanges.empty;
+  get needsUpdate(): boolean {
+    return this.hasDirtyChanges || this.xtextStateId === undefined;
+  }
+
+  get lockedForUpdate(): boolean {
+    return this.mutex.isLocked();
   }
 
-  get upToDate(): boolean {
-    return !this.locekdForUpdate && !this.hasDirtyChanges;
+  get hasPendingChanges(): boolean {
+    return this.lockedForUpdate || this.needsUpdate;
   }
 
   hasChangesSince(xtextStateId: string): boolean {
-    return (
-      this.xtextStateId !== xtextStateId ||
-      this.locekdForUpdate ||
-      this.hasDirtyChanges
-    );
+    return this.xtextStateId !== xtextStateId || this.hasPendingChanges;
   }
 
   /**
@@ -211,132 +129,102 @@ export default class UpdateStateTracker {
     );
   }
 
-  private newEmptyChangeSet(): ChangeSet {
-    return ChangeSet.of([], this.store.state.doc.length);
+  prepareDeltaUpdateExclusive(): Delta | undefined {
+    this.ensureLocked();
+    this.markDirtyChangesAsPendingExclusive();
+    if (this.pendingChanges === undefined || this.pendingChanges.empty) {
+      return undefined;
+    }
+    let minFromA = Number.MAX_SAFE_INTEGER;
+    let maxToA = 0;
+    let minFromB = Number.MAX_SAFE_INTEGER;
+    let maxToB = 0;
+    this.pendingChanges.iterChangedRanges((fromA, toA, fromB, toB) => {
+      minFromA = Math.min(minFromA, fromA);
+      maxToA = Math.max(maxToA, toA);
+      minFromB = Math.min(minFromB, fromB);
+      maxToB = Math.max(maxToB, toB);
+    });
+    return {
+      deltaOffset: minFromA,
+      deltaReplaceLength: maxToA - minFromA,
+      deltaText: this.store.state.doc.sliceString(minFromB, maxToB),
+    };
   }
 
-  private readonly pendingUpdate: PendingUpdate = {
-    prepareDeltaUpdateExclusive: (): Delta | undefined => {
-      this.pendingUpdate.extendPendingUpdateExclusive();
-      if (this.pendingChanges === undefined || this.pendingChanges.empty) {
-        return undefined;
-      }
-      let minFromA = Number.MAX_SAFE_INTEGER;
-      let maxToA = 0;
-      let minFromB = Number.MAX_SAFE_INTEGER;
-      let maxToB = 0;
-      this.pendingChanges.iterChangedRanges((fromA, toA, fromB, toB) => {
-        minFromA = Math.min(minFromA, fromA);
-        maxToA = Math.max(maxToA, toA);
-        minFromB = Math.min(minFromB, fromB);
-        maxToB = Math.max(maxToB, toB);
-      });
-      return {
-        deltaOffset: minFromA,
-        deltaReplaceLength: maxToA - minFromA,
-        deltaText: this.store.state.doc.sliceString(minFromB, maxToB),
-      };
-    },
-    extendPendingUpdateExclusive: (): void => {
-      if (!this.locekdForUpdate) {
-        throw new Error('Cannot update state without locking the mutex');
-      }
-      if (this.hasDirtyChanges) {
-        this.pendingChanges =
-          this.pendingChanges?.compose(this.dirtyChanges) ?? this.dirtyChanges;
-        this.dirtyChanges = this.newEmptyChangeSet();
-      }
-    },
-    applyBeforeDirtyChangesExclusive: (changeSpec: ChangeSpec): void => {
-      if (!this.locekdForUpdate) {
-        throw new Error('Cannot update state without locking the mutex');
-      }
-      const pendingChanges =
+  prepareFullTextUpdateExclusive(): void {
+    this.ensureLocked();
+    this.markDirtyChangesAsPendingExclusive();
+  }
+
+  private markDirtyChangesAsPendingExclusive(): void {
+    if (!this.lockedForUpdate) {
+      throw new Error('Cannot update state without locking the mutex');
+    }
+    if (this.hasDirtyChanges) {
+      this.pendingChanges =
         this.pendingChanges?.compose(this.dirtyChanges) ?? this.dirtyChanges;
-      const revertChanges = pendingChanges.invert(this.store.state.doc);
-      const applyBefore = ChangeSet.of(changeSpec, revertChanges.newLength);
-      const redoChanges = pendingChanges.map(applyBefore.desc);
-      const changeSet = revertChanges.compose(applyBefore).compose(redoChanges);
-      this.store.dispatch({
-        changes: changeSet,
-        // Keep the current set of dirty changes (but update them according the re-formatting)
-        // and to not add the formatting the dirty changes.
-        effects: [setDirtyChanges.of(redoChanges)],
-      });
-    },
-  };
+      this.dirtyChanges = this.newEmptyChangeSet();
+    }
+  }
 
-  private readonly lockedState: LockedState = {
-    updateExclusive: (
-      callback: (pendingUpdate: PendingUpdate) => Promise<string | undefined>,
-    ): Promise<void> => {
-      return this.lockedState.withUpdateExclusive<void>(
-        async (pendingUpdate) => {
-          const newStateId = await callback(pendingUpdate);
-          return { newStateId, data: undefined };
-        },
-      );
-    },
-    withUpdateExclusive: async <T>(
-      callback: (pendingUpdate: PendingUpdate) => Promise<StateUpdateResult<T>>,
-    ): Promise<T> => {
-      if (!this.locekdForUpdate) {
-        throw new Error('Cannot update state without locking the mutex');
-      }
-      if (this.pendingChanges !== undefined) {
-        throw new Error('Delta updates are not reentrant');
-      }
-      let newStateId: string | undefined;
-      let data: T;
-      try {
-        ({ newStateId, data } = await callback(this.pendingUpdate));
-      } catch (e) {
-        log.error('Error while update', e);
-        this.cancelUpdate();
-        throw e;
-      }
-      if (newStateId === undefined) {
-        this.cancelUpdate();
-      } else {
-        this.xtextStateId = newStateId;
-        this.pendingChanges = undefined;
-      }
-      return data;
-    },
-  };
+  private newEmptyChangeSet(): ChangeSet {
+    return ChangeSet.of([], this.store.state.doc.length);
+  }
 
-  private cancelUpdate(): void {
-    if (this.pendingChanges === undefined) {
-      return;
+  setStateIdExclusive(
+    newStateId: string,
+    remoteChanges?: ChangeSpec | undefined,
+  ): void {
+    this.ensureLocked();
+    if (remoteChanges !== undefined) {
+      this.applyRemoteChangesExclusive(remoteChanges);
     }
-    this.dirtyChanges = this.pendingChanges.compose(this.dirtyChanges);
+    this.xtextStateId = newStateId;
     this.pendingChanges = undefined;
   }
 
-  runExclusive<T>(
-    callback: (lockedState: LockedState) => Promise<T>,
-  ): Promise<T> {
+  private applyRemoteChangesExclusive(changeSpec: ChangeSpec): void {
+    const pendingChanges =
+      this.pendingChanges?.compose(this.dirtyChanges) ?? this.dirtyChanges;
+    const revertChanges = pendingChanges.invert(this.store.state.doc);
+    const applyBefore = ChangeSet.of(changeSpec, revertChanges.newLength);
+    const redoChanges = pendingChanges.map(applyBefore.desc);
+    const changeSet = revertChanges.compose(applyBefore).compose(redoChanges);
+    this.store.dispatch({
+      changes: changeSet,
+      // Keep the current set of dirty changes (but update them according the re-formatting)
+      // and to not add the formatting the dirty changes.
+      effects: [setDirtyChanges.of(redoChanges)],
+    });
+  }
+
+  private ensureLocked(): void {
+    if (!this.lockedForUpdate) {
+      throw new Error('Cannot update state without locking the mutex');
+    }
+  }
+
+  runExclusive<T>(callback: () => Promise<T>): Promise<T> {
     return this.mutex.runExclusive(async () => {
-      if (this.pendingChanges !== undefined) {
-        throw new Error('Update is pending before entering critical section');
-      }
-      const result = await callback(this.lockedState);
-      if (this.pendingChanges !== undefined) {
-        throw new Error('Update is pending after entering critical section');
+      try {
+        return await callback();
+      } finally {
+        if (this.pendingChanges !== undefined) {
+          this.dirtyChanges = this.pendingChanges.compose(this.dirtyChanges);
+          this.pendingChanges = undefined;
+        }
       }
-      return result;
     });
   }
 
-  runExclusiveHighPriority<T>(
-    callback: (lockedState: LockedState) => Promise<T>,
-  ): Promise<T> {
+  runExclusiveHighPriority<T>(callback: () => Promise<T>): Promise<T> {
     this.mutex.cancel();
     return this.runExclusive(callback);
   }
 
   async runExclusiveWithRetries<T>(
-    callback: (lockedState: LockedState) => Promise<T>,
+    callback: () => Promise<T>,
     maxRetries = 5,
   ): Promise<T> {
     let retries = 0;