Introduce SessionService::TerminationNotification as best-effort attempt to notify on exit

Author: nbauernfeind

FishUtil/src/main/java/io/deephaven/util/process/FatalErrorReporter.java

@@ -21,6 +21,14 @@ public interface FatalErrorReporter extends Thread.UncaughtExceptionHandler {
 
     @FunctionalInterface
     interface Interceptor {
-        void intercept(@NotNull String message, @NotNull Throwable throwable);
+        /**
+         * Report a fatal error.
+         *
+         * @param message the message
+         * @param throwable the throwable
+         * @param isFromUncaught true iff called from
+         *        {@link java.lang.Thread.UncaughtExceptionHandler#uncaughtException(Thread, Throwable)}.
+         */
+        void intercept(@NotNull String message, @NotNull Throwable throwable, boolean isFromUncaught);
     }
 }

FishUtil/src/main/java/io/deephaven/util/process/FatalErrorReporterBase.java

@@ -33,7 +33,7 @@ private final class FatalException extends RuntimeException {
 
     @Override
     public final void report(@NotNull final String message, @NotNull final Throwable throwable) {
-        interceptors.forEach(interceptor -> interceptor.intercept(message, throwable));
+        interceptors.forEach(interceptor -> interceptor.intercept(message, throwable, false));
         reportImpl(message, throwable, false);
     }
 
@@ -56,7 +56,7 @@ public final void reportAsync(@NotNull final String message) {
     @Override
     public final void uncaughtException(@NotNull final Thread thread, @NotNull final Throwable throwable) {
         final String message = "Uncaught exception in thread " + thread.getName();
-        interceptors.forEach(interceptor -> interceptor.intercept(message, throwable));
+        interceptors.forEach(interceptor -> interceptor.intercept(message, throwable, true));
         reportImpl(message, throwable, true);
     }
 

grpc-api/src/main/java/io/deephaven/grpc_api/console/ConsoleServiceGrpcImpl.java

@@ -15,6 +15,7 @@
 import io.deephaven.db.util.VariableProvider;
 import io.deephaven.db.v2.DynamicNode;
 import io.deephaven.figures.FigureWidgetTranslator;
+import io.deephaven.grpc_api.session.SessionCloseableObserver;
 import io.deephaven.grpc_api.session.SessionService;
 import io.deephaven.grpc_api.session.SessionState;
 import io.deephaven.grpc_api.session.SessionState.ExportBuilder;
@@ -402,40 +403,16 @@ public void fetchFigure(FetchFigureRequest request, StreamObserver<FetchFigureRe
         });
     }
 
-    private class LogBufferStreamAdapter implements Closeable, LogBufferRecordListener {
-        private final SessionState session;
+    private static class LogBufferStreamAdapter extends SessionCloseableObserver<LogSubscriptionData>
+            implements LogBufferRecordListener {
         private final LogSubscriptionRequest request;
-        private final StreamObserver<LogSubscriptionData> responseObserver;
-        private boolean isClosed = false;
 
         public LogBufferStreamAdapter(
                 final SessionState session,
                 final LogSubscriptionRequest request,
                 final StreamObserver<LogSubscriptionData> responseObserver) {
-            this.session = session;
+            super(session, responseObserver);
             this.request = request;
-            this.responseObserver = responseObserver;
-            session.addOnCloseCallback(this);
-            ((ServerCallStreamObserver<LogSubscriptionData>) responseObserver).setOnCancelHandler(this::tryClose);
-        }
-
-        @Override
-        public void close() {
-            synchronized (this) {
-                if (isClosed) {
-                    return;
-                }
-                isClosed = true;
-            }
-
-            safelyExecute(() -> logBuffer.unsubscribe(this));
-            safelyExecuteLocked(responseObserver, responseObserver::onCompleted);
-        }
-
-        private void tryClose() {
-            if (session.removeOnCloseCallback(this)) {
-                close();
-            }
         }
 
         @Override
@@ -465,9 +442,9 @@ public void record(LogBufferRecord record) {
                 synchronized (responseObserver) {
                     responseObserver.onNext(payload);
                 }
-            } catch (Throwable t) {
+            } catch (Throwable ignored) {
                 // we are ignoring exceptions here deliberately, and just shutting down
-                tryClose();
+                close();
             }
         }
     }

grpc-api/src/main/java/io/deephaven/grpc_api/runner/DeephavenApiServer.java

