Skip to content

Commit 07a7505

Browse files
rahul2393rahul2393
authored
fix(grpc-gcp): prevent hot-channel skew (#14198)
1 parent 4d93800 commit 07a7505

8 files changed

Lines changed: 1649 additions & 542 deletions

File tree

‎grpc-gcp-java/pom.xml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@
6363
<project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
6464
<site.installationModule>grpc-gcp</site.installationModule>
6565
<api-common.version>2.68.0-SNAPSHOT</api-common.version><!-- {x-version-update:api-common:current} -->
66+
<awaitility.version>4.3.0</awaitility.version>
6667
<auto-value.version>1.11.0</auto-value.version>
6768
<error-prone-annotations.version>2.50.0</error-prone-annotations.version>
6869
<google-http-client.version>2.2.0</google-http-client.version>
@@ -172,6 +173,12 @@
172173
<version>${junit.version}</version>
173174
<scope>test</scope>
174175
</dependency>
176+
<dependency>
177+
<groupId>org.awaitility</groupId>
178+
<artifactId>awaitility</artifactId>
179+
<version>${awaitility.version}</version>
180+
<scope>test</scope>
181+
</dependency>
175182
<dependency>
176183
<groupId>com.google.truth</groupId>
177184
<artifactId>truth</artifactId>

‎grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpClientCall.java‎

Lines changed: 51 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import com.google.cloud.grpc.proto.AffinityConfig;
2020
import com.google.common.base.MoreObjects;
21+
import com.google.common.base.Preconditions;
2122
import io.grpc.Attributes;
2223
import io.grpc.CallOptions;
2324
import io.grpc.ClientCall;
@@ -30,7 +31,7 @@
3031
import java.util.Collections;
3132
import java.util.List;
3233
import java.util.Queue;
33-
import java.util.concurrent.atomic.AtomicBoolean;
34+
import java.util.concurrent.atomic.AtomicInteger;
3435
import javax.annotation.Nullable;
3536
import javax.annotation.concurrent.GuardedBy;
3637

@@ -52,7 +53,8 @@ public class GcpClientCall<ReqT, RespT> extends ClientCall<ReqT, RespT> {
5253
private ClientCall<ReqT, RespT> delegateCall = null;
5354
private List<String> keys = null;
5455
private boolean received = false;
55-
private final AtomicBoolean decremented = new AtomicBoolean(false);
56+
// 0 = not counted, 1 = counted, 2 = finished.
57+
private final AtomicInteger countState = new AtomicInteger();
5658

5759
@GuardedBy("this")
5860
private final Queue<Runnable> calls = new ArrayDeque<>();
@@ -123,16 +125,23 @@ public void sendMessage(ReqT message) {
123125
delegateChannelRef = delegateChannel.getChannelRef(key);
124126
}
125127
delegateChannelRef.activeStreamsCountIncr();
128+
Preconditions.checkState(countState.compareAndSet(0, 1));
126129

127130
// Create the client call and do the previous operations.
128131
CallOptions callOptionsWithChannelId =
129132
callOptions.withOption(GcpManagedChannel.CHANNEL_ID_KEY, delegateChannelRef.getId());
130-
delegateCall =
131-
delegateChannelRef.getChannel().newCall(methodDescriptor, callOptionsWithChannelId);
132-
for (Runnable call : calls) {
133-
call.run();
133+
try {
134+
delegateCall =
135+
delegateChannelRef.getChannel().newCall(methodDescriptor, callOptionsWithChannelId);
136+
for (Runnable call : calls) {
137+
call.run();
138+
}
139+
} catch (RuntimeException | Error failure) {
140+
finishCount(Status.fromThrowable(failure), true);
141+
throw failure;
142+
} finally {
143+
calls.clear();
134144
}
135-
calls.clear();
136145
started = true;
137146
}
138147
}
@@ -165,12 +174,16 @@ public String toString() {
165174
}
166175

167176
private void checkedCancel(@Nullable String message, @Nullable Throwable cause) {
168-
if (!decremented.getAndSet(true)) {
169-
delegateChannelRef.activeStreamsCountDecr(startNanos, Status.CANCELLED, true);
170-
}
177+
finishCount(Status.CANCELLED, true);
171178
delegateCall.cancel(message, cause);
172179
}
173180

181+
private void finishCount(Status status, boolean cancelled) {
182+
if (countState.compareAndSet(1, 2)) {
183+
delegateChannelRef.activeStreamsCountDecr(startNanos, status, cancelled);
184+
}
185+
}
186+
174187
private void checkSendMessage(Runnable call) {
175188
synchronized (this) {
176189
if (started) {
@@ -188,9 +201,7 @@ private Listener<RespT> getListener(final Listener<RespT> responseListener) {
188201
// Decrement the stream number by one when the call is closed.
189202
@Override
190203
public void onClose(Status status, Metadata trailers) {
191-
if (!decremented.getAndSet(true)) {
192-
delegateChannelRef.activeStreamsCountDecr(startNanos, status, false);
193-
}
204+
finishCount(status, false);
194205
// If the operation completed successfully, bind/unbind the affinity key.
195206
if (keys != null && status.getCode() == Status.Code.OK) {
196207
if (affinity.getCommand() == AffinityConfig.Command.UNBIND) {
@@ -231,9 +242,10 @@ public static class SimpleGcpClientCall<ReqT, RespT> extends ForwardingClientCal
231242
private final ClientCall<ReqT, RespT> delegateCall;
232243
@Nullable private final String affinityKey;
233244
private final boolean unbindOnComplete;
234-
private long startNanos = 0;
245+
private long startNanos;
235246

236-
private final AtomicBoolean decremented = new AtomicBoolean(false);
247+
// 0 = not counted, 1 = counted, 2 = finished.
248+
private final AtomicInteger countState = new AtomicInteger();
237249

238250
protected SimpleGcpClientCall(
239251
GcpManagedChannel delegateChannel,
@@ -247,8 +259,16 @@ protected SimpleGcpClientCall(
247259
// Set the actual channel ID in callOptions so downstream interceptors can access it.
248260
CallOptions callOptionsWithChannelId =
249261
callOptions.withOption(GcpManagedChannel.CHANNEL_ID_KEY, channelRef.getId());
250-
this.delegateCall =
251-
channelRef.getChannel().newCall(methodDescriptor, callOptionsWithChannelId);
262+
startNanos = System.nanoTime();
263+
channelRef.activeStreamsCountIncr();
264+
Preconditions.checkState(countState.compareAndSet(0, 1));
265+
try {
266+
this.delegateCall =
267+
channelRef.getChannel().newCall(methodDescriptor, callOptionsWithChannelId);
268+
} catch (RuntimeException | Error failure) {
269+
finishCount(Status.fromThrowable(failure), true);
270+
throw failure;
271+
}
252272
}
253273

254274
@Override
@@ -258,16 +278,12 @@ protected ClientCall<ReqT, RespT> delegate() {
258278

259279
@Override
260280
public void start(Listener<RespT> responseListener, Metadata headers) {
261-
startNanos = System.nanoTime();
262-
263281
Listener<RespT> listener =
264282
new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(
265283
responseListener) {
266284
@Override
267285
public void onClose(Status status, Metadata trailers) {
268-
if (!decremented.getAndSet(true)) {
269-
channelRef.activeStreamsCountDecr(startNanos, status, false);
270-
}
286+
finishCount(status, false);
271287
// Unbind the affinity key when the caller explicitly requests it
272288
// (e.g., on terminal RPCs like Commit or Rollback) to prevent
273289
// unbounded growth of the affinity map.
@@ -284,20 +300,28 @@ public void onMessage(RespT message) {
284300
}
285301
};
286302

287-
channelRef.activeStreamsCountIncr();
288-
delegateCall.start(listener, headers);
303+
try {
304+
delegateCall.start(listener, headers);
305+
} catch (RuntimeException | Error failure) {
306+
finishCount(Status.fromThrowable(failure), true);
307+
throw failure;
308+
}
289309
}
290310

291311
@Override
292312
public void cancel(String message, Throwable cause) {
293-
if (!decremented.getAndSet(true)) {
294-
channelRef.activeStreamsCountDecr(startNanos, Status.CANCELLED, true);
295-
}
313+
finishCount(Status.CANCELLED, true);
296314
// Always unbind on cancel — the transaction is being abandoned.
297315
if (affinityKey != null) {
298316
delegateChannel.unbind(Collections.singletonList(affinityKey));
299317
}
300318
delegateCall.cancel(message, cause);
301319
}
320+
321+
private void finishCount(Status status, boolean cancelled) {
322+
if (countState.compareAndSet(1, 2)) {
323+
channelRef.activeStreamsCountDecr(startNanos, status, cancelled);
324+
}
325+
}
302326
}
303327
}

0 commit comments

Comments
 (0)