Add thunk action abort signal cancellation

Author: jcobis

packages/compass-collection/src/modules/collection-tab.ts

@@ -79,6 +79,7 @@ type CollectionThunkAction<R, A extends AnyAction = AnyAction> = ThunkAction<
     preferences: PreferencesAccess;
     connectionInfoRef: ConnectionInfoRef;
     fakerSchemaGenerationAbortControllerRef: { current?: AbortController };
+    schemaAnalysisAbortControllerRef: { current?: AbortController };
   },
   A
 >;
@@ -101,6 +102,7 @@ export enum CollectionActions {
   SchemaAnalysisStarted = 'compass-collection/SchemaAnalysisStarted',
   SchemaAnalysisFinished = 'compass-collection/SchemaAnalysisFinished',
   SchemaAnalysisFailed = 'compass-collection/SchemaAnalysisFailed',
+  SchemaAnalysisCanceled = 'compass-collection/SchemaAnalysisCanceled',
   SchemaAnalysisReset = 'compass-collection/SchemaAnalysisReset',
   MockDataGeneratorModalOpened = 'compass-collection/MockDataGeneratorModalOpened',
   MockDataGeneratorModalClosed = 'compass-collection/MockDataGeneratorModalClosed',
@@ -139,6 +141,10 @@ interface SchemaAnalysisFailedAction {
   error: Error;
 }
 
+interface SchemaAnalysisCanceledAction {
+  type: CollectionActions.SchemaAnalysisCanceled;
+}
+
 interface MockDataGeneratorModalOpenedAction {
   type: CollectionActions.MockDataGeneratorModalOpened;
 }
@@ -263,6 +269,20 @@ const reducer: Reducer<CollectionState, Action> = (
     };
   }
 
