Skip to content

Commit aefc6d0

Browse files
kannanjgithubkannanjgithub
committed
Implement new draining logic as per updates in grfc.
1 parent 3390e5c commit aefc6d0

3 files changed

Lines changed: 1147 additions & 359 deletions

File tree

‎xds/src/main/java/io/grpc/xds/ExternalProcessorClientInterceptor.java‎

Lines changed: 130 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,14 @@ private static class DataPlaneClientCall
316316
final AtomicBoolean pendingHalfClose = new AtomicBoolean(false);
317317
final AtomicBoolean bodyMessageSentToExtProc = new AtomicBoolean(false);
318318
private final AtomicBoolean downstreamCancelled = new AtomicBoolean(false);
319+
final AtomicBoolean requestDraining = new AtomicBoolean(false);
320+
final AtomicBoolean responseDraining = new AtomicBoolean(false);
321+
final AtomicBoolean requestDrainComplete = new AtomicBoolean(false);
322+
final AtomicBoolean responseDrainComplete = new AtomicBoolean(false);
323+
final AtomicBoolean requestEosSent = new AtomicBoolean(false);
324+
final AtomicBoolean responseTrailersSent = new AtomicBoolean(false);
325+
final AtomicBoolean requestBodySentToExtProc = new AtomicBoolean(false);
326+
final AtomicBoolean responseBodySentToExtProc = new AtomicBoolean(false);
319327

320328
protected DataPlaneClientCall(
321329
DataPlaneDelayedCall<InputStream, InputStream> delayedCall,
@@ -502,10 +510,25 @@ public void onNext(ProcessingResponse response) {
502510
}
503511
}
504512

