Merge pull request #5178 from microsoft/bugfix-cancel-login2

fix a bug of canceling login doesn't cancel background job

Author: andxu

PluginsAndFeatures/azure-toolkit-for-intellij/src/com/microsoft/intellij/ui/DeviceLoginUI.java

@@ -14,14 +14,15 @@
 import lombok.Setter;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
-import reactor.core.Disposable;
+
+import java.util.concurrent.Future;
 
 
 public class DeviceLoginUI implements IDeviceLoginUI {
     private DeviceLoginWindow deviceLoginWindow;
 
     @Setter
-    private Disposable disposable;
+    private Future future;
 
     @Nullable
     @Override
@@ -45,8 +46,8 @@ public void closePrompt() {
 
     @Override
     public void cancel() {
-        if (disposable != null) {
-            this.disposable.dispose();
+        if (future != null) {
+            this.future.cancel(true);
         }
     }
 }

PluginsAndFeatures/azure-toolkit-for-intellij/src/com/microsoft/intellij/ui/SignInWindow.java

@@ -10,6 +10,7 @@
 import com.intellij.openapi.fileChooser.FileChooser;
 import com.intellij.openapi.fileChooser.FileChooserDescriptor;
 import com.intellij.openapi.fileChooser.FileChooserDescriptorFactory;
+import com.intellij.openapi.progress.ProcessCanceledException;
 import com.intellij.openapi.progress.ProgressIndicator;
 import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.project.Project;
@@ -49,7 +50,6 @@
 import reactor.core.scheduler.Schedulers;
 import reactor.util.function.Tuple2;
 import rx.Single;
-import rx.exceptions.Exceptions;
 
 import javax.swing.*;
 import java.awt.*;
@@ -58,6 +58,7 @@
 import java.util.List;
 import java.util.*;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Stream;
 
@@ -66,7 +67,6 @@
 public class SignInWindow extends AzureDialogWrapper {
     private static final Logger LOGGER = Logger.getInstance(SignInWindow.class);
     private static final String SIGN_IN_ERROR = "Sign In Error";
-    private static final String USER_CANCEL = "user cancel";
 
     private JPanel contentPane;
 
@@ -215,21 +215,26 @@ public Single<AuthMethodDetails> login() {
     }
 
     private static AuthMethodDetails checkCanceled(ProgressIndicator indicator, Mono<? extends AuthMethodDetails> mono) {
-        CompletableFuture<AuthMethodDetails> loginFuture = new CompletableFuture<>();
+        CompletableFuture<? extends AuthMethodDetails> future = mono.toFuture();
         Disposable disposable = Flux.interval(Duration.ofSeconds(1)).map(ts -> {
-            if (indicator != null && indicator.isCanceled()) {
-                IllegalStateException e = new IllegalStateException(USER_CANCEL);
-                loginFuture.completeExceptionally(e);
-                throw e;
+            if (indicator != null) {
+                indicator.checkCanceled();
             }
             return 1;
-        }).onErrorResume(e -> Mono.empty()).subscribe();
-        mono.doOnError(loginFuture::completeExceptionally).doOnSuccess(loginFuture::complete).doFinally(e -> disposable.dispose()).subscribe();
+        }).onErrorResume(e -> {
+            if (e instanceof ProcessCanceledException) {
+                future.cancel(true);
+            }
+            return Mono.empty();
+        }).subscribeOn(Schedulers.boundedElastic()).subscribe();
         try {
-            return loginFuture.get();
-        } catch (Throwable ex) {
-            Throwable e = Exceptions.getFinalCause(ex);
-            throw new AzureToolkitRuntimeException("Cannot login due to error: " + ex.getMessage(), ex);
+            return future.get();
+        } catch (CancellationException e) {
+            return null;
+        } catch (Throwable e) {
+            throw new AzureToolkitRuntimeException("Cannot login due to error: " + e.getMessage(), e);
+        } finally {
+            disposable.dispose();
         }
     }
 
@@ -334,25 +339,22 @@ private synchronized AuthMethodDetails doDeviceLogin() {
             final IDeviceLoginUI deviceLoginUI = CommonSettings.getUiFactory().getDeviceLoginUI();
             final AzureAccount az = com.microsoft.azure.toolkit.lib.Azure.az(AzureAccount.class);
             final Account account = az.loginAsync(AuthType.DEVICE_CODE, true).block();
-            Disposable subscribe = account.continueLogin().doOnCancel(() -> {
-                deviceCodeLoginFuture.completeExceptionally(new IllegalStateException("user cancel"));
-            }).doOnSuccess(ac -> {
-                deviceCodeLoginFuture.complete(fromAccountEntity(ac.getEntity()));
-            }).doOnError(deviceCodeLoginFuture::completeExceptionally).doFinally(signal -> {
-                deviceLoginUI.closePrompt();
-            }).subscribe();
-            deviceLoginUI.setDisposable(subscribe);
+
+            CompletableFuture<AuthMethodDetails> future =
+                    account.continueLogin().map(ac -> fromAccountEntity(ac.getEntity())).doFinally(signal -> {
+                        deviceLoginUI.closePrompt();
+                    }).toFuture();
+            deviceLoginUI.setFuture(future);
             if (ApplicationManager.getApplication().isDispatchThread()) {
                 deviceLoginUI.promptDeviceCode(((DeviceCodeAccount) account).getDeviceCode());
             } else {
                 AzureTaskManager.getInstance().runAndWait(() ->
                                                               deviceLoginUI.promptDeviceCode(((DeviceCodeAccount) account).getDeviceCode()));
             }
-            return Mono.fromFuture(deviceCodeLoginFuture).block();
+            return future.get();
 
         } catch (Exception ex) {
-            if (ex instanceof IllegalStateException && USER_CANCEL.equals(ex.getMessage())) {
-            } else {
+            if (!(ex instanceof CancellationException)) {
                 ex.printStackTrace();
                 ErrorWindow.show(project, ex.getMessage(), SIGN_IN_ERROR);
             }

Utils/azuretools-core/src/com/microsoft/azuretools/adauth/IDeviceLoginUI.java

@@ -11,7 +11,8 @@
 import com.microsoft.aad.adal4j.AuthenticationResult;
 import com.microsoft.aad.adal4j.DeviceCode;
 import com.microsoft.azuretools.azurecommons.helpers.Nullable;
-import reactor.core.Disposable;
+
+import java.util.concurrent.Future;
 
 public interface IDeviceLoginUI {
     @Nullable
@@ -26,7 +27,7 @@ default void closePrompt() {
 
     }
 
-    default void setDisposable(Disposable disposable) {
+    default void setFuture(Future future) {
 
     }