+  if (
+    isAction<SchemaAnalysisCanceledAction>(
+      action,
+      CollectionActions.SchemaAnalysisCanceled
+    )
+  ) {
+    return {
+      ...state,
+      schemaAnalysis: {
+        status: SCHEMA_ANALYSIS_STATE_INITIAL,
+      },
+    };
+  }
+
   if (
     isAction<MockDataGeneratorModalOpenedAction>(
       action,
@@ -524,7 +544,11 @@ export const openMockDataGeneratorModal = (): CollectionThunkAction<
 export const analyzeCollectionSchema = (): CollectionThunkAction<
   Promise<void>
 > => {
-  return async (dispatch, getState, { dataService, preferences, logger }) => {
+  return async (
+    dispatch,
+    getState,
+    { dataService, preferences, logger, schemaAnalysisAbortControllerRef }
+  ) => {
     const { schemaAnalysis, namespace } = getState();
     const analysisStatus = schemaAnalysis.status;
     if (analysisStatus === SCHEMA_ANALYSIS_STATE_ANALYZING) {
@@ -534,6 +558,10 @@ export const analyzeCollectionSchema = (): CollectionThunkAction<
       return;
     }
 
+    // Create abort controller for this analysis
+    const abortController = new AbortController();
+    schemaAnalysisAbortControllerRef.current = abortController;
+
     try {
       logger.debug('Schema analysis started.');
 
@@ -552,8 +580,15 @@ export const analyzeCollectionSchema = (): CollectionThunkAction<
         driverOptions,
         {
           fallbackReadPreference: 'secondaryPreferred',
+          abortSignal: abortController.signal,
         }
       );
+
+      // Check if analysis was aborted after sampling
+      if (abortController.signal.aborted) {
+        logger.debug('Schema analysis was aborted during sampling');
+        return;
+      }
       if (sampleDocuments.length === 0) {
         logger.debug(NO_DOCUMENTS_ERROR);
         dispatch({
@@ -565,6 +600,13 @@ export const analyzeCollectionSchema = (): CollectionThunkAction<
 
       // Analyze sampled documents
       const schemaAccessor = await analyzeDocuments(sampleDocuments);
+
+      // Check if analysis was aborted after document analysis
+      if (abortController.signal.aborted) {
+        logger.debug('Schema analysis was aborted during document analysis');
+        return;
+      }
+
       const schema = await schemaAccessor.getInternalSchema();
 
       // Filter out internal fields from the schema
@@ -583,13 +625,29 @@ export const analyzeCollectionSchema = (): CollectionThunkAction<
         maxNestingDepth,
         validationRules,
       };
+
+      // Final check before dispatching results
+      if (abortController.signal.aborted) {
+        logger.debug('Schema analysis was aborted before completion');
+        return;
+      }
+
       dispatch({
         type: CollectionActions.SchemaAnalysisFinished,
         processedSchema,
         sampleDocument: sampleDocuments[0],
         schemaMetadata,
       });
     } catch (err: any) {
+      // Check if the error is due to cancellation
+      if (isCancelError(err) || abortController.signal.aborted) {
+        logger.debug('Schema analysis was aborted');
+        dispatch({
+          type: CollectionActions.SchemaAnalysisCanceled,
+        });
+        return;
+      }
+
       logger.log.error(
         mongoLogId(1_001_000_363),
         'Collection',
@@ -603,6 +661,23 @@ export const analyzeCollectionSchema = (): CollectionThunkAction<
         type: CollectionActions.SchemaAnalysisFailed,
         error: err as Error,
       });
+    } finally {
+      // Clean up abort controller
+      schemaAnalysisAbortControllerRef.current = undefined;
+    }
+  };
+};
+
+export const cancelSchemaAnalysis = (): CollectionThunkAction<void> => {
+  return (
+    _dispatch,
+    _getState,
+    { schemaAnalysisAbortControllerRef, logger }
+  ) => {
+    if (schemaAnalysisAbortControllerRef.current) {
+      logger.debug('Canceling schema analysis');
+      schemaAnalysisAbortControllerRef.current.abort();
+      schemaAnalysisAbortControllerRef.current = undefined;
     }
   };
 };

packages/compass-collection/src/stores/collection-tab.spec.ts

@@ -134,7 +134,7 @@ describe('Collection Tab Content store', function () {
         logger,
         preferences,
       },
-      { on() {}, cleanup() {} } as any
+      { on() {}, cleanup() {}, addCleanup() {} } as any
     ));
     await waitFor(() => {
       expect(store.getState())
@@ -417,4 +417,28 @@ describe('Collection Tab Content store', function () {
       expect(analyzeCollectionSchemaStub).to.not.have.been.called;
     });
   });
+
+  describe('schema analysis cancellation', function () {
+    it('should cancel schema analysis when cancelSchemaAnalysis is dispatched', async function () {
+      const getAssignment = sandbox.spy(() =>
+        Promise.resolve(
+          createMockAssignment(ExperimentTestGroup.mockDataGeneratorVariant)
+        )
+      );
+      const assignExperiment = sandbox.spy(() => Promise.resolve(null));
+
+      const store = await configureStore(undefined, undefined, {
+        getAssignment,
+        assignExperiment,
+      });
+
+      // Dispatch cancel action
+      store.dispatch(collectionTabModule.cancelSchemaAnalysis() as any);
+
+      // Verify the state is reset to initial
+      expect((store.getState() as any).schemaAnalysis.status).to.equal(
+        'initial'
+      );
+    });
+  });
 });

packages/compass-collection/src/stores/collection-tab.ts

@@ -7,6 +7,7 @@ import reducer, {
   selectTab,
   collectionMetadataFetched,
   analyzeCollectionSchema,
+  cancelSchemaAnalysis,
 } from '../modules/collection-tab';
 import { MockDataGeneratorStep } from '../components/mock-data-generator-modal/types';
 
@@ -58,7 +59,7 @@ export type CollectionTabServices = {
 export function activatePlugin(
   { namespace, editViewName, tabId }: CollectionTabOptions,
   services: CollectionTabServices,
-  { on, cleanup }: ActivateHelpers
+  { on, cleanup, addCleanup }: ActivateHelpers
 ): {
   store: ReturnType<typeof createStore>;
   deactivate: () => void;
@@ -84,6 +85,9 @@ export function activatePlugin(
   const fakerSchemaGenerationAbortControllerRef = {
     current: undefined,
   };
+  const schemaAnalysisAbortControllerRef = {
+    current: undefined,
+  };
   const store = createStore(
     reducer,
     {
@@ -113,6 +117,7 @@ export function activatePlugin(
         logger,
         preferences,
         fakerSchemaGenerationAbortControllerRef,
+        schemaAnalysisAbortControllerRef,
       })
     )
   );
@@ -190,6 +195,9 @@ export function activatePlugin(
     }
   });
 
+  // Cancel schema analysis when plugin is deactivated
+  addCleanup(() => store.dispatch(cancelSchemaAnalysis()));
+
   return {
     store,
     deactivate: cleanup,