Improve image upload cleanup logic, reduce e2e flake (#2674)

fix image upload cleanup logic + e2es

Author: david-crespo

app/forms/image-upload.tsx

@@ -184,12 +184,6 @@ export function Component() {
   const queryClient = useApiQueryClient()
   const { project } = useProjectSelector()
 
-  // Note: abort currently only works if it fires during the upload file step.
-  // We could make it work between the other steps by calling
-  // `abortController.throwIfAborted()` after each one. We could technically
-  // plumb through the signal to the requests themselves, but they complete so
-  // quickly it's probably not necessary.
-
   // The state in this component is very complex because we are doing a bunch of
   // requests in order, all of which can fail, plus the whole thing can be
   // aborted. We have the usual form state, plus an additional validation step
@@ -250,8 +244,19 @@ export function Component() {
   // separate so we can distinguish between cleanup due to error vs. cleanup after success
   const stopImportCleanup = useApiMutation('diskBulkWriteImportStop')
   const finalizeDiskCleanup = useApiMutation('diskFinalizeImport')
-  const deleteDiskCleanup = useApiMutation('diskDelete')
-  const deleteSnapshotCleanup = useApiMutation('snapshotDelete')
+  // in production these invalidations are unlikely to matter, but they help a
+  // lot in the tests when we check the disk list after canceling to make sure
+  // the temp resources got deleted
+  const deleteDiskCleanup = useApiMutation('diskDelete', {
+    onSuccess() {
+      queryClient.invalidateQueries('diskList')
+    },
+  })
+  const deleteSnapshotCleanup = useApiMutation('snapshotDelete', {
+    onSuccess() {
+      queryClient.invalidateQueries('snapshotList')
+    },
+  })
 
   const cleanupMutations = [
     stopImportCleanup,
@@ -270,31 +275,36 @@ export function Component() {
   const snapshot = useRef<Snapshot | null>(null)
   const disk = useRef<Disk | null>(null)
 
-  // if closeModal runs during bulk upload due to a cancel, cancelEverything
-  // causes an abort of the bulk upload, which throws an error to onSubmit's
-  // catch, which calls `cleanup`. so when we call cleanup here, it will be a
-  // double cleanup. we could get rid of this one, but for the rare cancel *not*
-  // during bulk upload we will still want to call cleanup. rather than
-  // coordinating when to cleanup, we make cleanup idempotent by having it check
-  // whether it has already been run, or more concretely before each action,
-  // check whether it needs to be done
   function closeModal() {
     if (allDone) {
       backToImages()
       return
     }
 
-    // if we're still going, need to confirm cancelation. if we have an error,
+    // if we're still going, need to confirm cancellation. if we have an error,
     // everything is already stopped
     if (modalError || confirm('Are you sure? Closing the modal will cancel the upload.')) {
+      // Note we don't run cleanup() here -- cancelEverything triggers an
+      // abort, which gets caught by the try/catch in the onSubmit on the upload
+      // form, which does the cleanup. We used to call cleanup here and used
+      // error-prone state logic to avoid it running twice.
+      //
+      // Because we are working with a closed-over copy of allDone, there is
+      // a possibility that the upload finishes while the user is looking at
+      // the confirm modal, in which case cancelEverything simply won't do
+      // anything. The finally{} in onSubmit clears out the abortController so
+      // cancelEverything() is a noop.
       cancelEverything()
-      // TODO: probably shouldn't await this, but we do need to catch errors
-      cleanup()
       resetMainFlow()
       setModalOpen(false)
     }
   }
 
+  // Aborting works for steps other than file upload despite the
+  // signal not being used directly in the former because we call
+  // `abortController.throwIfAborted()` after each step. We could technically
+  // plumb through the signal to the requests themselves, but they complete so
+  // quickly it's probably not necessary.
   function cancelEverything() {
     abortController.current?.abort(ABORT_ERROR)
   }
@@ -306,14 +316,8 @@ export function Component() {
     setSyntheticUploadState(initSyntheticState)
   }
 
-  const cleaningUp = useRef(false)
-
   /** If a snapshot or disk was created, clean it up*/
   async function cleanup() {
-    // don't run if already running
-    if (cleaningUp.current) return
-    cleaningUp.current = true
-
     if (snapshot.current) {
       await deleteSnapshotCleanup.mutateAsync({ path: { snapshot: snapshot.current.id } })
       snapshot.current = null
@@ -335,7 +339,6 @@ export function Component() {
       await deleteDiskCleanup.mutateAsync({ path: { disk: disk.current.id } })
       disk.current = null
     }
-    cleaningUp.current = false
   }
 
   async function onSubmit({
@@ -500,10 +503,6 @@ export function Component() {
       title="Upload image"
       onDismiss={backToImages}
       onSubmit={async (values) => {
-        // every submit needs its own AbortController because they can't be
-        // reset
-        abortController.current = new AbortController()
-
         setFormError(null)
 
         // check that image name isn't taken before starting the whole thing
@@ -532,16 +531,28 @@ export function Component() {
           return
         }
 
+        // every submit needs its own AbortController because they can't be
+        // reset
+        abortController.current = new AbortController()
+
         try {
           await onSubmit(values)
         } catch (e) {
           if (e !== ABORT_ERROR) {
             setModalError('Something went wrong. Please try again.')
+            // abort anything in flight in case
+            cancelEverything()
           }
-          cancelEverything()
           // user canceled
           await cleanup()
           // TODO: if we get here, show failure state in the upload modal
+        } finally {
+          // Clear the abort controller. This is aimed at the case where the
+          // user clicks cancel and then stares at the confirm modal without
+          // clicking for so long that the upload manages to finish, which means
+          // there's no longer anything to cancel. If abortController is gone,
+          // cancelEverything is a noop.
+          abortController.current = null
         }
       }}
       loading={formLoading}

mock-api/msw/handlers.ts

@@ -180,7 +180,7 @@ export const handlers = makeHandlers({
       ),
     }
   },
-  diskBulkWriteImportStart: ({ path, query }) => {
+  async diskBulkWriteImportStart({ path, query }) {
     const disk = lookup.disk({ ...path, ...query })
 
     if (disk.name === 'import-start-500') throw 500
@@ -189,20 +189,21 @@ export const handlers = makeHandlers({
       throw 'Can only enter state importing_from_bulk_write from import_ready'
     }
 
-    // throw 400
+    await delay(1000) // slow it down for the tests
 
     db.diskBulkImportState.set(disk.id, { blocks: {} })
     disk.state = { state: 'importing_from_bulk_writes' }
     return 204
   },
-  diskBulkWriteImportStop: ({ path, query }) => {
+  async diskBulkWriteImportStop({ path, query }) {
     const disk = lookup.disk({ ...path, ...query })
 
     if (disk.name === 'import-stop-500') throw 500
 
     if (disk.state.state !== 'importing_from_bulk_writes') {
       throw 'Can only stop import for disk in state importing_from_bulk_write'
     }
+    await delay(1000) // slow it down for the tests
 
     db.diskBulkImportState.delete(disk.id)
     disk.state = { state: 'import_ready' }

test/e2e/image-upload.e2e.ts

@@ -123,41 +123,45 @@ test.describe('Image upload', () => {
     await expectVisible(page, [fileRequired])
   })
 
-  test('cancel', async ({ page, browserName }) => {
-    // eslint-disable-next-line playwright/no-skipped-test
-    test.skip(browserName === 'webkit', 'safari. stop this')
-
-    await fillForm(page, 'new-image')
+  const cancelStates = [
+    'Put disk in import mode',
+    'Upload image file',
+    'Get disk out of import mode',
+    // 'Finalize disk and create snapshot',
+  ]
 
-    await page.getByRole('button', { name: 'Upload image' }).click()
+  for (const state of cancelStates) {
+    test(`cancel in state '${state}'`, async ({ page }) => {
+      await fillForm(page, 'new-image')
 
-    const progressModal = page.getByRole('dialog', { name: 'Image upload progress' })
-    await expect(progressModal).toBeVisible()
+      await page.getByRole('button', { name: 'Upload image' }).click()
 
-    // wait to be in the middle of upload
-    const uploadStep = page.getByTestId('upload-step: Upload image file')
-    await expect(uploadStep).toHaveAttribute('data-status', 'running')
+      const progressModal = page.getByRole('dialog', { name: 'Image upload progress' })
+      await expect(progressModal).toBeVisible()
 
-    // form is disabled and semi-hidden
-    // await expectNotVisible(page, ['role=textbox[name="Name"]'])
+      // wait to be in the middle of upload
+      const uploadStep = page.getByTestId(`upload-step: ${state}`)
+      await expect(uploadStep).toHaveAttribute('data-status', 'running')
 
-    page.on('dialog', (dialog) => dialog.accept()) // click yes on the are you sure prompt
-    await progressModal.getByRole('button', { name: 'Cancel' }).click()
+      // form is disabled and semi-hidden
+      // await expectNotVisible(page, ['role=textbox[name="Name"]'])
 
-    // modal has closed
-    await expect(progressModal).toBeHidden()
+      page.on('dialog', (dialog) => dialog.accept()) // click yes on the are you sure prompt
+      await progressModal.getByRole('button', { name: 'Cancel' }).click()
 
-    // form's back
-    await expectVisible(page, ['role=textbox[name="Name"]'])
+      // modal has closed
+      await expect(progressModal).toBeHidden()
 
-    // get out of the form
-    await page.click('text=Cancel')
+      // form's back
+      await expect(page.getByRole('textbox', { name: 'Name' })).toBeVisible()
 
-    // TODO: go to disks and make sure the tmp one got cleaned up
-    // await page.click('role=link[name="Disks"]')
-    // await expectVisible(page, ['role=cell[name="disk-1"]'])
-    // await expectNotVisible(page, ['role=cell[name=tmp]'])
-  })
+      // get out of the form and go to the disks page to check it's not there
+      await page.getByRole('button', { name: 'Cancel' }).click()
+      await page.getByRole('link', { name: 'Disks' }).click()
+      await expect(page.getByRole('cell', { name: 'disk-1', exact: true })).toBeVisible()
+      await expect(page.getByRole('cell', { name: 'tmp' })).toBeHidden()
+    })
+  }
 
   // testing the onFocusOutside fix
   test('cancel canceling', async ({ page, browserName }) => {