Start incremental migration from rust-cpython to PyO3 (pantsbuild#12110)

## Rust-cpython vs PyO3

While developing my talk on Rust native extensions, I found PyO3 has had lots of traction recently, including Python cryptography now using it. PyO3 started as a fork, but has since diverged: https://pyo3.rs/v0.13.2/rust_cpython.html.

We went with rust-cpython instead of PyO3 for Rust FFI in pantsbuild#9593 because, at the time, PyO3 required the Rust nightly compiler. While this was a huge win over CFFI, PyO3 would now bring us several benefits:

* Uses procedural macros, rather than a custom `macro_rules!` DSL. This is more natural to work with and better understood by IDEs, rustfmt, etc.
* Great documentation, including a book: https://pyo3.rs/v0.13.2/index.html.
* Better error handling, especially due to not requiring `Python`. See https://pyo3.rs/v0.13.2/rust_cpython.html#error-handling.
~* Can use the Python stable ABI again, meaning we only need to build one wheel per platform, rather than platform x interpreter. (Altho we would need to stop using `PyBuffer` API.) See https://pyo3.rs/v0.13.2/building_and_distribution.html#py_limited_apiabi3.~ We're unlikely to use this due to performance hit.

The community has been very responsive too: PyO3/pyo3#1625, and they have an active Gitter room.

## Incremental migration

To reduce risk and complexity, we use an incremental migration to PyO3 by building two distinct native extensions. At first, `native_engine_pyo3` will only have self-contained functionality, e.g. `PyStubCAS` and, soon, the Nailgun server.

Because the migration is incremental, we can more safely improve our FFI implementation along the way, rather than merely preserving the status quo. For example, this PR makes the `PyStubCAS` API more ergnonomic. 

When we reach critical mass with PyO3, `native_engine_pyo3` will be renamed to `native_engine` and `native_engine` will be changed to `native_engine_cpython`. (Or, we'll remove rust-cpython in one big swoop). This will result in some churn in our Python files (due to import module changing), but I argue that's worth the reduction in risk.

.github/workflows/test-cron.yaml

@@ -102,6 +102,8 @@ jobs:
           '
         path: 'src/python/pants/engine/internals/native_engine.so
 
+          src/python/pants/engine/internals/native_engine_pyo3.so
+
           src/python/pants/engine/internals/native_engine.so.metadata'
     - name: Bootstrap Pants
       run: './pants --version
@@ -131,6 +133,8 @@ jobs:
         name: native_engine.so.${{ matrix.python-version }}.${{ runner.os }}
         path: 'src/python/pants/engine/internals/native_engine.so
 
+          src/python/pants/engine/internals/native_engine_pyo3.so
+
           src/python/pants/engine/internals/native_engine.so.metadata'
     - if: '!contains(env.COMMIT_MESSAGE, ''[ci skip-rust]'')'
       name: Test and Lint Rust
@@ -242,6 +246,8 @@ jobs:
           '
         path: 'src/python/pants/engine/internals/native_engine.so
 
+          src/python/pants/engine/internals/native_engine_pyo3.so
+
           src/python/pants/engine/internals/native_engine.so.metadata'
     - name: Bootstrap Pants
       run: './pants --version
@@ -253,6 +259,8 @@ jobs:
         name: native_engine.so.${{ matrix.python-version }}.${{ runner.os }}
         path: 'src/python/pants/engine/internals/native_engine.so
 
+          src/python/pants/engine/internals/native_engine_pyo3.so
+
           src/python/pants/engine/internals/native_engine.so.metadata'
     - env:
         TMPDIR: ${{ runner.temp }}

.github/workflows/test.yaml

@@ -102,6 +102,8 @@ jobs:
           '
         path: 'src/python/pants/engine/internals/native_engine.so
 
+          src/python/pants/engine/internals/native_engine_pyo3.so
+
           src/python/pants/engine/internals/native_engine.so.metadata'
     - name: Bootstrap Pants
       run: './pants --version
@@ -131,6 +133,8 @@ jobs:
         name: native_engine.so.${{ matrix.python-version }}.${{ runner.os }}
         path: 'src/python/pants/engine/internals/native_engine.so
 
+          src/python/pants/engine/internals/native_engine_pyo3.so
+
           src/python/pants/engine/internals/native_engine.so.metadata'
     - if: '!contains(env.COMMIT_MESSAGE, ''[ci skip-rust]'')'
       name: Test and Lint Rust
@@ -241,6 +245,8 @@ jobs:
           '
         path: 'src/python/pants/engine/internals/native_engine.so
 
+          src/python/pants/engine/internals/native_engine_pyo3.so
+
           src/python/pants/engine/internals/native_engine.so.metadata'
     - name: Bootstrap Pants
       run: './pants --version
@@ -252,6 +258,8 @@ jobs:
         name: native_engine.so.${{ matrix.python-version }}.${{ runner.os }}
         path: 'src/python/pants/engine/internals/native_engine.so
 
+          src/python/pants/engine/internals/native_engine_pyo3.so
+
           src/python/pants/engine/internals/native_engine.so.metadata'
     - env:
         TMPDIR: ${{ runner.temp }}

build-support/bin/generate_github_workflows.py

@@ -34,6 +34,7 @@
 
 NATIVE_ENGINE_SO_FILES = [
     "src/python/pants/engine/internals/native_engine.so",
+    "src/python/pants/engine/internals/native_engine_pyo3.so",
     "src/python/pants/engine/internals/native_engine.so.metadata",
 ]
 

build-support/bin/rust/bootstrap_code.sh

@@ -5,6 +5,9 @@
 
 REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && cd ../../.. && pwd -P)"
 
+# shellcheck source=build-support/common.sh
+source "${REPO_ROOT}/build-support/common.sh"
+
 # Defines:
 # + NATIVE_ROOT: The Rust code directory, ie: src/rust/engine.
 # + MODE: Whether to run in debug or release mode.
@@ -28,7 +31,9 @@ case "${KERNEL}" in
 esac
 
 readonly NATIVE_ENGINE_BINARY="native_engine.so"
+readonly NATIVE_ENGINE_BINARY_PYO3="native_engine_pyo3.so"
 readonly NATIVE_ENGINE_RESOURCE="${REPO_ROOT}/src/python/pants/engine/internals/${NATIVE_ENGINE_BINARY}"
+readonly NATIVE_ENGINE_RESOURCE_PYO3="${REPO_ROOT}/src/python/pants/engine/internals/${NATIVE_ENGINE_BINARY_PYO3}"
 readonly NATIVE_ENGINE_RESOURCE_METADATA="${NATIVE_ENGINE_RESOURCE}.metadata"
 
 function _build_native_code() {
@@ -37,15 +42,21 @@ function _build_native_code() {
   echo "${NATIVE_ROOT}/target/${MODE}/libengine.${LIB_EXTENSION}"
 }
 
+function _build_native_code_pyo3() {
+  # NB: See engine_pyo3/Cargo.toml with regard to the `extension-module` feature.
+  "${REPO_ROOT}/cargo" build --features=engine_pyo3/extension-module ${MODE_FLAG} -p engine_pyo3 || die
+  echo "${NATIVE_ROOT}/target/${MODE}/libengine_pyo3.${LIB_EXTENSION}"
+}
+
 function bootstrap_native_code() {
   # We expose a safety valve to skip compilation iff the user already has `native_engine.so`. This
   # can result in using a stale `native_engine.so`, but we trust that the user knows what
   # they're doing.
   if [[ "${SKIP_NATIVE_ENGINE_SO_BOOTSTRAP}" == "true" ]]; then
-    if [[ ! -f "${NATIVE_ENGINE_RESOURCE}" ]]; then
-      die "You requested to override bootstrapping native_engine.so via the env var" \
-        "SKIP_NATIVE_ENGINE_SO_BOOTSTRAP, but the file does not exist at" \
-        "${NATIVE_ENGINE_RESOURCE}. This is not safe to do."
+    if [[ ! -f "${NATIVE_ENGINE_RESOURCE}" || ! -f "${NATIVE_ENGINE_RESOURCE_PYO3}" ]]; then
+      die "You requested to override bootstrapping native_engine.so and native_engine_pyo3.so via the env var" \
+        "SKIP_NATIVE_ENGINE_SO_BOOTSTRAP, but the files do not exist at" \
+        "${NATIVE_ENGINE_RESOURCE} and ${NATIVE_ENGINE_BINARY_PYO3}. This is not safe to do."
     fi
     return
   fi
@@ -56,14 +67,20 @@ function bootstrap_native_code() {
   if [[ -f "${NATIVE_ENGINE_RESOURCE_METADATA}" ]]; then
     engine_version_in_metadata="$(sed -n 's/^engine_version: //p' "${NATIVE_ENGINE_RESOURCE_METADATA}")"
   fi
-  if [[ ! -f "${NATIVE_ENGINE_RESOURCE}" || "${engine_version_calculated}" != "${engine_version_in_metadata}" ]]; then
-    echo "Building native engine"
-    local -r native_binary="$(_build_native_code)"
+  if [[ ! -f "${NATIVE_ENGINE_RESOURCE}" || ! -f "${NATIVE_ENGINE_RESOURCE_PYO3}" || "${engine_version_calculated}" != "${engine_version_in_metadata}" ]]; then
+    banner "Building native engine..."
+    banner "Building native_engine.so"
+    local -r native_binary="$(_build_native_code)" || die
+    banner "Building native_engine_pyo3.so"
+    local -r native_binary_pyo3="$(_build_native_code_pyo3)" || die
 
     # If bootstrapping the native engine fails, don't attempt to run pants
     # afterwards.
     if [[ ! -f "${native_binary}" ]]; then
-      die "Failed to build native engine."
+      die "Failed to build native engine, file missing at ${native_binary}."
+    fi
+    if [[ ! -f "${native_binary_pyo3}" ]]; then
+      die "Failed to build native engine, file missing at ${native_binary_pyo3}."
     fi
 
     # Pick up Cargo.lock changes if any caused by the `cargo build`.
@@ -72,8 +89,9 @@ function bootstrap_native_code() {
     # Create the native engine resource.
     # NB: On Mac Silicon, for some reason, first removing the old native_engine.so is necessary to avoid the Pants
     #  process from being killed when recompiling.
-    rm -f "${NATIVE_ENGINE_RESOURCE}"
+    rm -f "${NATIVE_ENGINE_RESOURCE}" "${NATIVE_ENGINE_RESOURCE_PYO3}"
     cp "${native_binary}" "${NATIVE_ENGINE_RESOURCE}"
+    cp "${native_binary_pyo3}" "${NATIVE_ENGINE_RESOURCE_PYO3}"
 
     # Create the accompanying metadata file.
     local -r metadata_file=$(mktemp -t pants.native_engine.metadata.XXXXXX)

cargo

@@ -9,6 +9,7 @@ source "${REPO_ROOT}/build-support/common.sh"
 
 PY="$(determine_python)"
 export PY
+export PYO3_PYTHON="${PY}"
 export PYTHON_SYS_EXECUTABLE="${PY}" # Consumed by the cpython crate.
 
 if ! command -v rustup &> /dev/null; then

src/python/pants/engine/internals/.gitignore

@@ -1,2 +1,3 @@
 /native_engine.so
+/native_engine_pyo3.so
 /native_engine.so.metadata

src/python/pants/engine/internals/BUILD

@@ -23,13 +23,13 @@ python_tests(
 
 python_library(
   name='native',
-  sources=['native_engine.pyi'],
+  sources=['native_engine.pyi', 'native_engine_pyo3.pyi'],
   dependencies=[':native_engine'],
 )
 
 resources(
   name='native_engine',
-  sources=['native_engine.so', 'native_engine.so.metadata'],
+  sources=['native_engine.so', 'native_engine_pyo3.so', 'native_engine.so.metadata'],
 )
 
 resources(

src/python/pants/engine/internals/native_engine.pyi

@@ -212,15 +212,6 @@ class PyTasks:
 class PyTypes:
     def __init__(self, **kwargs: Any) -> None: ...
 
-class PyStubCASBuilder:
-    def always_errors(self) -> None: ...
-    def build(self, executor: PyExecutor) -> PyStubCAS: ...
-
-class PyStubCAS:
-    @classmethod
-    def builder(cls) -> PyStubCASBuilder: ...
-    def address(self) -> str: ...
-
 class PyStdioDestination:
     pass
 

src/python/pants/engine/internals/native_engine_pyo3.pyi

@@ -0,0 +1,19 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+# TODO: black and flake8 disagree about the content of this file:
+#   see https://github.com/psf/black/issues/1548
+# flake8: noqa: E302
+
+class PyExecutor:
+    def __init__(self, core_threads: int, max_threads: int) -> None: ...
+
+class PyStubCASBuilder:
+    def always_errors(self) -> PyStubCASBuilder: ...
+    def build(self, executor: PyExecutor) -> PyStubCAS: ...
+
+class PyStubCAS:
+    @classmethod
+    def builder(cls) -> PyStubCASBuilder: ...
+    @property
+    def address(self) -> str: ...