TableServiceAsyncTest: Eliminate Compiler Overhead Variance; Close QueryCompiler's JavaFileManager

engine/context/src/main/java/io/deephaven/engine/context/QueryCompiler.java

@@ -113,7 +113,7 @@ private QueryCompiler(
         try {
             urls[0] = (classDestination.toURI().toURL());
         } catch (MalformedURLException e) {
-            throw new RuntimeException("", e);
+            throw new UncheckedDeephavenException(e);
         }
         this.ucl = new WritableURLClassLoader(urls, parentClassLoaderToUse);
 
@@ -183,7 +183,8 @@ public static void writeClass(final File destinationDirectory, final String clas
         ensureDirectories(parentDir,
                 () -> "Unable to create missing destination directory " + parentDir.getAbsolutePath());
         if (!destinationFile.createNewFile()) {
-            throw new RuntimeException("Unable to create destination file " + destinationFile.getAbsolutePath());
+            throw new UncheckedDeephavenException(
+                    "Unable to create destination file " + destinationFile.getAbsolutePath());
         }
         final ByteArrayOutputStream byteOutStream = new ByteArrayOutputStream(data.length);
         byteOutStream.write(data, 0, data.length);
@@ -274,7 +275,7 @@ private static void ensureDirectories(final File file, final Supplier<String> ru
         // (and therefore mkdirs() would return false), but still get the directory we need (and therefore exists()
         // would return true)
         if (!file.mkdirs() && !file.isDirectory()) {
-            throw new RuntimeException(runtimeErrMsg.get());
+            throw new UncheckedDeephavenException(runtimeErrMsg.get());
         }
     }
 
@@ -396,7 +397,7 @@ private void addClassSource(File classSourceDirectory) {
         try {
             ucl.addURL(classSourceDirectory.toURI().toURL());
         } catch (MalformedURLException e) {
-            throw new RuntimeException("", e);
+            throw new UncheckedDeephavenException(e);
         }
     }
 
@@ -425,7 +426,7 @@ private Class<?> compileHelper(@NotNull final String className,
         try {
             digest = MessageDigest.getInstance("SHA-256");
         } catch (NoSuchAlgorithmException e) {
-            throw new RuntimeException("Unable to create SHA-256 hashing digest", e);
+            throw new UncheckedDeephavenException("Unable to create SHA-256 hashing digest", e);
         }
         final String basicHashText =
                 ByteUtils.byteArrToHex(digest.digest(classBody.getBytes(StandardCharsets.UTF_8)));
@@ -645,7 +646,8 @@ private void maybeCreateClass(String className, String code, String packageName,
 
         final String[] splitPackageName = packageName.split("\\.");
         if (splitPackageName.length == 0) {
-            throw new RuntimeException(String.format("packageName %s expected to have at least one .", packageName));
+            throw new UncheckedDeephavenException(String.format(
+                    "packageName %s expected to have at least one .", packageName));
         }
         final String[] truncatedSplitPackageName = Arrays.copyOf(splitPackageName, splitPackageName.length - 1);
 
@@ -689,23 +691,39 @@ private void maybeCreateClassHelper(String fqClassName, String finalCode, String
 
         final JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
         if (compiler == null) {
-            throw new RuntimeException("No Java compiler provided - are you using a JRE instead of a JDK?");
+            throw new UncheckedDeephavenException("No Java compiler provided - are you using a JRE instead of a JDK?");
         }
 
         final String classPathAsString = getClassPath() + File.pathSeparator + getJavaClassPath();
         final List<String> compilerOptions = Arrays.asList("-d", tempDirAsString, "-cp", classPathAsString);
 
-        final StandardJavaFileManager fileManager = compiler.getStandardFileManager(null, null, null);
+        final JavaFileManager fileManager = compiler.getStandardFileManager(null, null, null);
 
-        final boolean result = compiler.getTask(compilerOutput,
-                fileManager,
-                null,
-                compilerOptions,
-                null,
-                Collections.singletonList(new JavaSourceFromString(fqClassName, finalCode)))
-                .call();
+        boolean result = false;
+        boolean exceptionThrown = false;
+        try {
+            result = compiler.getTask(compilerOutput,
+                    fileManager,
+                    null,
+                    compilerOptions,
+                    null,
+                    Collections.singletonList(new JavaSourceFromString(fqClassName, finalCode)))
+                    .call();
+        } catch (final Throwable err) {
+            exceptionThrown = true;
+            throw err;
+        } finally {
+            try {
+                fileManager.close();
+            } catch (final IOException ioe) {
+                if (!exceptionThrown) {
+                    // noinspection ThrowFromFinallyBlock
+                    throw new UncheckedIOException("Could not close JavaFileManager", ioe);
+                }
+            }
+        }
         if (!result) {
-            throw new RuntimeException("Error compiling class " + fqClassName + ":\n" + compilerOutput);
+            throw new UncheckedDeephavenException("Error compiling class " + fqClassName + ":\n" + compilerOutput);
         }
         // The above has compiled into e.g.
         // /tmp/workspace/cache/classes/temporaryCompilationDirectory12345/io/deephaven/test/cm12862183232603186v52_0/{various
@@ -737,7 +755,7 @@ private void maybeCreateClassHelper(String fqClassName, String finalCode, String
     private Pair<Boolean, String> tryCompile(File basePath, Collection<File> javaFiles) throws IOException {
         final JavaCompiler compiler = ToolProvider.getSystemJavaCompiler();
         if (compiler == null) {
-            throw new RuntimeException("No Java compiler provided - are you using a JRE instead of a JDK?");
+            throw new UncheckedDeephavenException("No Java compiler provided - are you using a JRE instead of a JDK?");
         }
 
         final File outputDirectory = Files.createTempDirectory("temporaryCompilationDirectory").toFile();
@@ -828,7 +846,7 @@ private static String getJavaClassPath() {
                     }
                 }
             } catch (IOException e) {
-                throw new RuntimeException("Error extract manifest file from " + javaClasspath + ".\n", e);
+                throw new UncheckedIOException("Error extract manifest file from " + javaClasspath + ".\n", e);
             }
         }
         return javaClasspath;

java-client/session-dagger/src/test/java/io/deephaven/client/DeephavenSessionTestBase.java

@@ -3,8 +3,11 @@
  */
 package io.deephaven.client;
 
+import io.deephaven.base.verify.Require;
 import io.deephaven.client.impl.Session;
+import io.deephaven.client.impl.SessionImpl;
 import io.deephaven.server.runner.DeephavenApiServerTestBase;
+import io.deephaven.server.session.SessionState;
 import io.grpc.ManagedChannel;
 import org.junit.After;
 import org.junit.Before;
@@ -17,6 +20,7 @@ public abstract class DeephavenSessionTestBase extends DeephavenApiServerTestBas
 
     private ScheduledExecutorService sessionScheduler;
     protected Session session;
+    protected SessionState serverSessionState;
 
     @Override
     @Before
@@ -25,8 +29,12 @@ public void setUp() throws Exception {
         ManagedChannel channel = channelBuilder().build();
         register(channel);
         sessionScheduler = Executors.newScheduledThreadPool(2);
-        session = DaggerDeephavenSessionRoot.create().factoryBuilder().managedChannel(channel)
-                .scheduler(sessionScheduler).build().newSession();
+        final SessionImpl clientSessionImpl =
+                DaggerDeephavenSessionRoot.create().factoryBuilder().managedChannel(channel)
+                        .scheduler(sessionScheduler).build().newSession();
+        session = clientSessionImpl;
+        serverSessionState = Require.neqNull(server().sessionService().getSessionForToken(
+                clientSessionImpl._hackBearerHandler().getCurrentToken()), "SessionState");
     }
 
     @Override

java-client/session-dagger/src/test/java/io/deephaven/client/impl/TableServiceAsyncTest.java

@@ -5,7 +5,12 @@
 
 import io.deephaven.client.DeephavenSessionTestBase;
 import io.deephaven.client.impl.TableService.TableHandleFuture;
+import io.deephaven.engine.context.ExecutionContext;
+import io.deephaven.engine.table.Table;
+import io.deephaven.engine.testutil.TstUtils;
+import io.deephaven.engine.util.TableTools;
 import io.deephaven.qst.table.TableSpec;
+import io.deephaven.util.SafeCloseable;
 import org.junit.Test;
 
 import java.time.Duration;
@@ -19,15 +24,15 @@
 public class TableServiceAsyncTest extends DeephavenSessionTestBase {
 
     private static final Duration GETTIME = Duration.ofSeconds(15);
-    private static final int CHAIN_OPS = 50;
+    private static final int CHAIN_OPS = 250;
     private static final int CHAIN_ROWS = 1000;
 
     @Test(timeout = 20000)
     public void longChainAsyncExportOnlyLast() throws ExecutionException, InterruptedException, TimeoutException {
         final List<TableSpec> longChain = createLongChain();
         final TableSpec longChainLast = longChain.get(longChain.size() - 1);
         try (final TableHandle handle = get(session.executeAsync(longChainLast))) {
-            checkSucceeded(handle);
+            checkSucceeded(handle, CHAIN_OPS);
         }
     }
 
@@ -36,9 +41,10 @@ public void longChainAsyncExportAll() throws ExecutionException, InterruptedExce
         final List<TableSpec> longChain = createLongChain();
         final List<? extends TableHandleFuture> futures = session.executeAsync(longChain);
         try {
+            int chainLength = 0;
             for (final TableHandleFuture future : futures) {
                 try (final TableHandle handle = get(future)) {
-                    checkSucceeded(handle);
+                    checkSucceeded(handle, ++chainLength);
                 }
             }
         } catch (final Throwable t) {
@@ -55,7 +61,7 @@ public void longChainAsyncExportAllCancelAllButLast()
         // Cancel or close all but the last one
         TableService.TableHandleFuture.cancelOrClose(futures.subList(0, futures.size() - 1), true);
         try (final TableHandle lastHandle = get(futures.get(futures.size() - 1))) {
-            checkSucceeded(lastHandle);
+            checkSucceeded(lastHandle, CHAIN_OPS);
         }
     }
 
@@ -68,7 +74,7 @@ public void immediatelyCompletedFromCachedTableServices()
         try (final TableHandle ignored = get(tableService.executeAsync(longChainLast))) {
             for (int i = 0; i < 1000; ++i) {
                 try (final TableHandle handle = get(tableService.executeAsync(longChainLast))) {
-                    checkSucceeded(handle);
+                    checkSucceeded(handle, CHAIN_OPS);
                 }
             }
         }
@@ -79,22 +85,30 @@ private static TableHandle get(TableHandleFuture future)
         return future.getOrCancel(GETTIME);
     }
 
-    private static void checkSucceeded(TableHandle x) {
+    private void checkSucceeded(TableHandle x, int chainLength) {
         assertThat(x.isSuccessful()).isTrue();
+        try (final SafeCloseable ignored = getExecutionContext().open()) {
+            final Table result = serverSessionState.<Table>getExport(x.exportId().id()).get();
+            ExecutionContext.getContext().getQueryScope().putParam("ChainLength", chainLength);
+            final Table expected = TableTools.emptyTable(CHAIN_ROWS).update("Current = ii - 1 + ChainLength");
+            TstUtils.assertTableEquals(expected, result);
+        }
     }
 
     private static List<TableSpec> createLongChain() {
         return createLongChain(CHAIN_OPS, CHAIN_ROWS);
     }
 
-    private static List<TableSpec> createLongChain(int numColumns, int numRows) {
-        final List<TableSpec> longChain = new ArrayList<>(numColumns);
-        for (int i = 0; i < numColumns; ++i) {
+    private static List<TableSpec> createLongChain(int chainLength, int numRows) {
+        final List<TableSpec> longChain = new ArrayList<>(chainLength);
+        for (int i = 0; i < chainLength; ++i) {
             if (i == 0) {
-                longChain.add(TableSpec.empty(numRows).view("I_0=ii"));
+                longChain.add(TableSpec.empty(numRows).view("Current = ii"));
             } else {
                 final TableSpec prev = longChain.get(i - 1);
-                longChain.add(prev.updateView("I_" + i + " = 1 + I_" + (i - 1)));
+                // Note: it's important that this formula is constant with respect to "i", otherwise we'll spend a lot
+                // of time compiling formulas
+                longChain.add(prev.updateView("Current = 1 + Current"));
             }
         }
         return longChain;

java-client/session/src/main/java/io/deephaven/client/impl/BearerHandler.java

@@ -1,5 +1,6 @@
 package io.deephaven.client.impl;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.grpc.CallCredentials;
 import io.grpc.CallOptions;
 import io.grpc.Channel;
@@ -14,6 +15,7 @@
 
 import java.util.Objects;
 import java.util.Optional;
+import java.util.UUID;
 import java.util.concurrent.Executor;
 
 import static io.deephaven.client.impl.Authentication.AUTHORIZATION_HEADER;
@@ -56,6 +58,11 @@ public void setBearerToken(String bearerToken) {
         }
     }
 
+    @VisibleForTesting
+    public UUID getCurrentToken() {
+        return UUID.fromString(bearerToken);
+    }
+
     private void handleMetadata(Metadata metadata) {
         parseBearerToken(metadata).ifPresent(BearerHandler.this::setBearerToken);
     }

java-client/session/src/main/java/io/deephaven/client/impl/SessionImpl.java

@@ -123,8 +123,8 @@ private SessionImpl(SessionImplConfig config, DeephavenChannel bearerChannel, Du
                 pingFrequency.toNanos(), pingFrequency.toNanos(), TimeUnit.NANOSECONDS);
     }
 
-    // exposed for Flight
-    BearerHandler _hackBearerHandler() {
+    // exposed for Flight and testing
+    public BearerHandler _hackBearerHandler() {
         return bearerHandler;
     }
 

server/src/main/java/io/deephaven/server/runner/DeephavenApiServer.java

@@ -98,7 +98,7 @@ public GrpcServer server() {
     }
 
     @VisibleForTesting
-    SessionService sessionService() {
+    public SessionService sessionService() {
         return sessionService;
     }
 

server/src/test/java/io/deephaven/server/runner/DeephavenApiServerTestBase.java

@@ -6,6 +6,7 @@
 import dagger.BindsInstance;
 import dagger.Component;
 import io.deephaven.client.ClientDefaultsModule;
+import io.deephaven.engine.context.ExecutionContext;
 import io.deephaven.engine.context.TestExecutionContext;
 import io.deephaven.engine.liveness.LivenessScope;
 import io.deephaven.engine.liveness.LivenessScopeStack;
@@ -85,7 +86,8 @@ interface Builder {
     @Rule
     public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
 
-    private SafeCloseable executionContext;
+    private ExecutionContext executionContext;
+    private SafeCloseable executionContextCloseable;
 
     private LogBuffer logBuffer;
     private SafeCloseable scopeCloseable;
@@ -127,7 +129,8 @@ public void setUp() throws Exception {
                 .injectFields(this);
 
         final PeriodicUpdateGraph updateGraph = server.getUpdateGraph().cast();
-        executionContext = TestExecutionContext.createForUnitTests().withUpdateGraph(updateGraph).open();
+        executionContext = TestExecutionContext.createForUnitTests().withUpdateGraph(updateGraph);
+        executionContextCloseable = executionContext.open();
         if (updateGraph.isUnitTestModeAllowed()) {
             updateGraph.enableUnitTestMode();
             updateGraph.resetForUnitTests(false);
@@ -153,7 +156,7 @@ public void tearDown() throws Exception {
         if (updateGraph.isUnitTestModeAllowed()) {
             updateGraph.resetForUnitTests(true);
         }
-        executionContext.close();
+        executionContextCloseable.close();
 
         scheduler.shutdown();
     }
@@ -170,6 +173,10 @@ public ScriptSession getScriptSession() {
         return scriptSessionProvider.get();
     }
 
+    public ExecutionContext getExecutionContext() {
+        return executionContext;
+    }
+
     /**
      * The session token expiration
      *