[core-lro] Set isCancelled when operation status is cancelled (#21893)

* [core-lro] Set isCancelled when status is cancelled

* don't check for isCanceled in TA test

* fix lint

* address feedback and handle cancellation uniformly

* address feedback

* add tests

* edit

* revert behavioral change

* Update sdk/textanalytics/ai-text-analytics/package.json

Co-authored-by: Will Temple <witemple@microsoft.com>

sdk/core/core-lro/CHANGELOG.md

@@ -1,15 +1,11 @@
 # Release History
 
-## 2.3.0 (Unreleased)
+## 2.3.0-beta.1 (2022-05-18)
 
 ### Features Added
 
 - `lroEngine` now supports cancellation of the long-running operation.
 
-### Breaking Changes
-
-### Bugs Fixed
-
 ### Other Changes
 
 - Removed the unused dependency `@azure/core-tracing`.

sdk/core/core-lro/package.json

@@ -2,7 +2,7 @@
   "name": "@azure/core-lro",
   "author": "Microsoft Corporation",
   "sdk-type": "client",
-  "version": "2.3.0",
+  "version": "2.3.0-beta.1",
   "description": "Isomorphic client library for supporting long-running operations in node.js and browser.",
   "tags": [
     "isomorphic",

sdk/core/core-lro/src/lroEngine/bodyPolling.ts

@@ -1,39 +1,31 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import {
-  LroBody,
-  LroResponse,
-  LroStatus,
-  RawResponse,
-  failureStates,
-  successStates,
-} from "./models";
-import { isUnexpectedPollingResponse } from "./requestUtils";
-
-function getProvisioningState(rawResponse: RawResponse): string {
-  const { properties, provisioningState } = (rawResponse.body as LroBody) ?? {};
-  const state: string | undefined = properties?.provisioningState ?? provisioningState;
-  return typeof state === "string" ? state.toLowerCase() : "succeeded";
-}
-
-export function isBodyPollingDone(rawResponse: RawResponse): boolean {
-  const state = getProvisioningState(rawResponse);
-  if (isUnexpectedPollingResponse(rawResponse) || failureStates.includes(state)) {
-    throw new Error(`The long running operation has failed. The provisioning state: ${state}.`);
-  }
-  return successStates.includes(state);
-}
+import { LroResponse, LroStatus } from "./models";
+import { getProvisioningState, isCanceled, isPollingDone } from "./requestUtils";
+import { PollOperationState } from "../pollOperation";
 
 /**
  * Creates a polling strategy based on BodyPolling which uses the provisioning state
  * from the result to determine the current operation state
  */
-export function processBodyPollingOperationResult<TResult>(
-  response: LroResponse<TResult>
-): LroStatus<TResult> {
-  return {
-    ...response,
-    done: isBodyPollingDone(response.rawResponse),
+export function processBodyPollingOperationResult<
+  TResult,
+  TState extends PollOperationState<TResult>
+>(state: TState): (response: LroResponse<TResult>) => LroStatus<TResult> {
+  return (response: LroResponse<TResult>): LroStatus<TResult> => {
+    const status = getProvisioningState(response.rawResponse);
+    return {
+      ...response,
+      done:
+        isCanceled({
+          state,
+          status,
+        }) ||
+        isPollingDone({
+          rawResponse: response.rawResponse,
+          status,
+        }),
+    };
   };
 }

sdk/core/core-lro/src/lroEngine/locationPolling.ts

@@ -8,21 +8,20 @@ import {
   LroResponse,
   LroStatus,
   RawResponse,
-  failureStates,
-  successStates,
 } from "./models";
-import { isUnexpectedPollingResponse } from "./requestUtils";
+import { isCanceled, isPollingDone } from "./requestUtils";
+import { PollOperationState } from "../pollOperation";
 
-function isPollingDone(rawResponse: RawResponse): boolean {
-  if (isUnexpectedPollingResponse(rawResponse) || rawResponse.statusCode === 202) {
-    return false;
-  }
+function getStatus(rawResponse: RawResponse): string {
   const { status } = (rawResponse.body as LroBody) ?? {};
-  const state = typeof status === "string" ? status.toLowerCase() : "succeeded";
-  if (isUnexpectedPollingResponse(rawResponse) || failureStates.includes(state)) {
-    throw new Error(`The long running operation has failed. The provisioning state: ${state}.`);
+  return typeof status === "string" ? status.toLowerCase() : "succeeded";
+}
+
+function isLocationPollingDone(rawResponse: RawResponse, status: string): boolean {
+  if (rawResponse.statusCode === 202) {
+    return false;
   }
-  return successStates.includes(state);
+  return isPollingDone({ rawResponse, status });
 }
 
 /**
@@ -44,13 +43,24 @@ async function sendFinalRequest<TResult>(
   }
 }
 
-export function processLocationPollingOperationResult<TResult>(
+export function processLocationPollingOperationResult<
+  TResult,
+  TState extends PollOperationState<TResult>
+>(
   lro: LongRunningOperation<TResult>,
+  state: TState,
   resourceLocation?: string,
   lroResourceLocationConfig?: LroResourceLocationConfig
 ): (response: LroResponse<TResult>) => LroStatus<TResult> {
   return (response: LroResponse<TResult>): LroStatus<TResult> => {
-    if (isPollingDone(response.rawResponse)) {
+    const status = getStatus(response.rawResponse);
+    if (
+      isCanceled({
+        state,
+        status,
+      }) ||
+      isLocationPollingDone(response.rawResponse, status)
+    ) {
       if (resourceLocation === undefined) {
         return { ...response, done: true };
       } else {

sdk/core/core-lro/src/lroEngine/models.ts

@@ -33,18 +33,12 @@ export interface LroEngineOptions<TResult, TState> {
   isDone?: (lastResponse: unknown, state: TState) => boolean;
 
   /**
-   * A function to cancel the LRO.
+   * A function that takes the mutable state as input and attempts to cancel the
+   * LRO.
    */
   cancel?: (state: TState) => Promise<void>;
 }
 
-export const successStates = ["succeeded"];
-export const failureStates = ["failed", "canceled", "cancelled"];
-/**
- * The LRO states that signal that the LRO has completed.
- */
-export const terminalStates = successStates.concat(failureStates);
-
 /**
  * The potential location of the result of the LRO if specified by the LRO extension in the swagger.
  */

sdk/core/core-lro/src/lroEngine/operation.ts

@@ -86,7 +86,12 @@ export class GenericPollOperation<TResult, TState extends PollOperationState<TRe
               ...response,
               done: isDone(response.flatResponse, this.state),
             })
-          : createGetLroStatusFromResponse(this.lro, state.config, this.lroResourceLocationConfig);
+          : createGetLroStatusFromResponse(
+              this.lro,
+              state.config,
+              this.state,
+              this.lroResourceLocationConfig
+            );
         this.poll = createPoll(this.lro);
       }
       if (!state.pollingURL) {
@@ -122,8 +127,13 @@ export class GenericPollOperation<TResult, TState extends PollOperationState<TRe
   }
 
   async cancel(): Promise<PollOperation<TState, TResult>> {
-    this.state.isCancelled = true;
     await this.cancelOp?.(this.state);
+    /**
+     * When `cancelOperation` is called, polling stops so it is important that
+     * `isCancelled` is set now because the polling logic will not be able to
+     * set it itself because it will not fire.
+     */
+    this.state.isCancelled = true;
     return this;
   }
 

sdk/core/core-lro/src/lroEngine/requestUtils.ts

@@ -1,7 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
-import { LroConfig, RawResponse } from "./models";
+import { LroBody, LroConfig, RawResponse } from "./models";
+import { PollOperationState } from "../pollOperation";
 
 /**
  * Detects where the continuation token is and returns it. Notice that azure-asyncoperation
@@ -106,3 +107,33 @@ export function isUnexpectedPollingResponse(rawResponse: RawResponse): boolean {
   }
   return false;
 }
+
+export function isCanceled<TResult, TState extends PollOperationState<TResult>>(operation: {
+  state: TState;
+  status: string;
+}): boolean {
+  const { state, status } = operation;
+  if (["canceled", "cancelled"].includes(status)) {
+    state.isCancelled = true;
+    throw new Error(`The long-running operation has been canceled.`);
+  }
+  return false;
+}
+
+export function isSucceededStatus(status: string): boolean {
+  return status === "succeeded";
+}
+
+export function isPollingDone(result: { rawResponse: RawResponse; status: string }): boolean {
+  const { rawResponse, status } = result;
+  if (isUnexpectedPollingResponse(rawResponse) || status === "failed") {
+    throw new Error(`The long-running operation has failed.`);
+  }
+  return isSucceededStatus(status);
+}
+
+export function getProvisioningState(rawResponse: RawResponse): string {
+  const { properties, provisioningState } = (rawResponse.body as LroBody) ?? {};
+  const state: string | undefined = properties?.provisioningState ?? provisioningState;
+  return typeof state === "string" ? state.toLowerCase() : "succeeded";
+}

sdk/core/core-lro/src/lroEngine/stateMachine.ts

@@ -11,30 +11,39 @@ import {
   PollerConfig,
   ResumablePollOperationState,
 } from "./models";
-import { getPollingUrl, inferLroMode, isUnexpectedInitialResponse } from "./requestUtils";
-import { isBodyPollingDone, processBodyPollingOperationResult } from "./bodyPolling";
+import {
+  getPollingUrl,
+  getProvisioningState,
+  inferLroMode,
+  isPollingDone,
+  isUnexpectedInitialResponse,
+} from "./requestUtils";
+import { PollOperationState } from "../pollOperation";
 import { logger } from "./logger";
+import { processBodyPollingOperationResult } from "./bodyPolling";
 import { processLocationPollingOperationResult } from "./locationPolling";
 import { processPassthroughOperationResult } from "./passthrough";
 
 /**
  * creates a stepping function that maps an LRO state to another.
  */
-export function createGetLroStatusFromResponse<TResult>(
+export function createGetLroStatusFromResponse<TResult, TState extends PollOperationState<TResult>>(
   lroPrimitives: LongRunningOperation<TResult>,
   config: LroConfig,
+  state: TState,
   lroResourceLocationConfig?: LroResourceLocationConfig
 ): GetLroStatusFromResponse<TResult> {
   switch (config.mode) {
     case "Location": {
       return processLocationPollingOperationResult(
         lroPrimitives,
+        state,
         config.resourceLocation,
         lroResourceLocationConfig
       );
     }
     case "Body": {
-      return processBodyPollingOperationResult;
+      return processBodyPollingOperationResult(state);
     }
     default: {
       return processPassthroughOperationResult;
@@ -103,7 +112,11 @@ export function createInitializeState<TResult>(
     /** short circuit polling if body polling is done in the initial request */
     if (
       state.config.mode === undefined ||
-      (state.config.mode === "Body" && isBodyPollingDone(state.initialRawResponse))
+      (state.config.mode === "Body" &&
+        isPollingDone({
+          rawResponse: state.initialRawResponse,
+          status: getProvisioningState(state.initialRawResponse),
+        }))
     ) {
       state.result = response.flatResponse as TResult;
       state.isCompleted = true;