stub: use the closedTrailers in StatusException (#12259)

This commit is contained in:
MV Shiva 2025-08-06 12:24:33 +05:30 committed by GitHub
parent a40c8cf5a4
commit 7040417eee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 32 additions and 21 deletions

View File

@ -29,6 +29,7 @@ import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue; import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException; import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;
import java.util.logging.Level; import java.util.logging.Level;
import java.util.logging.Logger; import java.util.logging.Logger;
@ -69,7 +70,7 @@ public final class BlockingClientCall<ReqT, RespT> {
private final ThreadSafeThreadlessExecutor executor; private final ThreadSafeThreadlessExecutor executor;
private boolean writeClosed; private boolean writeClosed;
private volatile Status closedStatus; // null if not closed private AtomicReference<CloseState> closeState = new AtomicReference<>();
BlockingClientCall(ClientCall<ReqT, RespT> call, ThreadSafeThreadlessExecutor executor) { BlockingClientCall(ClientCall<ReqT, RespT> call, ThreadSafeThreadlessExecutor executor) {
this.call = call; this.call = call;
@ -120,22 +121,22 @@ public final class BlockingClientCall<ReqT, RespT> {
logger.finer("Client Blocking read had value: " + bufferedValue); logger.finer("Client Blocking read had value: " + bufferedValue);
} }
Status currentClosedStatus; CloseState currentCloseState;
if (bufferedValue != null) { if (bufferedValue != null) {
call.request(1); call.request(1);
return bufferedValue; return bufferedValue;
} else if ((currentClosedStatus = closedStatus) == null) { } else if ((currentCloseState = closeState.get()) == null) {
throw new IllegalStateException( throw new IllegalStateException(
"The message disappeared... are you reading from multiple threads?"); "The message disappeared... are you reading from multiple threads?");
} else if (!currentClosedStatus.isOk()) { } else if (!currentCloseState.status.isOk()) {
throw currentClosedStatus.asException(); throw currentCloseState.status.asException(currentCloseState.trailers);
} else { } else {
return null; return null;
} }
} }
boolean skipWaitingForRead() { boolean skipWaitingForRead() {
return closedStatus != null || !buffer.isEmpty(); return closeState.get() != null || !buffer.isEmpty();
} }
/** /**
@ -148,11 +149,11 @@ public final class BlockingClientCall<ReqT, RespT> {
* @throws StatusException If the stream was closed in an error state * @throws StatusException If the stream was closed in an error state
*/ */
public boolean hasNext() throws InterruptedException, StatusException { public boolean hasNext() throws InterruptedException, StatusException {
executor.waitAndDrain((x) -> !x.buffer.isEmpty() || x.closedStatus != null, this); executor.waitAndDrain((x) -> !x.buffer.isEmpty() || x.closeState.get() != null, this);
Status currentClosedStatus = closedStatus; CloseState currentCloseState = closeState.get();
if (currentClosedStatus != null && !currentClosedStatus.isOk()) { if (currentCloseState != null && !currentCloseState.status.isOk()) {
throw currentClosedStatus.asException(); throw currentCloseState.status.asException(currentCloseState.trailers);
} }
return !buffer.isEmpty(); return !buffer.isEmpty();
@ -221,17 +222,16 @@ public final class BlockingClientCall<ReqT, RespT> {
} }
Predicate<BlockingClientCall<ReqT, RespT>> predicate = Predicate<BlockingClientCall<ReqT, RespT>> predicate =
(x) -> x.call.isReady() || x.closedStatus != null; (x) -> x.call.isReady() || x.closeState.get() != null;
executor.waitAndDrainWithTimeout(waitForever, endNanoTime, predicate, this); executor.waitAndDrainWithTimeout(waitForever, endNanoTime, predicate, this);
Status savedClosedStatus = closedStatus; CloseState savedCloseState = closeState.get();
if (savedClosedStatus == null) { if (savedCloseState == null || savedCloseState.status == null) {
call.sendMessage(request); call.sendMessage(request);
return true; return true;
} else if (savedClosedStatus.isOk()) { } else if (savedCloseState.status.isOk()) {
return false; return false;
} else { } else {
// Propagate any errors returned from the server throw savedCloseState.status.asException(savedCloseState.trailers);
throw savedClosedStatus.asException();
} }
} }
@ -274,7 +274,8 @@ public final class BlockingClientCall<ReqT, RespT> {
@VisibleForTesting @VisibleForTesting
Status getClosedStatus() { Status getClosedStatus() {
drainQuietly(); drainQuietly();
return closedStatus; CloseState state = closeState.get();
return (state == null) ? null : state.status;
} }
/** /**
@ -317,7 +318,7 @@ public final class BlockingClientCall<ReqT, RespT> {
* @return True if writes haven't been closed and the server hasn't closed the stream * @return True if writes haven't been closed and the server hasn't closed the stream
*/ */
private boolean isWriteLegal() { private boolean isWriteLegal() {
return !writeClosed && closedStatus == null; return !writeClosed && closeState.get() == null;
} }
ClientCall.Listener<RespT> getListener() { ClientCall.Listener<RespT> getListener() {
@ -335,15 +336,25 @@ public final class BlockingClientCall<ReqT, RespT> {
private final class QueuingListener extends ClientCall.Listener<RespT> { private final class QueuingListener extends ClientCall.Listener<RespT> {
@Override @Override
public void onMessage(RespT value) { public void onMessage(RespT value) {
Preconditions.checkState(closedStatus == null, "ClientCall already closed"); Preconditions.checkState(closeState.get() == null, "ClientCall already closed");
buffer.add(value); buffer.add(value);
} }
@Override @Override
public void onClose(Status status, Metadata trailers) { public void onClose(Status status, Metadata trailers) {
Preconditions.checkState(closedStatus == null, "ClientCall already closed"); CloseState newCloseState = new CloseState(status, trailers);
closedStatus = status; boolean wasSet = closeState.compareAndSet(null, newCloseState);
Preconditions.checkState(wasSet, "ClientCall already closed");
} }
} }
private static final class CloseState {
final Status status;
final Metadata trailers;
CloseState(Status status, Metadata trailers) {
this.status = Preconditions.checkNotNull(status, "status");
this.trailers = trailers;
}
}
} }