|
4 | 4 | import io.grpc.ConnectivityState; |
5 | 5 | import io.grpc.ManagedChannel; |
6 | 6 | import java.util.concurrent.CountDownLatch; |
| 7 | +import java.util.concurrent.Executors; |
| 8 | +import java.util.concurrent.ScheduledFuture; |
7 | 9 | import java.util.concurrent.TimeUnit; |
8 | 10 | import lombok.extern.slf4j.Slf4j; |
9 | 11 |
|
@@ -32,65 +34,58 @@ public static void monitorChannelState( |
32 | 34 | ConnectivityState currentState = channel.getState(true); |
33 | 35 | log.info("Channel state changed to: {}", currentState); |
34 | 36 | if (currentState == ConnectivityState.READY) { |
35 | | - onConnectionReady.run(); |
| 37 | + if (onConnectionReady != null) { |
| 38 | + onConnectionReady.run(); |
| 39 | + } else { |
| 40 | + log.debug("onConnectionReady is null"); |
| 41 | + } |
36 | 42 | } else if (currentState == ConnectivityState.TRANSIENT_FAILURE |
37 | 43 | || currentState == ConnectivityState.SHUTDOWN) { |
38 | | - onConnectionLost.run(); |
| 44 | + if (onConnectionLost != null) { |
| 45 | + onConnectionLost.run(); |
| 46 | + } else { |
| 47 | + log.debug("onConnectionLost is null"); |
| 48 | + } |
39 | 49 | } |
40 | 50 | // Re-register the state monitor to watch for the next state transition. |
41 | 51 | monitorChannelState(currentState, channel, onConnectionReady, onConnectionLost); |
42 | 52 | }); |
43 | 53 | } |
44 | 54 |
|
45 | 55 | /** |
46 | | - * Waits for the channel to reach a desired state within a specified timeout period. |
| 56 | + * Waits for the channel to reach the desired connectivity state within the specified timeout. |
47 | 57 | * |
48 | | - * @param channel the ManagedChannel to monitor. |
49 | | - * @param desiredState the ConnectivityState to wait for. |
50 | | - * @param connectCallback callback invoked when the desired state is reached. |
51 | | - * @param timeout the maximum amount of time to wait. |
52 | | - * @param unit the time unit of the timeout. |
53 | | - * @throws InterruptedException if the current thread is interrupted while waiting. |
| 58 | + * @param desiredState the desired {@link ConnectivityState} to wait for |
| 59 | + * @param channel the {@link ManagedChannel} to monitor |
| 60 | + * @param connectCallback the {@link Runnable} to execute when the desired state is reached |
| 61 | + * @param timeout the maximum time to wait |
| 62 | + * @param unit the time unit of the timeout argument |
| 63 | + * @throws InterruptedException if the current thread is interrupted while waiting |
| 64 | + * @throws GeneralError if the desired state is not reached within the timeout |
54 | 65 | */ |
55 | 66 | public static void waitForDesiredState( |
56 | | - ManagedChannel channel, |
57 | 67 | ConnectivityState desiredState, |
58 | | - Runnable connectCallback, |
59 | | - long timeout, |
60 | | - TimeUnit unit) |
61 | | - throws InterruptedException { |
62 | | - waitForDesiredState(channel, desiredState, connectCallback, new CountDownLatch(1), timeout, unit); |
63 | | - } |
64 | | - |
65 | | - private static void waitForDesiredState( |
66 | 68 | ManagedChannel channel, |
67 | | - ConnectivityState desiredState, |
68 | 69 | Runnable connectCallback, |
69 | | - CountDownLatch latch, |
70 | 70 | long timeout, |
71 | 71 | TimeUnit unit) |
72 | 72 | throws InterruptedException { |
73 | | - channel.notifyWhenStateChanged(ConnectivityState.SHUTDOWN, () -> { |
74 | | - try { |
75 | | - ConnectivityState state = channel.getState(true); |
76 | | - log.debug("Channel state changed to: {}", state); |
| 73 | + CountDownLatch latch = new CountDownLatch(1); |
77 | 74 |
|
78 | | - if (state == desiredState) { |
79 | | - connectCallback.run(); |
80 | | - latch.countDown(); |
81 | | - return; |
82 | | - } |
83 | | - waitForDesiredState(channel, desiredState, connectCallback, latch, timeout, unit); |
84 | | - } catch (InterruptedException e) { |
85 | | - Thread.currentThread().interrupt(); |
86 | | - log.error("Thread interrupted while waiting for desired state", e); |
87 | | - } catch (Exception e) { |
88 | | - log.error("Error occurred while waiting for desired state", e); |
| 75 | + Runnable waitForStateTask = () -> { |
| 76 | + ConnectivityState currentState = channel.getState(true); |
| 77 | + if (currentState == desiredState) { |
| 78 | + connectCallback.run(); |
| 79 | + latch.countDown(); |
89 | 80 | } |
90 | | - }); |
| 81 | + }; |
| 82 | + |
| 83 | + ScheduledFuture<?> scheduledFuture = Executors.newSingleThreadScheduledExecutor() |
| 84 | + .scheduleWithFixedDelay(waitForStateTask, 0, 100, TimeUnit.MILLISECONDS); |
91 | 85 |
|
92 | | - // Await the latch or timeout for the state change |
93 | | - if (!latch.await(timeout, unit)) { |
| 86 | + boolean success = latch.await(timeout, unit); |
| 87 | + scheduledFuture.cancel(true); |
| 88 | + if (!success) { |
94 | 89 | throw new GeneralError(String.format( |
95 | 90 | "Deadline exceeded. Condition did not complete within the %d " + "deadline", timeout)); |
96 | 91 | } |
|
0 commit comments