m-trieu commented on code in PR #31902:
URL: https://github.com/apache/beam/pull/31902#discussion_r1730027564
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/AbstractWindmillStream.java:
##########
@@ -290,12 +333,83 @@ private void setLastError(String error) {
lastErrorTime.set(DateTime.now());
}
+ protected abstract void shutdownInternal();
+
public static class WindmillStreamShutdownException extends RuntimeException
{
public WindmillStreamShutdownException(String message) {
super(message);
}
}
+ /**
+ * Request observer that allows resetting its internal delegate using the
given {@link
+ * #requestObserverSupplier}.
+ */
+ @ThreadSafe
+ private static class ResettableRequestObserver<RequestT> implements
StreamObserver<RequestT> {
+ private final Supplier<StreamObserver<RequestT>> requestObserverSupplier;
+
+ @GuardedBy("delegateRequestObserver")
+ private final AtomicReference<StreamObserver<RequestT>>
delegateRequestObserver;
+
+ @GuardedBy("delegateRequestObserver")
+ /* Indicates if onCompleted() has been called for the current
delegateRequestObserver instance.
+ Reset to false when reset() is called. */
+ private volatile boolean isClosed;
+
+ private ResettableRequestObserver(Supplier<StreamObserver<RequestT>>
requestObserverSupplier) {
+ this.requestObserverSupplier = requestObserverSupplier;
+ this.delegateRequestObserver = new AtomicReference<>();
+ this.isClosed = false;
+ }
+
+ private StreamObserver<RequestT> delegate() {
+ synchronized (delegateRequestObserver) {
+ if (delegateRequestObserver.get() == null) {
+ throw new NullPointerException(
+ "requestObserver cannot be null. Missing a call to startStream()
to initialize.");
+ }
+
+ return delegateRequestObserver.get();
+ }
+ }
+
+ private void reset() {
+ synchronized (delegateRequestObserver) {
+ delegateRequestObserver.set(requestObserverSupplier.get());
+ isClosed = false;
+ }
+ }
+
+ @Override
+ public void onNext(RequestT requestT) {
+ delegate().onNext(requestT);
+ }
+
+ @Override
+ public void onError(Throwable throwable) {
+ delegate().onError(throwable);
+ }
+
+ @Override
+ public void onCompleted() {
+ synchronized (delegateRequestObserver) {
+ if (!isClosed) {
+ // onCompleted() can only be called once for each StreamObserver
instance, or else an
+ // IllegalStateException is thrown.
+ delegate().onCompleted();
+ isClosed = true;
+ }
+ }
+ }
+
+ public boolean isClosed() {
Review Comment:
removed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
