JiT test suite

cli/BUILD

@@ -65,7 +65,13 @@ ts_test_suite(
         "index_project_test.ts",
         "index_compile_test.ts",
         "index_run_e2e_test.ts",
-        "util_test.ts"
+        "tests/jit/index_jit_main_test.ts",
+        "tests/jit/index_jit_advanced_test.ts",
+        "tests/jit/index_jit_dependency_test.ts",
+        "tests/jit/index_jit_runtime_test.ts",
+        "util_test.ts",
+        "tests/jit/jit_build_test.ts",
+        "tests/jit/jit_run_test.ts",
     ],
     data = [
         ":node_modules",

cli/api/BUILD

@@ -1,6 +1,6 @@
+load("//tools:ts_library.bzl", "ts_library")
 load("//testing:index.bzl", "ts_test_suite")
 load("//tools:node_modules.bzl", "node_modules")
-load("//tools:ts_library.bzl", "ts_library")
 
 package(default_visibility = ["//visibility:public"])
 

cli/api/commands/base_worker.ts

@@ -4,9 +4,10 @@ export abstract class BaseWorker<TResponse, TMessage = any> {
   protected constructor(private readonly loaderPath: string) {}
 
   protected async runWorker(
-    timeoutMillis: number,
+    timeoutMillis: number | undefined,
     onBoot: (child: ChildProcess) => void,
-    onMessage: (message: TMessage, child: ChildProcess, resolve: (res: TResponse) => void, reject: (err: Error) => void) => void
+    onMessage: (message: TMessage, child: ChildProcess, resolve: (res: TResponse) => void, reject: (err: Error) => void) => void,
+    onCancel?: (cancel: () => void) => void
   ): Promise<TResponse> {
     const forkScript = this.resolveScript();
     const child = fork(forkScript, [], {
@@ -22,20 +23,28 @@ export abstract class BaseWorker<TResponse, TMessage = any> {
           return;
         }
         completed = true;
-        clearTimeout(timeout);
+        if (timeout !== undefined) {
+          clearTimeout(timeout);
+        }
         child.kill("SIGKILL");
         fn();
       };
 
-      const timeout = setTimeout(() => {
-        terminate(() =>
-          reject(new Error(
-            `Compilation timed out after ${timeoutMillis / 1000} seconds. ` +
-            `To allow more time, re-run with a longer --timeout ` +
-            `(e.g. --timeout=2m, --timeout=1h).`
-          ))
-        );
-      }, timeoutMillis);
+      const timeout = timeoutMillis !== undefined
+        ? setTimeout(() => {
+            terminate(() =>
+              reject(new Error(
+                `Compilation timed out after ${timeoutMillis / 1000} seconds. ` +
+                `To allow more time, re-run with a longer --timeout ` +
+                `(e.g. --timeout=2m, --timeout=1h).`
+              ))
+            );
+          }, timeoutMillis)
+        : undefined;
+
+      onCancel?.(() =>
+        terminate(() => reject(new Error("Run cancelled while worker was in flight.")))
+      );
 
       child.on("message", (message: any) => {
         if (message.type === "worker_booted") {

cli/api/commands/build.ts

@@ -69,7 +69,8 @@ export class Builder {
       runConfig: this.runConfig,
       warehouseState: this.warehouseState,
       declarationTargets: this.prunedGraph.declarations.map(declaration => declaration.target),
-      actions
+      actions,
+      jitData: this.prunedGraph.jitData
     });
   }
 
@@ -108,11 +109,16 @@ export class Builder {
   private toPartialExecutionAction(
     action: dataform.ITable | dataform.IOperation | dataform.IAssertion
   ) {
-    return dataform.ExecutionAction.create({
+    const jitCode = (action as any).jitCode;
+    const executionAction = dataform.ExecutionAction.create({
       target: action.target,
       fileName: action.fileName,
       dependencyTargets: action.dependencyTargets,
       actionDescriptor: action.actionDescriptor
     });
+    if (jitCode && !action.disabled) {
+      executionAction.jitCode = jitCode;
+    }
+    return executionAction;
   }
 }

cli/api/commands/jit/compiler.ts

@@ -5,7 +5,6 @@ import { BaseWorker } from "df/cli/api/commands/base_worker";
 import { handleDbRequest } from "df/cli/api/commands/jit/rpc";
 import { IDbAdapter, IDbClient } from "df/cli/api/dbadapters";
 import { IBigQueryExecutionOptions } from "df/cli/api/dbadapters/bigquery";
-import { DEFAULT_COMPILATION_TIMEOUT_MILLIS } from "df/cli/api/utils/constants";
 import { dataform } from "df/protos/ts";
 
 export interface IJitWorkerMessage {
@@ -26,16 +25,18 @@ export class JitCompileChildProcess extends BaseWorker<
     projectDir: string,
     dbadapter: IDbAdapter,
     dbclient: IDbClient,
-    timeoutMillis: number = DEFAULT_COMPILATION_TIMEOUT_MILLIS,
-    options?: IBigQueryExecutionOptions
+    timeoutMillis?: number,
+    options?: IBigQueryExecutionOptions,
+    onCancel?: (cancel: () => void) => void
   ): Promise<dataform.IJitCompilationResponse> {
     return await new JitCompileChildProcess().run(
       request,
       projectDir,
       dbadapter,
       dbclient,
       timeoutMillis,
-      options
+      options,
+      onCancel
     );
   }
 
@@ -48,8 +49,9 @@ export class JitCompileChildProcess extends BaseWorker<
     projectDir: string,
     dbadapter: IDbAdapter,
     dbclient: IDbClient,
-    timeoutMillis: number,
-    options?: IBigQueryExecutionOptions
+    timeoutMillis: number | undefined,
+    options: IBigQueryExecutionOptions | undefined,
+    onCancel: ((cancel: () => void) => void) | undefined
   ): Promise<dataform.IJitCompilationResponse> {
     return await this.runWorker(
       timeoutMillis,
@@ -68,7 +70,8 @@ export class JitCompileChildProcess extends BaseWorker<
         } else if (message.type === "jit_error") {
           reject(new Error(message.error));
         }
-      }
+      },
+      onCancel
     );
   }
 
@@ -90,7 +93,7 @@ export class JitCompileChildProcess extends BaseWorker<
       child.send({
         type: "rpc_response",
         correlationId: message.correlationId,
-        response
+        response: Array.from(response)
       });
     } catch (e) {
       child.send({