Support the multiplexing worker

Author: tek

decls/haskell_rules.bzl

@@ -70,6 +70,7 @@ haskell_binary = prelude_rule(
             "_worker_srcs": attrs.list(attrs.source(), default = worker_srcs),
             "_worker_deps": attrs.list(attrs.dep(), default = ["@prelude//haskell/worker:{}".format(pkg) for pkg in worker_libs]),
             "_worker_compiler_flags": attrs.list(attrs.string(), default = worker_flags),
+            "_worker_plugin": attrs.dep(default = "@prelude//haskell/worker:ghc-persistent-worker-plugin"),
         }
     ),
 )
@@ -195,6 +196,7 @@ haskell_library = prelude_rule(
             "_worker_srcs": attrs.list(attrs.source(), default = worker_srcs),
             "_worker_deps": attrs.list(attrs.dep(), default = ["@prelude//haskell/worker:{}".format(pkg) for pkg in worker_libs]),
             "_worker_compiler_flags": attrs.list(attrs.string(), default = worker_flags),
+            "_worker_plugin": attrs.dep(default = "@prelude//haskell/worker:ghc-persistent-worker-plugin"),
         }
     ),
 )

haskell/compile.bzl

@@ -438,11 +438,21 @@ def _common_compile_module_args(
     sources: list[Artifact],
     direct_deps_info: list[HaskellLibraryInfoTSet],
     pkgname: str | None = None,
+    worker_plugin: Dependency | None = None,
 ) -> CommonCompileModuleArgs:
     command = cmd_args(ghc_wrapper)
     command.add("--ghc", haskell_toolchain.compiler)
     command.add("--ghc-dir", haskell_toolchain.ghc_dir)
 
+    if haskell_toolchain.use_worker and worker_plugin != None:
+        if pkgname == None:
+            warning("Module {} has no 'pkgname', plugin worker will break".format(label))
+        else:
+            package_db = pkg_deps.providers[DynamicHaskellPackageDbInfo].packages
+            db = package_db[worker_plugin[HaskellToolchainLibrary].name]
+            command.add("--plugin-db", db.value.db)
+        command.add("--worker-target-id", pkgname)
+
     # Some rules pass in RTS (e.g. `+RTS ... -RTS`) options for GHC, which can't
     # be parsed when inside an argsfile.
     command.add(haskell_toolchain.compiler_flags)
@@ -754,6 +764,7 @@ def _dynamic_do_compile_impl(actions, md_file, pkg_deps, arg, direct_deps_by_nam
         link_style = arg.link_style,
         direct_deps_info = arg.direct_deps_info,
         pkgname = arg.pkgname,
+        worker_plugin = arg.worker_plugin,
     )
 
     md = md_file.read_json()
@@ -814,6 +825,7 @@ def compile(
         enable_profiling: bool,
         enable_haddock: bool,
         md_file: Artifact,
+        worker_plugin: Dependency | None,
         worker: WorkerInfo | None = None,
         pkgname: str | None = None) -> CompileResultInfo:
     artifact_suffix = get_artifact_suffix(link_style, enable_profiling)
@@ -871,6 +883,7 @@ def compile(
             srcs_envs = ctx.attrs.srcs_envs,
             toolchain_deps_by_name = toolchain_deps_by_name,
             worker = worker,
+            worker_plugin = worker_plugin,
         ),
     ))
 

haskell/haskell.bzl

@@ -648,6 +648,7 @@ def _build_haskell_lib(
         md_file = md_file,
         pkgname = pkgname,
         worker = _persistent_worker(ctx),
+        worker_plugin = ctx.attrs._worker_plugin if ctx.label.cell != "prelude" and ctx.attrs._haskell_toolchain[HaskellToolchainInfo].use_worker else None,
     )
     solibs = {}
     artifact_suffix = get_artifact_suffix(link_style, enable_profiling)
@@ -1211,13 +1212,19 @@ def haskell_binary_impl(ctx: AnalysisContext) -> list[Provider]:
 
     md_file = target_metadata(ctx, sources = ctx.attrs.srcs)
 
+    # Provisional hack to have a worker ID
+    libname = repr(ctx.label.path).replace("//", "_").replace("/", "_") + "_" + ctx.label.name
+    pkgname = libname.replace("_", "-")
+
     compiled = compile(
         ctx,
         link_style,
         enable_profiling = enable_profiling,
         enable_haddock = False,
         md_file = md_file,
         worker = _persistent_worker(ctx),
+        pkgname = pkgname,
+        worker_plugin = ctx.attrs._worker_plugin if hasattr(ctx.attrs, "_worker_plugin") else None,
     )
 
     haskell_toolchain = ctx.attrs._haskell_toolchain[HaskellToolchainInfo]

haskell/worker/worker.bzl

@@ -1,15 +1,22 @@
 worker_libs = [
     "base",
+    "binary",
     "bytestring",
     "containers",
     "deepseq",
     "exceptions",
+    "extra",
     "filepath",
     "ghc",
+    "ghc-persistent-worker-plugin",
     "grpc-haskell",
+    "network",
+    "process",
     "proto3-suite",
     "proto3-wire",
+    "stm",
     "text",
+    "transformers",
     "vector",
     "unix",
 ]
@@ -22,8 +29,12 @@ worker_srcs = [
     "@prelude//haskell/worker/impl/plugin/src:Internal/Error.hs",
     "@prelude//haskell/worker/impl/plugin/src:Internal/Log.hs",
     "@prelude//haskell/worker/impl/plugin/src:Internal/Session.hs",
+    "@prelude//haskell/worker/impl/server/app:Server.hs",
+    "@prelude//haskell/worker/impl/server/app:Pool.hs",
+    "@prelude//haskell/worker/impl/server/app:Worker.hs",
+    "@prelude//haskell/worker/impl/comm/src:Message.hs",
     "@prelude//haskell/worker/impl/buck-worker:Args.hs",
-    "@prelude//haskell/worker/impl/buck-worker:Main.hs",
+    "@prelude//haskell/worker/impl/buck-worker-2:Main.hs",
     "@prelude//haskell/worker/impl/buck-worker:BuckWorker.hs",
 ]
 
@@ -36,6 +47,6 @@ worker_flags = [
     "-XDuplicateRecordFields",
     "-XOverloadedRecordDot",
     "-XStrictData",
-    "-XNoFieldSelectors",
+    # "-XNoFieldSelectors",
     "-XLambdaCase",
 ]