From 18970e6ef310fda38e577de2a30c71c6c44f1e66 Mon Sep 17 00:00:00 2001
From: ZHANG Dapeng <zdapeng@google.com>
Date: Tue, 1 Aug 2017 10:26:02 -0700
Subject: [PATCH] core: fix ConnectivityStateManager is already disabled bug

---
 .../internal/ConnectivityStateManager.java    |  9 +++-
 .../io/grpc/internal/ManagedChannelImpl.java  |  4 +-
 .../ConnectivityStateManagerTest.java         | 45 +++++++++++++++++++
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/io/grpc/internal/ConnectivityStateManager.java b/core/src/main/java/io/grpc/internal/ConnectivityStateManager.java
index 62165f6a41..58afaff1bf 100644
--- a/core/src/main/java/io/grpc/internal/ConnectivityStateManager.java
+++ b/core/src/main/java/io/grpc/internal/ConnectivityStateManager.java
@@ -63,7 +63,7 @@ final class ConnectivityStateManager {
    */
   void gotoState(@Nonnull ConnectivityState newState) {
     checkNotNull(newState, "newState");
-    checkState(state != null, "ConnectivityStateManager is already disabled");
+    checkState(!isDisabled(), "ConnectivityStateManager is already disabled");
     gotoNullableState(newState);
   }
 
@@ -102,6 +102,13 @@ final class ConnectivityStateManager {
     gotoNullableState(null);
   }
 
+  /**
+   * This method is threadsafe.
+   */
+  boolean isDisabled() {
+    return state == null;
+  }
+
   private static final class Listener {
     final Runnable callback;
     final Executor executor;
diff --git a/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java b/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
index bbfbdb0ea3..7294d60e8b 100644
--- a/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
+++ b/core/src/main/java/io/grpc/internal/ManagedChannelImpl.java
@@ -462,7 +462,9 @@ public final class ManagedChannelImpl extends ManagedChannel implements WithLogI
     channelExecutor.executeLater(new Runnable() {
       @Override
       public void run() {
-        channelStateManager.gotoState(SHUTDOWN);
+        if (!channelStateManager.isDisabled()) {
+          channelStateManager.gotoState(SHUTDOWN);
+        }
       }
     });
 
diff --git a/core/src/test/java/io/grpc/internal/ConnectivityStateManagerTest.java b/core/src/test/java/io/grpc/internal/ConnectivityStateManagerTest.java
index 66adb6f0b1..04d87ca464 100644
--- a/core/src/test/java/io/grpc/internal/ConnectivityStateManagerTest.java
+++ b/core/src/test/java/io/grpc/internal/ConnectivityStateManagerTest.java
@@ -17,12 +17,15 @@
 package io.grpc.internal;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import com.google.common.util.concurrent.MoreExecutors;
 import io.grpc.ConnectivityState;
 import java.util.LinkedList;
 import java.util.concurrent.Executor;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
@@ -31,6 +34,9 @@ import org.junit.runners.JUnit4;
  */
 @RunWith(JUnit4.class)
 public class ConnectivityStateManagerTest {
+  @Rule
+  public final ExpectedException thrown = ExpectedException.none();
+
   private final FakeClock executor = new FakeClock();
   private final ConnectivityStateManager state = new ConnectivityStateManager();
   private final LinkedList<ConnectivityState> sink = new LinkedList<ConnectivityState>();
@@ -223,4 +229,43 @@ public class ConnectivityStateManagerTest {
     assertEquals(1, sink.size());
     assertEquals(ConnectivityState.READY, sink.poll());
   }
+
+  @Test
+  public void disable() {
+    state.disable();
+    assertTrue(state.isDisabled());
+
+    thrown.expect(UnsupportedOperationException.class);
+    thrown.expectMessage("Channel state API is not implemented");
+    state.getState();
+  }
+
+  @Test
+  public void disableThenDisable() {
+    state.disable();
+    state.disable();
+    assertTrue(state.isDisabled());
+
+    thrown.expect(UnsupportedOperationException.class);
+    thrown.expectMessage("Channel state API is not implemented");
+    state.getState();
+  }
+
+  @Test
+  public void disableThenGotoReady() {
+    state.disable();
+
+    thrown.expect(IllegalStateException.class);
+    thrown.expectMessage("ConnectivityStateManager is already disabled");
+    state.gotoState(ConnectivityState.READY);
+  }
+
+  @Test
+  public void shutdownThenReady() {
+    state.gotoState(ConnectivityState.SHUTDOWN);
+    assertEquals(ConnectivityState.SHUTDOWN, state.getState());
+
+    state.gotoState(ConnectivityState.READY);
+    assertEquals(ConnectivityState.SHUTDOWN, state.getState());
+  }
 }