505-
if (response.getRequestDrain()) {
506-
extProcStreamState.set(ExtProcStreamState.DRAINING);
507-
halfCloseExtProcStream();
508-
activateCall();
513+
if (response.getRequestDrainRequests()) {
514+
if (!requestEosSent.get() && requestDraining.compareAndSet(false, true)) {
515+
activateCall();
516+
sendToExtProc(ProcessingRequest.newBuilder()
517+
.setRequestBody(HttpBody.newBuilder()
518+
.setDrainComplete(true)
519+
.build())
520+
.build());
521+
}
522+
}
523+
if (response.getRequestDrainResponses()) {
524+
if (!responseTrailersSent.get() && responseDraining.compareAndSet(false, true)) {
525+
activateCall();
526+
sendToExtProc(ProcessingRequest.newBuilder()
527+
.setResponseBody(HttpBody.newBuilder()
528+
.setDrainComplete(true)
529+
.build())
530+
.build());
531+
}
509532
}
510533

511534
// 1. Client Headers
@@ -599,6 +622,29 @@ public void onError(Throwable t) {
599622

600623
@Override
601624
public void onCompleted() {
625+
boolean requestDrainRequired = currentProcessingMode.getRequestBodyMode() == ProcessingMode.BodySendMode.GRPC
626+
&& requestBodySentToExtProc.get()
627+
&& !requestDrainComplete.get()
628+
&& !requestEosSent.get();
629+
630+
boolean responseDrainRequired = currentProcessingMode.getResponseBodyMode() == ProcessingMode.BodySendMode.GRPC
631+
&& responseBodySentToExtProc.get()
632+
&& !responseDrainComplete.get()
633+
&& !responseTrailersSent.get();
634+
635+
if (!config.getObservabilityMode() && (requestDrainRequired || responseDrainRequired)) {
636+
if (markExtProcStreamFailed(extProcStreamState)) {
637+
synchronized (streamLock) {
638+
extProcClientCallRequestObserver = null;
639+
}
640+
cancelDownstream("External processor stream completed without drain",
641+
Status.INTERNAL.withDescription("External processor stream completed without drain")
642+
.asRuntimeException());
643+
wrappedListener.proceedWithClose();
644+
}
645+
return;
646+
}
647+
602648
if (markExtProcStreamCompleted(extProcStreamState)) {
603649
handleFailOpen(wrappedListener);
604650
}
@@ -719,24 +765,32 @@ private void internalOnError(Throwable t) {
719765
}
720766
}
721767

722-
private void halfCloseExtProcStream() {
723-
synchronized (streamLock) {
724-
if (!extProcStreamState.get().isCompleted() && extProcClientCallRequestObserver != null) {
725-
extProcClientCallRequestObserver.onCompleted();
726-
}
727-
}
728-
}
768+
729769

730770
private void onReadyNotify() {
731771
wrappedListener.onReadyNotify();
732772
}
733773

734-
private boolean isSidecarReady() {
774+
private boolean isRequestSidecarReady() {
775+
ExtProcStreamState state = extProcStreamState.get();
776+
if (state.isCompleted()) {
777+
return true;
778+
}
779+
if (requestDraining.get()) {
780+
return false;
781+
}
782+
synchronized (streamLock) {
783+
ClientCallStreamObserver<ProcessingRequest> observer = extProcClientCallRequestObserver;
784+
return observer != null && observer.isReady();
785+
}
786+
}
787+
788+
private boolean isResponseSidecarReady() {
735789
ExtProcStreamState state = extProcStreamState.get();
736790
if (state.isCompleted()) {
737791
return true;
738792
}
739-
if (state.isDraining()) {
793+
if (responseDraining.get()) {
740794
return false;
741795
}
742796
synchronized (streamLock) {
@@ -756,7 +810,7 @@ public boolean isReady() {
756810
if (dataPlaneCallState.get() == DataPlaneCallState.IDLE && !config.getObservabilityMode()) {
757811
return false;
758812
}
759-
boolean sidecarReady = isSidecarReady();
813+
boolean sidecarReady = isRequestSidecarReady();
760814
if (config.getObservabilityMode()) {
761815
return super.isReady() && sidecarReady;
762816
}
@@ -774,7 +828,7 @@ public void request(int numMessages) {
774828
super.request(numMessages);
775829
return;
776830
}
777-
if (!isSidecarReady()) {
831+
if (!isResponseSidecarReady()) {
778832
pendingRequests.addAndGet(numMessages);
779833
return;
780834
}
@@ -788,19 +842,19 @@ public void sendMessage(InputStream message) {
788842
return;
789843
}
790844

791-
if (passThroughMode.get()) {
845+
if (passThroughMode.get() || requestDrainComplete.get()) {
792846
super.sendMessage(message);
793847
return;
794848
}
795849

796850
synchronized (streamLock) {
797-
if (passThroughMode.get()) {
851+
if (passThroughMode.get() || requestDrainComplete.get()) {
798852
super.sendMessage(message);
799853
return;
800854
}
801855

802856
ExtProcStreamState state = extProcStreamState.get();
803-
if (state.isDraining() || state.isCompleted()) {
857+
if (requestDraining.get() || state.isCompleted()) {
804858
if (currentProcessingMode.getRequestBodyMode() == ProcessingMode.BodySendMode.NONE) {
805859
super.sendMessage(message);
806860
return;
@@ -829,6 +883,7 @@ public void sendMessage(InputStream message) {
829883
.setEndOfStream(false)
830884
.build())
831885
.build());
886+
requestBodySentToExtProc.set(true);
832887
bodyMessageSentToExtProc.set(true);
833888

834889
if (config.getObservabilityMode()) {
@@ -851,7 +906,7 @@ private void proceedWithHalfClose() {
851906
@Override
852907
public void halfClose() {
853908
clientHalfCloseStartNanos = System.nanoTime();
854-
if (passThroughMode.get()) {
909+
if (passThroughMode.get() || requestDrainComplete.get()) {
855910
if (requestSideClosed.compareAndSet(false, true)) {
856911
proceedWithHalfClose();
857912
}
@@ -869,7 +924,7 @@ public void halfClose() {
869924
return;
870925
}
871926

872-
if (extProcStreamState.get().isDraining()) {
927+
if (requestDraining.get()) {
873928
boolean canProceed = false;
874929
synchronized (streamLock) {
875930
if (currentProcessingMode.getRequestBodyMode() == ProcessingMode.BodySendMode.NONE
@@ -898,6 +953,7 @@ public void halfClose() {
898953
.setEndOfStreamWithoutMessage(true)
899954
.build())
900955
.build());
956+
requestEosSent.set(true);
901957
}
902958

903959
private void cancelDownstream(@Nullable String message, @Nullable Throwable cause) {
@@ -925,7 +981,9 @@ private void handleRequestBodyResponse(BodyResponse bodyResponse) {
925981
BodyMutation mutation = bodyResponse.getResponse().getBodyMutation();
926982
if (mutation.hasStreamedResponse()) {
927983
StreamedBodyResponse streamed = mutation.getStreamedResponse();
928-
if (!streamed.getEndOfStreamWithoutMessage()) {
984+
if (streamed.getDrainComplete()) {
985+
handleRequestDrainComplete();
986+
} else if (!streamed.getEndOfStreamWithoutMessage()) {
929987
super.sendMessage(new KnownLengthInputStream(streamed.getBody()));
930988
}
931989
if (streamed.getEndOfStream() || streamed.getEndOfStreamWithoutMessage()) {
@@ -943,7 +1001,11 @@ private void handleResponseBodyResponse(
9431001
BodyMutation mutation = bodyResponse.getResponse().getBodyMutation();
9441002
if (mutation.hasStreamedResponse()) {
9451003
StreamedBodyResponse streamed = mutation.getStreamedResponse();
946-
listener.onExternalBody(streamed.getBody());
1004+
if (streamed.getDrainComplete()) {
1005+
handleResponseDrainComplete();
1006+
} else {
1007+
listener.onExternalBody(streamed.getBody());
1008+
}
9471009
}
9481010
}
9491011
}
@@ -991,6 +1053,20 @@ private void drainPendingDrainingMessages() {
9911053
}
9921054
}
9931055

1056+
private void handleRequestDrainComplete() {
1057+
if (requestDraining.compareAndSet(true, false)) {
1058+
requestDrainComplete.set(true);
1059+
drainPendingDrainingMessages();
1060+
}
1061+
}
1062+
1063+
private void handleResponseDrainComplete() {
1064+
if (responseDraining.compareAndSet(true, false)) {
1065+
responseDrainComplete.set(true);
1066+
wrappedListener.unblockAfterResponseDrain();
1067+
}
1068+
}
1069+
9941070
private void handleFailOpen(DataPlaneListener listener) {
9951071
activateCall();
9961072
drainPendingRequests();
@@ -1111,7 +1187,8 @@ public void onHeaders(Metadata headers) {
11111187
|| dataPlaneClientCall.getCurrentProcessingMode().getResponseHeaderMode()
11121188
== ProcessingMode.HeaderSendMode.DEFAULT;
11131189

1114-
if (dataPlaneClientCall.getExtProcStreamState().get().isDraining() && sendResponseHeaders) {
1190+
1191+
if (dataPlaneClientCall.responseDraining.get() && sendResponseHeaders) {
11151192
this.savedHeaders = headers;
11161193
return;
11171194
}
@@ -1144,7 +1221,7 @@ public void onMessage(InputStream message) {
11441221
return;
11451222
}
11461223

1147-
boolean checkDrain = dataPlaneClientCall.getExtProcStreamState().get().isDraining()
1224+
boolean checkDrain = dataPlaneClientCall.responseDraining.get()
11481225
&& dataPlaneClientCall.getCurrentProcessingMode().getResponseBodyMode()
11491226
== ProcessingMode.BodySendMode.GRPC;
11501227

@@ -1174,6 +1251,7 @@ public void onMessage(InputStream message) {
11741251
try {
11751252
ByteString bodyByteString = ByteString.readFrom(message);
11761253
sendResponseBodyToExtProc(bodyByteString, false);
1254+
dataPlaneClientCall.responseBodySentToExtProc.set(true);
11771255
dataPlaneClientCall.bodyMessageSentToExtProc.set(true);
11781256

11791257
if (dataPlaneClientCall.getConfig().getObservabilityMode()) {
@@ -1198,8 +1276,13 @@ public void onClose(Status status, Metadata trailers) {
11981276
&& (!dataPlaneClientCall.getConfig().getFailureModeAllow()
11991277
|| dataPlaneClientCall.bodyMessageSentToExtProc.get())) {
12001278
if (markDataPlaneCallClosed(dataPlaneClientCall.dataPlaneCallState)) {
1201-
proceedWithClose(Status.INTERNAL.withDescription("External processor stream failed")
1202-
.withCause(status.getCause()), new Metadata());
1279+
Status finalStatus = Status.INTERNAL.withCause(status.getCause());
1280+
if (status.getDescription() != null) {
1281+
finalStatus = finalStatus.withDescription(status.getDescription());
1282+
} else {
1283+
finalStatus = finalStatus.withDescription("External processor stream failed");
1284+
}
1285+
proceedWithClose(finalStatus, new Metadata());
12031286
}
12041287
return;
12051288
}
@@ -1225,7 +1308,7 @@ public void onClose(Status status, Metadata trailers) {
12251308
dataPlaneClientCall.getCurrentProcessingMode().getResponseTrailerMode()
12261309
== ProcessingMode.HeaderSendMode.SEND;
12271310

1228-
if (dataPlaneClientCall.getExtProcStreamState().get().isDraining() && sendResponseTrailers) {
1311+
if (dataPlaneClientCall.responseDraining.get() && sendResponseTrailers) {
12291312
return;
12301313
}
12311314

@@ -1245,7 +1328,7 @@ void proceedWithHeaders() {
12451328
proceedWithHeaders(savedHeaders);
12461329
synchronized (savedMessages) {
12471330
savedHeaders = null;
1248-
if (!dataPlaneClientCall.getExtProcStreamState().get().isDraining()) {
1331+
if (!dataPlaneClientCall.responseDraining.get()) {
12491332
InputStream msg;
12501333
while ((msg = savedMessages.poll()) != null) {
12511334
onMessage(msg);
@@ -1302,6 +1385,15 @@ void unblockAfterStreamComplete() {
13021385
proceedWithClose();
13031386
}
13041387

1388+
void unblockAfterResponseDrain() {
1389+
synchronized (savedMessages) {
1390+
inboundPassThrough = true;
1391+
}
1392+
proceedWithHeaders();
1393+
proceedWithSavedMessages();
1394+
proceedWithClose();
1395+
}
1396+
13051397
private void proceedWithSavedMessages() {
13061398
synchronized (savedMessages) {
13071399
InputStream msg;
@@ -1319,6 +1411,14 @@ private void triggerCloseHandshake() {
13191411
return;
13201412
}
13211413

1414+
if (dataPlaneClientCall.responseDrainComplete.get()) {
1415+
proceedWithClose();
1416+
if (!dataPlaneClientCall.getConfig().getObservabilityMode()) {
1417+
dataPlaneClientCall.closeExtProcStream();
1418+
}
1419+
return;
1420+
}
1421+
13221422
boolean sendResponseHeaders =
13231423
dataPlaneClientCall.getCurrentProcessingMode().getResponseHeaderMode()
13241424
== ProcessingMode.HeaderSendMode.SEND
@@ -1340,6 +1440,7 @@ private void triggerCloseHandshake() {
13401440
.setEndOfStream(true)
13411441
.build())
13421442
.build());
1443+
dataPlaneClientCall.responseTrailersSent.set(true);
13431444
} else {
13441445
proceedWithClose();
13451446
if (!dataPlaneClientCall.getConfig().getObservabilityMode()) {
@@ -1356,6 +1457,7 @@ private void triggerCloseHandshake() {
13561457
dataPlaneClientCall.getConfig().getForwardRulesConfig()))
13571458
.build())
13581459
.build());
1460+
dataPlaneClientCall.responseTrailersSent.set(true);
13591461
} else {
13601462
proceedWithClose();
13611463
if (!dataPlaneClientCall.getConfig().getObservabilityMode()) {

0 commit comments

Comments
 (0)