BarrageSession Subscription/Snapshot Methods now Return Future<Table>

extensions/barrage/src/main/java/io/deephaven/extensions/barrage/table/BarrageTable.java

@@ -396,6 +396,24 @@ private synchronized void realRefresh() {
         }
     }
 
+    public void unsubscribe() {
+        unsubscribed = true;
+
+        if (!isRefreshing()) {
+            try {
+                realRefresh();
+            } catch (Throwable err) {
+                if (viewportChangedCallback != null) {
+                    viewportChangedCallback.onError(err);
+                    viewportChangedCallback = null;
+                }
+                throw err;
+            }
+        } else {
+            doWakeup();
+        }
+    }
+
     private void cleanup() {
         unsubscribed = true;
         if (stats != null) {

extensions/barrage/src/main/java/io/deephaven/extensions/barrage/util/GrpcUtil.java

@@ -9,7 +9,10 @@
 import io.deephaven.internal.log.LoggerFactory;
 import io.deephaven.util.function.ThrowingRunnable;
 import io.grpc.StatusRuntimeException;
+import io.grpc.stub.ClientCallStreamObserver;
 import io.grpc.stub.StreamObserver;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.UUID;
 
@@ -37,7 +40,7 @@ public static StatusRuntimeException securelyWrapError(final Logger log, final T
      * @param observer the stream that will be used in the runnable
      * @param runner the runnable to execute safely
      */
-    public static void safelyExecuteLocked(final StreamObserver<?> observer,
+    private static void safelyExecuteLocked(final StreamObserver<?> observer,
             final ThrowingRunnable<Exception> runner) {
         try {
             // noinspection SynchronizationOnLocalVariableOrMethodParameter
@@ -87,22 +90,46 @@ public static void safelyComplete(StreamObserver<?> observer) {
     /**
      * Writes an error to the observer in a try/catch block to minimize damage caused by failing observer call.
      * <p>
-     * </p>
      * This will always synchronize on the observer to ensure thread safety when interacting with the grpc response
      * stream.
      */
-    public static void safelyError(final StreamObserver<?> observer, final Code statusCode, final String msg) {
+    public static void safelyError(
+            @NotNull final StreamObserver<?> observer,
+            final Code statusCode,
+            @NotNull final String msg) {
         safelyError(observer, Exceptions.statusRuntimeException(statusCode, msg));
     }
 
     /**
      * Writes an error to the observer in a try/catch block to minimize damage caused by failing observer call.
      * <p>
-     * </p>
      * This will always synchronize on the observer to ensure thread safety when interacting with the grpc response
      * stream.
      */
-    public static void safelyError(final StreamObserver<?> observer, StatusRuntimeException exception) {
+    public static void safelyError(
+            @NotNull final StreamObserver<?> observer,
+            @NotNull final StatusRuntimeException exception) {
         safelyExecuteLocked(observer, () -> observer.onError(exception));
     }
+
+    /**
+     * Cancels the observer in a try/catch block to minimize damage caused by failing observer call.
+     * <p>
+     * This will always synchronize on the observer to ensure thread safety when interacting with the grpc response
+     * stream.
+     * <p>
+     * It is recommended that at least one of {@code message} or {@code cause} to be non-{@code null}, to provide useful
+     * debug information. Both argument being null may log warnings and result in suboptimal performance. Also note that
+     * the provided information will not be sent to the server.
+     *
+     * @param observer the stream that will be used in the runnable
+     * @param message if not {@code null}, will appear as the description of the CANCELLED status
+     * @param cause if not {@code null}, will appear as the cause of the CANCELLED status
+     */
+    public static void safelyCancel(
+            @NotNull final ClientCallStreamObserver<?> observer,
+            @Nullable final String message,
+            @Nullable final Throwable cause) {
+        safelyExecuteLocked(observer, () -> observer.cancel(message, cause));
+    }
 }

java-client/barrage/src/main/java/io/deephaven/client/impl/BarrageSnapshotImpl.java

@@ -37,6 +37,7 @@
 import java.nio.ByteBuffer;
 import java.util.BitSet;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 import java.util.concurrent.locks.Condition;
 
 /**
@@ -64,7 +65,9 @@ public class BarrageSnapshotImpl extends ReferenceCountedLivenessNode implements
     private volatile boolean completed = false;
     private volatile Throwable exceptionWhileCompleting = null;
 
-    private volatile boolean connected = true;
+    private volatile int connected = 1;
+    private static final AtomicIntegerFieldUpdater<BarrageSnapshotImpl> CONNECTED_UPDATER =
+            AtomicIntegerFieldUpdater.newUpdater(BarrageSnapshotImpl.class, "connected");
 
     private boolean prevUsed = false;
 
@@ -126,8 +129,7 @@ public void onNext(final BarrageMessage barrageMessage) {
                 return;
             }
             try (barrageMessage) {
-                final Listener localResultTable = resultTable;
-                if (!connected || localResultTable == null) {
+                if (connected == 0) {
                     return;
                 }
 
@@ -144,7 +146,7 @@ public void onNext(final BarrageMessage barrageMessage) {
 
                 rowsReceived += resultSize;
 
-                localResultTable.handleBarrageMessage(barrageMessage);
+                resultTable.handleBarrageMessage(barrageMessage);
             }
         }
 
@@ -154,17 +156,18 @@ public void onError(final Throwable t) {
                     .append(": Error detected in snapshot: ")
                     .append(t).endl();
 
-            final Listener localResultTable = resultTable;
-            if (!connected || localResultTable == null) {
+            if (connected == 0) {
                 return;
             }
-            localResultTable.handleBarrageError(t);
-            handleDisconnect();
+
+            // this error will always be propagated to our CheckForCompletion#onError callback
+            resultTable.handleBarrageError(t);
         }
 
         @Override
         public void onCompleted() {
-            handleDisconnect();
+            // this will always be propagated to our CheckForCompletion#onClose callback
+            resultTable.unsubscribe();
         }
     }
 
@@ -194,7 +197,7 @@ public BarrageTable partialTable(
             RowSet viewport, BitSet columns, boolean reverseViewport, boolean blockUntilComplete)
             throws InterruptedException {
         synchronized (this) {
-            if (!connected) {
+            if (connected == 0) {
                 throw new UncheckedDeephavenException(this + " is not connected");
             }
 
@@ -274,41 +277,37 @@ protected void destroy() {
         close();
     }
 
-    private void handleDisconnect() {
-        if (!connected) {
-            return;
-        }
+    private synchronized void signalCompletion() {
+        signalCompletion(null);
+    }
 
+    /**
+     * This method will only be invoked once using the CAS of {@code connected} from 1 to 0 as a guard.
+     */
+    private synchronized void signalCompletion(@Nullable final Throwable t) {
         completed = true;
-        signalCompletion();
-        cleanup();
-    }
 
-    private synchronized void signalCompletion() {
+        if (t != null) {
+            exceptionWhileCompleting = t;
+        }
+
         if (completedCondition != null) {
             resultTable.getUpdateGraph().requestSignal(completedCondition);
         }
 
         notifyAll();
+
+        tableHandle.close();
     }
 
     @Override
-    public synchronized void close() {
-        if (!connected) {
+    public void close() {
+        if (!CONNECTED_UPDATER.compareAndSet(this, 1, 0)) {
             return;
         }
 
-        exceptionWhileCompleting = new RequestCancelledException("BarrageSnapshotImpl closed");
-        signalCompletion();
-        GrpcUtil.safelyExecuteLocked(observer, () -> {
-            observer.cancel("BarrageSnapshotImpl closed", exceptionWhileCompleting);
-        });
-        cleanup();
-    }
-
-    private void cleanup() {
-        this.connected = false;
-        this.tableHandle.close();
+        signalCompletion(new RequestCancelledException("BarrageSnapshotImpl closed"));
+        GrpcUtil.safelyCancel(observer, "BarrageSnapshotImpl closed", exceptionWhileCompleting);
     }
 
     @Override
@@ -429,12 +428,17 @@ public boolean viewportChanged(@Nullable RowSet rowSet, @Nullable BitSet columns
 
         @Override
         public void onError(Throwable t) {
-            exceptionWhileCompleting = t;
-            signalCompletion();
+            if (!CONNECTED_UPDATER.compareAndSet(BarrageSnapshotImpl.this, 1, 0)) {
+                return;
+            }
+            signalCompletion(t);
         }
 
         @Override
         public void onClose() {
+            if (!CONNECTED_UPDATER.compareAndSet(BarrageSnapshotImpl.this, 1, 0)) {
+                return;
+            }
             signalCompletion();
         }
     }

java-client/barrage/src/main/java/io/deephaven/client/impl/BarrageSubscriptionImpl.java

@@ -39,6 +39,7 @@
 import java.io.InputStream;
 import java.nio.ByteBuffer;
 import java.util.BitSet;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.ScheduledExecutorService;
 
@@ -80,7 +81,9 @@ public class BarrageSubscriptionImpl extends ReferenceCountedLivenessNode implem
     private boolean isSnapshot;
 
     private volatile Condition completedCondition;
-    private volatile boolean completed;
+    private volatile int completed;
+    private static final AtomicIntegerFieldUpdater<BarrageSubscriptionImpl> COMPLETED_UPDATER =
+            AtomicIntegerFieldUpdater.newUpdater(BarrageSubscriptionImpl.class, "completed");
     private volatile Throwable exceptionWhileCompleting;
 
     private volatile boolean connected = true;
@@ -173,7 +176,7 @@ public void onCompleted() {
 
     @Override
     public boolean isCompleted() {
-        return completed;
+        return completed == 1;
     }
 
     @Override
@@ -239,7 +242,7 @@ private boolean checkIfCompleteOrThrow() {
         if (exceptionWhileCompleting != null) {
             throw new UncheckedDeephavenException("Error while handling subscription:", exceptionWhileCompleting);
         }
-        return completed;
+        return completed == 1;
     }
 
     @Override
@@ -282,7 +285,17 @@ public BarrageTable blockUntilComplete() throws InterruptedException {
     }
 
     private synchronized void signalCompletion() {
-        completed = true;
+        signalCompletion(null);
+    }
+
+    private synchronized void signalCompletion(@Nullable final Throwable t) {
+        if (!COMPLETED_UPDATER.compareAndSet(BarrageSubscriptionImpl.this, 0, 1)) {
+            return;
+        }
+
+        if (t != null) {
+            exceptionWhileCompleting = t;
+        }
 
         // if we are building a snapshot via a growing viewport subscription, then cancel our subscription
         if (isSnapshot) {
@@ -341,13 +354,12 @@ private void handleDisconnect() {
     }
 
     @Override
-    public synchronized void close() {
+    public void close() {
         if (!connected) {
             return;
         }
 
-        exceptionWhileCompleting = new RequestCancelledException("BarrageSubscriptionImpl closed");
-        signalCompletion();
+        signalCompletion(new RequestCancelledException("BarrageSubscriptionImpl closed"));
         GrpcUtil.safelyComplete(observer);
         cleanup();
     }
@@ -485,7 +497,7 @@ public synchronized boolean viewportChanged(
                 @Nullable final RowSet serverViewport,
                 @Nullable final BitSet serverColumns,
                 final boolean serverReverseViewport) {
-            if (completed) {
+            if (completed == 1) {
                 return false;
             }
 
@@ -525,8 +537,7 @@ public synchronized boolean viewportChanged(
 
         @Override
         public void onError(Throwable t) {
-            exceptionWhileCompleting = t;
-            signalCompletion();
+            signalCompletion(t);
         }
 
         @Override