@@ -80,7 +80,7 @@ public static void startMain(PrintStream out, PrintStream err)
 
         // Close outstanding sessions to give any gRPCs closure.
         ProcessEnvironment.getGlobalShutdownManager().registerTask(ShutdownManager.OrderingCategory.MIDDLE,
-                sessionService::closeAllSessions);
+                sessionService::onShutdown);
 
         // Finally wait for gRPC to exit now.
         ProcessEnvironment.getGlobalShutdownManager().registerTask(ShutdownManager.OrderingCategory.LAST, () -> {

grpc-api/src/main/java/io/deephaven/grpc_api/session/SessionCloseableObserver.java

@@ -0,0 +1,45 @@
+package io.deephaven.grpc_api.session;
+
+import io.grpc.stub.ServerCallStreamObserver;
+import io.grpc.stub.StreamObserver;
+
+import java.io.Closeable;
+
+import static io.deephaven.grpc_api.util.GrpcUtil.safelyExecuteLocked;
+
+public abstract class SessionCloseableObserver<T> implements Closeable {
+    protected final SessionState session;
+    protected final StreamObserver<T> responseObserver;
+    private boolean isClosed = false;
+
+    public SessionCloseableObserver(
+            final SessionState session,
+            final StreamObserver<T> responseObserver) {
+        this.session = session;
+        this.responseObserver = responseObserver;
+        session.addOnCloseCallback(this);
+        ((ServerCallStreamObserver<T>) responseObserver).setOnCancelHandler(this::close);
+    }
+
+    @Override
+    public final void close() {
+        session.removeOnCloseCallback(this);
+
+        synchronized (this) {
+            if (isClosed) {
+                return;
+            }
+            isClosed = true;
+        }
+
+        onClose();
+        safelyExecuteLocked(responseObserver, responseObserver::onCompleted);
+    }
+
+    /**
+     * Override this to perform any additional specific clean up that must be performed.
+     */
+    void onClose() {
+
+    }
+}

grpc-api/src/main/java/io/deephaven/grpc_api/session/SessionService.java

@@ -4,21 +4,27 @@
 import com.google.protobuf.ByteString;
 import io.deephaven.db.tables.utils.DBDateTime;
 import io.deephaven.db.tables.utils.DBTimeUtils;
+import io.deephaven.grpc_api.util.GrpcUtil;
+import io.deephaven.proto.backplane.grpc.TerminationNotificationResponse;
 import io.deephaven.util.auth.AuthContext;
 import io.deephaven.grpc_api.util.Scheduler;
+import io.deephaven.util.process.ProcessEnvironment;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
+import io.grpc.stub.StreamObserver;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import javax.inject.Inject;
 import javax.inject.Named;
 import javax.inject.Singleton;
 import java.util.Deque;
+import java.util.List;
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedDeque;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 @Singleton
 public class SessionService {
@@ -36,6 +42,8 @@ public class SessionService {
     private boolean cleanupJobInstalled = false;
     private final SessionCleanupJob sessionCleanupJob = new SessionCleanupJob();
 
+    private final List<TerminationNotificationListener> terminationListeners = new CopyOnWriteArrayList<>();
+
     @Inject()
     public SessionService(final Scheduler scheduler, final SessionState.Factory sessionFactory,
             @Named("session.tokenExpireMs") final long tokenExpireMs) {
@@ -51,8 +59,53 @@ public SessionService(final Scheduler scheduler, final SessionState.Factory sess
 
         // Protect ourselves from rotation spam, but be loose enough that any reasonable refresh strategy works.
         this.tokenRotateMs = tokenExpireMs / 5;
+
+        ProcessEnvironment.getGlobalFatalErrorReporter().addInterceptor(this::onFatalError);
+    }
+
+    private synchronized void onFatalError(
+            @NotNull String message,
+            @NotNull Throwable throwable,
+            boolean isFromUncaught) {
+        final TerminationNotificationResponse.Builder builder =
+                TerminationNotificationResponse.newBuilder()
+                        .setAbnormalTermination(true)
+                        .setReason(message);
+
+        builder.addStackTrace(throwable.getMessage());
+        for (final StackTraceElement element : throwable.getStackTrace()) {
+            builder.addStackTrace(element.toString());
+        }
+
+        final TerminationNotificationResponse notification = builder.build();
+        terminationListeners.forEach(listener -> listener.sendMessage(notification));
+        terminationListeners.clear();
+    }
+
+    public synchronized void onShutdown() {
+        final TerminationNotificationResponse notification = TerminationNotificationResponse.newBuilder()
+                .setAbnormalTermination(false)
+                .build();
+        terminationListeners.forEach(listener -> listener.sendMessage(notification));
+        terminationListeners.clear();
+
+        closeAllSessions();
+    }
+
+    /**
+     * Add a listener who receives a single notification when this process is exiting and yet able to communicate with
+     * the observer.
+     *
+     * @param session the session the observer belongs to
+     * @param responseObserver the observer to notify
+     */
+    public void addTerminationListener(
+            final SessionState session,
+            final StreamObserver<TerminationNotificationResponse> responseObserver) {
+        terminationListeners.add(new TerminationNotificationListener(session, responseObserver));
     }
 
+
     /**
      * Create a new session object for the provided auth context.
      *
@@ -168,7 +221,7 @@ public SessionState getOptionalSession() {
 
     /**
      * Reduces the liveness of the session.
-     * 
+     *
      * @param session the session to close
      */
     public void closeSession(final SessionState session) {
@@ -178,7 +231,7 @@ public void closeSession(final SessionState session) {
         session.onExpired();
     }
 
-    public void closeAllSessions() {
+    private void closeAllSessions() {
         for (final TokenExpiration token : outstandingCookies) {
             // close all exports/resources acquired by the session
             token.session.onExpired();
@@ -239,4 +292,25 @@ public void run() {
             }
         }
     }
+
+    private final class TerminationNotificationListener
+            extends SessionCloseableObserver<TerminationNotificationResponse> {
+        public TerminationNotificationListener(
+                final SessionState session,
+                final StreamObserver<TerminationNotificationResponse> responseObserver) {
+            super(session, responseObserver);
+        }
+
+        @Override
+        void onClose() {
+            terminationListeners.remove(this);
+        }
+
+        void sendMessage(final TerminationNotificationResponse response) {
+            GrpcUtil.safelyExecuteLocked(responseObserver, () -> {
+                responseObserver.onNext(response);
+                responseObserver.onCompleted();
+            });
+        }
+    }
 }

grpc-api/src/main/java/io/deephaven/grpc_api/session/SessionServiceGrpcImpl.java

@@ -13,14 +13,17 @@
 import io.deephaven.io.logger.Logger;
 import io.deephaven.proto.backplane.grpc.*;
 import io.deephaven.util.auth.AuthContext;
+import io.deephaven.util.process.ProcessEnvironment;
 import io.grpc.*;
 import io.grpc.stub.ServerCallStreamObserver;
 import io.grpc.stub.StreamObserver;
 
 import javax.inject.Inject;
 import javax.inject.Singleton;
+import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 public class SessionServiceGrpcImpl extends SessionServiceGrpc.SessionServiceImplBase {
     // TODO (#997): use flight AuthConstants
@@ -199,6 +202,15 @@ public void exportNotifications(final ExportNotificationRequest request,
         });
     }
 
+    @Override
+    public void terminationNotification(TerminationNotificationRequest request,
+            StreamObserver<TerminationNotificationResponse> responseObserver) {
+        GrpcUtil.rpcWrapper(log, responseObserver, () -> {
+            final SessionState session = service.getCurrentSession();
+            service.addTerminationListener(session, responseObserver);
+        });
+    }
+
     @Singleton
     public static class AuthServerInterceptor implements ServerInterceptor {
         private final SessionService service;
@@ -222,4 +234,5 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(final ServerCall<Re
             return Contexts.interceptCall(newContext, serverCall, metadata, serverCallHandler);
         }
     }
+
 }

proto/proto-backplane-grpc/src/main/proto/deephaven/proto/session.proto

@@ -67,6 +67,12 @@ service SessionService {
    * updated exports. An export id of zero will be sent to indicate all pre-existing exports have been sent.
    */
   rpc ExportNotifications(ExportNotificationRequest) returns (stream ExportNotification) {}
+
+  /*
+   * Receive a best-effort message on-exit indicating why this worker is exiting. Reception of this message cannot be
+   * guaranteed.
+   */
+  rpc TerminationNotification(TerminationNotificationRequest) returns (TerminationNotificationResponse) {}
 }
 
 /*
@@ -172,3 +178,18 @@ message ExportNotification {
    */
   string dependent_handle = 4;
 }
+
+message TerminationNotificationRequest {
+  // Intentionally empty and is here for backwards compatibility should this API change.
+}
+
+message TerminationNotificationResponse {
+  // whether or not this termination is expected
+  bool abnormal_termination = 1;
+  // if additional information is available then provide it in this field
+  string reason = 2;
+  // if this is due to an exception, whether or not it was uncaught
+  bool is_from_uncaught_exception = 3;
+  // if the termination is caused by an exception of any sort, provide the detailed stack trace here
+  repeated string stack_trace = 4;
+}