m-trieu commented on code in PR #31902:
URL: https://github.com/apache/beam/pull/31902#discussion_r1760700251
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/AbstractWindmillStream.java:
##########
@@ -278,24 +329,72 @@ public String backendWorkerToken() {
}
@Override
- public void shutdown() {
+ public final void shutdown() {
+ // Don't lock here as isShutdown checks are used in the stream to free
blocked
+ // threads or as exit conditions to loops.
if (isShutdown.compareAndSet(false, true)) {
requestObserver()
.onError(new WindmillStreamShutdownException("Explicit call to
shutdown stream."));
+ shutdownInternal();
}
}
- private void setLastError(String error) {
- lastError.set(error);
- lastErrorTime.set(DateTime.now());
+ private void recordRestartReason(String error) {
+ lastRestartReason.set(error);
+ lastRestartTime.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("this")
+ private volatile @Nullable StreamObserver<RequestT>
delegateRequestObserver;
+
+ private ResettableRequestObserver(Supplier<StreamObserver<RequestT>>
requestObserverSupplier) {
+ this.requestObserverSupplier = requestObserverSupplier;
+ this.delegateRequestObserver = null;
+ }
+
+ private synchronized StreamObserver<RequestT> delegate() {
+ return Preconditions.checkNotNull(
+ delegateRequestObserver,
+ "requestObserver cannot be null. Missing a call to startStream() to
initialize.");
+ }
+
+ private synchronized void reset() {
+ delegateRequestObserver = requestObserverSupplier.get();
+ }
+
+ @Override
+ public void onNext(RequestT requestT) {
+ delegate().onNext(requestT);
+ }
+
+ @Override
+ public void onError(Throwable throwable) {
+ delegate().onError(throwable);
+ }
+
+ @Override
+ public synchronized void onCompleted() {
Review Comment:
done
--
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]
