feat(typescript): worker mode for ts_project (#2136)

* feat(typescript): worker mode for ts_project

* chore: cleanup and declare protobufjs dependency

* fix: do not pass --watch for non-worker mode

* fix: do not test worker on windows

* fix: log on standalone failures

* chore: docs improvements

Co-authored-by: Alex Eagle <alex.eagle@robinhood.com>
Co-authored-by: Dan Muller <mrmeku@stairwell.com>

Author: 3 people

examples/react_webpack/BUILD.bazel

@@ -14,6 +14,8 @@ sass(
 )
 
 ts_project(
+    # Experimental: Start a tsc daemon to watch for changes to make recompiles faster.
+    supports_workers = True,
     deps = [
         "@npm//@types",
         "@npm//csstype",

examples/react_webpack/tsconfig.json

@@ -1,6 +1,15 @@
 {
     "compilerOptions": {
         "jsx": "react",
-        "lib": ["ES2015", "DOM"]
-    }
-}
+        "lib": [
+            "ES2015",
+            "DOM"
+        ]
+    },
+    // When using ts_project in worker mode, we run outside the Bazel sandbox (unless using --worker_sandboxing).
+    // We list the files that should be part of this particular compilation to avoid TypeScript discovering others.
+    "include": [
+        "*.tsx",
+        "*.ts"
+    ]
+}

internal/node/node.bzl

@@ -37,7 +37,7 @@ def _trim_package_node_modules(package_name):
     for n in package_name.split("/"):
         if n == "node_modules":
             break
-        segments += [n]
+        segments.append(n)
     return "/".join(segments)
 
 def _compute_node_modules_root(ctx):
@@ -150,6 +150,9 @@ def _to_execroot_path(ctx, file):
 
     return file.path
 
+def _join(*elements):
+    return "/".join([f for f in elements if f])
+
 def _nodejs_binary_impl(ctx):
     node_modules_manifest = write_node_modules_manifest(ctx, link_workspace_root = ctx.attr.link_workspace_root)
     node_modules_depsets = []
@@ -250,7 +253,12 @@ fi
     expanded_args = [expand_location_into_runfiles(ctx, a, ctx.attr.data) for a in expanded_args]
 
     # Next expand predefined variables & custom variables
-    expanded_args = [ctx.expand_make_variables("templated_args", e, {}) for e in expanded_args]
+    rule_dir = _join(ctx.bin_dir.path, ctx.label.workspace_root, ctx.label.package)
+    additional_substitutions = {
+        "@D": rule_dir,
+        "RULEDIR": rule_dir,
+    }
+    expanded_args = [ctx.expand_make_variables("templated_args", e, additional_substitutions) for e in expanded_args]
 
     substitutions = {
         # TODO: Split up results of multifile expansions into separate args and qoute them with

packages/typescript/BUILD.bazel

@@ -100,6 +100,7 @@ pkg_npm(
         ":npm_version_check",
         "//packages/typescript/internal:BUILD",
         "//packages/typescript/internal:ts_project_options_validator.js",
+        "//packages/typescript/internal/worker",
     ] + select({
         # FIXME: fix stardoc on Windows; //packages/typescript:index.md generation fails with:
         #   ERROR: D:/b/62unjjin/external/npm_bazel_typescript/BUILD.bazel:36:1: Couldn't build file

packages/typescript/internal/BUILD.bazel

@@ -64,6 +64,7 @@ filegroup(
         "ts_config.bzl",
         "ts_project.bzl",
         "//packages/typescript/internal/devserver:package_contents",
+        "//packages/typescript/internal/worker:package_contents",
     ],
     visibility = ["//packages/typescript:__subpackages__"],
 )

packages/typescript/internal/ts_project.bzl

@@ -2,6 +2,7 @@
 
 load("@build_bazel_rules_nodejs//:providers.bzl", "DeclarationInfo", "NpmPackageInfo", "declaration_info", "js_module_info", "run_node")
 load("@build_bazel_rules_nodejs//internal/linker:link_node_modules.bzl", "module_mappings_aspect")
+load("@build_bazel_rules_nodejs//internal/node:node.bzl", "nodejs_binary")
 load(":ts_config.bzl", "TsConfigInfo", "write_tsconfig")
 
 _ValidOptionsInfo = provider()
@@ -13,6 +14,20 @@ _DEFAULT_TSC = (
     "//typescript/bin:tsc"
 )
 
+_DEFAULT_TSC_BIN = (
+    # BEGIN-INTERNAL
+    "@npm" +
+    # END-INTERNAL
+    "//:node_modules/typescript/bin/tsc"
+)
+
+_DEFAULT_TYPESCRIPT_MODULE = (
+    # BEGIN-INTERNAL
+    "@npm" +
+    # END-INTERNAL
+    "//typescript"
+)
+
 _ATTRS = {
     "args": attr.string_list(),
     "declaration_dir": attr.string(),
@@ -33,7 +48,14 @@ _ATTRS = {
     # if you swap out the `compiler` attribute (like with ngtsc)
     # that compiler might allow more sources than tsc does.
     "srcs": attr.label_list(allow_files = True, mandatory = True),
-    "tsc": attr.label(default = Label(_DEFAULT_TSC), executable = True, cfg = "host"),
+    "supports_workers": attr.bool(
+        doc = """Experimental! Use only with caution.
+
+Allows you to enable the Bazel Worker strategy for this project.
+This requires that the tsc binary support it.""",
+        default = False,
+    ),
+    "tsc": attr.label(default = Label(_DEFAULT_TSC), executable = True, cfg = "target"),
     "tsconfig": attr.label(mandatory = True, allow_single_file = [".json"]),
 }
 
@@ -56,6 +78,16 @@ def _join(*elements):
 
 def _ts_project_impl(ctx):
     arguments = ctx.actions.args()
+    execution_requirements = {}
+    progress_prefix = "Compiling TypeScript project"
+
+    if ctx.attr.supports_workers:
+        # Set to use a multiline param-file for worker mode
+        arguments.use_param_file("@%s", use_always = True)
+        arguments.set_param_file_format("multiline")
+        execution_requirements["supports-workers"] = "1"
+        execution_requirements["worker-key-mnemonic"] = "TsProject"
+        progress_prefix = "Compiling TypeScript project (worker mode)"
 
     generated_srcs = False
     for src in ctx.files.srcs:
@@ -162,7 +194,9 @@ def _ts_project_impl(ctx):
             arguments = [arguments],
             outputs = outputs,
             executable = "tsc",
-            progress_message = "Compiling TypeScript project %s [tsc -p %s]" % (
+            execution_requirements = execution_requirements,
+            progress_message = "%s %s [tsc -p %s]" % (
+                progress_prefix,
                 ctx.label,
                 ctx.file.tsconfig.short_path,
             ),
@@ -287,7 +321,10 @@ def ts_project_macro(
         emit_declaration_only = False,
         ts_build_info_file = None,
         tsc = None,
+        worker_tsc_bin = _DEFAULT_TSC_BIN,
+        worker_typescript_module = _DEFAULT_TYPESCRIPT_MODULE,
         validate = True,
+        supports_workers = False,
         declaration_dir = None,
         out_dir = None,
         root_dir = None,
@@ -453,8 +490,28 @@ def ts_project_macro(
             For example, `tsc = "@my_deps//typescript/bin:tsc"`
             Or you can pass a custom compiler binary instead.
 
+        worker_tsc_bin: Label of the TypeScript compiler binary to run when running in worker mode.
+
+            For example, `tsc = "@my_deps//node_modules/typescript/bin/tsc"`
+            Or you can pass a custom compiler binary instead.
+
+        worker_typescript_module: Label of the package containing all data deps of worker_tsc_bin.
+
+            For example, `tsc = "@my_deps//typescript"`
+
         validate: boolean; whether to check that the tsconfig settings match the attributes.
 
+        supports_workers: Experimental! Use only with caution.
+
+            Allows you to enable the Bazel Persistent Workers strategy for this project.
+            See https://docs.bazel.build/versions/master/persistent-workers.html
+
+            This requires that the tsc binary support a `--watch` option.
+
+            NOTE: this does not work on Windows yet.
+            We will silently fallback to non-worker mode on Windows regardless of the value of this attribute.
+            Follow https://github.com/bazelbuild/rules_nodejs/issues/2277 for progress on this feature.
+
         root_dir: a string specifying a subdirectory under the input package which should be consider the
             root directory of all the input files.
             Equivalent to the TypeScript --rootDir option.
@@ -559,6 +616,38 @@ def ts_project_macro(
             )
             extra_deps.append("_validate_%s_options" % name)
 
+    if supports_workers:
+        tsc_worker = "%s_worker" % name
+        protobufjs = (
+            # BEGIN-INTERNAL
+            "@npm" +
+            # END-INTERNAL
+            "//protobufjs"
+        )
+        nodejs_binary(
+            name = tsc_worker,
+            data = [
+                Label("//packages/typescript/internal/worker:worker"),
+                Label(worker_tsc_bin),
+                Label(worker_typescript_module),
+                Label(protobufjs),
+                tsconfig,
+            ],
+            entry_point = Label("//packages/typescript/internal/worker:worker_adapter"),
+            templated_args = [
+                "--nobazel_patch_module_resolver",
+                "$(execpath {})".format(Label(worker_tsc_bin)),
+                "--project",
+                "$(execpath {})".format(tsconfig),
+                # FIXME: should take out_dir into account
+                "--outDir",
+                "$(RULEDIR)",
+                # FIXME: what about other settings like declaration_dir, root_dir, etc
+            ],
+        )
+
+        tsc = ":" + tsc_worker
+
     typings_out_dir = declaration_dir if declaration_dir else out_dir
     tsbuildinfo_path = ts_build_info_file if ts_build_info_file else name + ".tsbuildinfo"
 
@@ -583,5 +672,9 @@ def ts_project_macro(
         buildinfo_out = tsbuildinfo_path if composite or incremental else None,
         tsc = tsc,
         link_workspace_root = link_workspace_root,
+        supports_workers = select({
+            "@bazel_tools//src/conditions:host_windows": False,
+            "//conditions:default": supports_workers,
+        }),
         **kwargs
     )

packages/typescript/internal/worker/BUILD.bazel

@@ -0,0 +1,55 @@
+# BEGIN-INTERNAL
+
+load("//internal/common:copy_to_bin.bzl", "copy_to_bin")
+load("//third_party/github.com/bazelbuild/bazel-skylib:rules/copy_file.bzl", "copy_file")
+
+# Copy the proto file to a matching third_party/... nested directory
+# so the runtime require() statements still work
+_worker_proto_dir = "third_party/github.com/bazelbuild/bazel/src/main/protobuf"
+
+genrule(
+    name = "copy_worker_js",
+    srcs = ["//packages/worker:npm_package"],
+    outs = ["worker.js"],
+    cmd = "cp $(execpath //packages/worker:npm_package)/index.js $@",
+    visibility = ["//visibility:public"],
+)
+
+copy_file(
+    name = "copy_worker_proto",
+    src = "@build_bazel_rules_typescript//%s:worker_protocol.proto" % _worker_proto_dir,
+    out = "%s/worker_protocol.proto" % _worker_proto_dir,
+    visibility = ["//visibility:public"],
+)
+
+copy_to_bin(
+    name = "worker_adapter",
+    srcs = [
+        "worker_adapter.js",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+filegroup(
+    name = "package_contents",
+    srcs = [
+        "BUILD.bazel",
+    ],
+    visibility = ["//packages/typescript:__subpackages__"],
+)
+
+# END-INTERNAL
+
+exports_files([
+    "worker_adapter.js",
+])
+
+filegroup(
+    name = "worker",
+    srcs = [
+        "third_party/github.com/bazelbuild/bazel/src/main/protobuf/worker_protocol.proto",
+        "worker.js",
+        "worker_adapter.js",
+    ],
+    visibility = ["//visibility:public"],
+)

packages/typescript/internal/worker/worker_adapter.js

@@ -0,0 +1,99 @@
+/**
+ * @fileoverview wrapper program around the TypeScript compiler, tsc
+ *
+ * It intercepts the Bazel Persistent Worker protocol, using it to remote-control tsc running as a
+ * child process. In between builds, the tsc process is stopped (akin to ctrl-z in a shell) and then
+ * resumed (akin to `fg`) when the inputs have changed.
+ *
+ * See https://medium.com/@mmorearty/how-to-create-a-persistent-worker-for-bazel-7738bba2cabb
+ * for more background (note, that is documenting a different implementation)
+ */
+const child_process = require('child_process');
+const MNEMONIC = 'TsProject';
+const worker = require('./worker');
+
+const workerArg = process.argv.indexOf('--persistent_worker')
+if (workerArg > 0) {
+  process.argv.splice(workerArg, 1, '--watch')
+
+  if (process.platform !== 'linux' && process.platform !== 'darwin') {
+    throw new Error(`Worker mode is only supported on linux and darwin, not ${process.platform}.
+        See https://github.com/bazelbuild/rules_nodejs/issues/2277`);
+  }
+}
+
+const [tscBin, ...tscArgs] = process.argv.slice(2);
+
+const child = child_process.spawn(
+    tscBin,
+    tscArgs,
+    {stdio: 'pipe'},
+);
+function awaitOneBuild() {
+  child.kill('SIGCONT')
+
+  let buffer = [];
+  return new Promise((res) => {
+    function awaitBuild(s) {
+      buffer.push(s);
+
+      if (s.includes('Watching for file changes.')) {
+        child.kill('SIGSTOP')
+
+        const success = s.includes('Found 0 errors.');
+        res(success);
+
+        child.stdout.removeListener('data', awaitBuild);
+
+        if (!success) {
+          console.error(
+              `\nError output from tsc worker:\n\n  ${
+                  buffer.slice(1).map(s => s.toString()).join('').replace(/\n/g, '\n  ')}`,
+          )
+        }
+
+        buffer = [];
+      }
+    };
+    child.stdout.on('data', awaitBuild);
+  });
+}
+
+async function main() {
+  // Bazel will pass a special argument to the program when it's running us as a worker
+  if (workerArg > 0) {
+    worker.log(`Running ${MNEMONIC} as a Bazel worker`);
+
+    worker.runWorkerLoop(awaitOneBuild);
+  } else {
+    // Running standalone so stdout is available as usual
+    console.log(`Running ${MNEMONIC} as a standalone process`);
+    console.error(
+        `Started a new process to perform this action. Your build might be misconfigured, try
+      --strategy=${MNEMONIC}=worker`);
+
+    const stdoutbuffer = [];
+    child.stdout.on('data', data => stdoutbuffer.push(data));
+
+    const stderrbuffer = [];
+    child.stderr.on('data', data => stderrbuffer.push(data));
+
+    child.on('exit', code => {
+      if (code !== 0) {
+        console.error(
+            `\nstdout from tsc:\n\n  ${
+                stdoutbuffer.map(s => s.toString()).join('').replace(/\n/g, '\n  ')}`,
+        )
+        console.error(
+            `\nstderr from tsc:\n\n  ${
+                stderrbuffer.map(s => s.toString()).join('').replace(/\n/g, '\n  ')}`,
+        )
+      }
+      process.exit(code)
+    });
+  }
+}
+
+if (require.main === module) {
+  main();
+}