feat: cross-language LTO to inline C TLS shim into Rust FFI

build-common/src/lib.rs

@@ -1,6 +1,10 @@
 // Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+use std::env;
+use std::path::PathBuf;
+use std::process::Command;
+
 #[cfg(not(feature = "cbindgen"))]
 pub fn generate_and_configure_header(_header_name: &str) {}
 #[cfg(not(feature = "cbindgen"))]
@@ -10,3 +14,31 @@ pub fn copy_and_configure_headers() {}
 mod cbindgen;
 #[cfg(feature = "cbindgen")]
 pub use crate::cbindgen::*;
+
+/// Locate the `gcc-ld/` shim directory shipped with the Rust toolchain.
+///
+/// This directory contains an `ld.lld` wrapper that delegates to `rust-lld`.
+/// Passing it via `-B` to the C compiler driver makes it discover rust-lld
+/// before any system-wide lld, which
+///
+/// 1. Avoids the need for a system-wide LLD install.
+/// 2. Picks a recent LLD that matches the Rust toolchain's LLVM version, as opposed to e.g. CentOS
+///    7' LLVM7 which is too old to handle TLSDESC
+pub fn find_rust_lld_dir() -> Option<PathBuf> {
+    let rustc = env::var("RUSTC").unwrap_or_else(|_| "rustc".into());
+    let target = env::var("TARGET").ok()?;
+
+    let output = Command::new(&rustc)
+        .arg("--print")
+        .arg("sysroot")
+        .output()
+        .ok()?;
+
+    let sysroot = std::str::from_utf8(&output.stdout).ok()?.trim();
+    let dir = PathBuf::from(sysroot)
+        .join("lib/rustlib")
+        .join(&target)
+        .join("bin/gcc-ld");
+
+    dir.join("ld.lld").exists().then_some(dir)
+}

libdd-otel-thread-ctx-ffi/README.md

@@ -0,0 +1,61 @@
+# libdd-otel-thread-ctx-ffi
+
+FFI bindings for the OTel thread-level context publisher. Exposes a C API for
+attaching, detaching, and updating per-thread OpenTelemetry context records
+that external readers (e.g. the eBPF profiler) can discover.
+
+Currently Linux-only (x86-64 and aarch64).
+
+## Optimized build (cross-language inlining)
+
+The OTel thread-level context sharing specification requires the use of the
+TLSDESC dialect for the thread-local variable that holds the current context.
+Because (stable) `rustc` doesn't currently provide a way to control the TLS
+dialect, we need to use a small C shim that defines the variable and expose a
+one-line getter. This unfortunately adds one level of indirection (a function
+call) when attaching or detaching a context.
+
+With the right toolchain, it's possible to use Link-Time Optimization (LTO) to
+inline the C wrapper at link time. The requirements are:
+
+- `clang` is available to compile the C shim to LLVM IR (version requirements
+  aren't clear -- tested with clang18 and clang20, but ideally the version
+  should be the same or close to the LLVM version shipped with `rustc`)
+- Either the Rust toolchain ships `lld` or there's a system-wide `lld` install
+  (Rust has been shipping `rust-lld` for a long time now, something like since
+  1.53+, however some musl-based distro like Alpine might have the Rust
+  toolchain without `rust-lld`)
+- `lld` version is at least 18.1 (TLSDESC support)
+
+**If those requirements are met, setting the environment variables
+`CARGO_TARGET_<TARGET>_RUSTFLAGS=-Clinker-plugin-lto -Clinker=clang` and
+`LIBDD_OTEL_THREAD_CTX_INLINE=1` when calling to `cargo` will trigger the
+optimized build where the C shim is inlined.** Here, `<TARGET>` is the target
+triple in screaming snake case.
+
+External environment variables are needed because cross-language LTO requires
+two `rustc` codegen flags (`-Clinker-plugin-lto` and `-Clinker=clang`) that
+cannot be set from a Cargo build script: they must come from `RUSTFLAGS` or
+`.cargo/config.toml`, which can't be entirely automated from Rust only. We
+advise to set those flags via the target-scoped
+`CARGO_TARGET_<TARGET>_RUSTFLAGS` env var so they don't leak to build scripts
+or proc-macros if cross-compiling.
+
+### Build script
+
+The `build-optimized.sh` wrapper script is provided as a convenience and as an
+example.
+
+#### Usage
+
+```bash
+./build-optimized.sh
+```
+
+The script auto-detects the host triple. To cross-compile:
+
+```bash
+./build-optimized.sh --target aarch64-unknown-linux-gnu
+```
+
+Extra arguments are forwarded to `cargo build`.

libdd-otel-thread-ctx-ffi/build-optimized.sh

@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+# Copyright 2026-Present Datadog, Inc. https://www.datadoghq.com/
+# SPDX-License-Identifier: Apache-2.0
+#
+# Build libdd-otel-thread-ctx-ffi with cross-language LTO so the C TLS shim is
+# inlined into the Rust FFI functions, eliminating a function-call indirection
+# on every TLS access.
+#
+# Requirements: clang, lld (rust-lld from the toolchain is used automatically).
+# The requirements are checked by the build.rs script.
+#
+# Usage:
+#   # auto-detect host triple
+#   ./build-optimized.sh
+#   # explicit target
+#   ./build-optimized.sh --target aarch64-unknown-linux-gnu
+#
+# Any extra arguments are forwarded to `cargo build`.
+set -euo pipefail
+
+# Parse --target from args, or auto-detect the host triple.
+TARGET=""
+EXTRA_ARGS=()
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --target)
+            TARGET="$2"; shift 2 ;;
+        --target=*)
+            TARGET="${1#--target=}"; shift ;;
+        *)
+            EXTRA_ARGS+=("$1"); shift ;;
+    esac
+done
+
+if [[ -z "$TARGET" ]]; then
+    TARGET=$(rustc -vV | sed -n 's/host: //p')
+fi
+
+# CARGO_TARGET_<TRIPLE>_RUSTFLAGS scopes the flags to the target only, keeping
+# build scripts and proc-macros unaffected.
+TARGET_ENV=$(echo "$TARGET" | tr 'a-z-' 'A-Z_')
+FLAGS_VAR="CARGO_TARGET_${TARGET_ENV}_RUSTFLAGS"
+EXISTING_FLAGS="${!FLAGS_VAR:-}"
+export "$FLAGS_VAR=${EXISTING_FLAGS:+$EXISTING_FLAGS }-Clinker-plugin-lto -Clinker=clang"
+export LIBDD_OTEL_THREAD_CTX_INLINE=1
+
+cargo build --release \
+    --target "$TARGET" \
+    -p libdd-otel-thread-ctx-ffi \
+    "${EXTRA_ARGS[@]}"
+
+# Sanity-check that the C shim was actually inlined, if `nm` is available.
+if ! command -v nm &>/dev/null; then
+		echo >&2 "WARNING: skipping sanity check that the C TLS shim was inlined (\`nm\` not found)"
+else
+    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+    REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
+    SO="$REPO_ROOT/target/$TARGET/release/liblibdd_otel_thread_ctx_ffi.so"
+
+    if [[ -f "$SO" ]]; then
+        if ! NM_OUTPUT=$(nm "$SO" 2>&1); then
+            echo >&2 "WARNING: command \`nm\` failed on $SO. Skipping sanity check that the C TLS shim was inlined."
+        elif echo "$NM_OUTPUT" | grep -q 'libdd_get_otel_thread_ctx'; then
+            echo >&2 "ERROR: build succeeded but the C TLS shim (libdd_get_otel_thread_ctx_v1) was NOT inlined."
+            echo >&2 "Cross-language LTO may not be working. Check that clang and lld versions are recent enough and compatible with the Rust toolchain's LLVM."
+            exit 1
+        fi
+    fi
+fi

libdd-otel-thread-ctx-ffi/build.rs

@@ -2,42 +2,85 @@
 // SPDX-License-Identifier: Apache-2.0
 extern crate build_common;
 
-use build_common::generate_and_configure_header;
-use std::env;
-use std::path::PathBuf;
-use std::process::Command;
+use build_common::{find_rust_lld_dir, generate_and_configure_header};
+use std::{env, fmt::Display, path::PathBuf, process::Command};
 
-/// Locate the `gcc-ld/` shim directory shipped with the Rust toolchain.
+#[derive(Debug, Eq, PartialEq, Ord, PartialOrd)]
+struct LldVersion {
+    major: u32,
+    minor: u32,
+}
+
+impl Display for LldVersion {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}.{}", self.major, self.minor)
+    }
+}
+
+/// Parse the major and minor version from `ld.lld --version` output.
 ///
-/// This directory contains an `ld.lld` wrapper that delegates to `rust-lld`.
-/// Passing it via `-B` to the C compiler driver makes it discover rust-lld
-/// before any system-wide lld, which
+/// Typical formats:
+///   "LLD 18.1.3 (compatible with GNU linkers)"
+///   "LLD 19.1.0"
+fn system_lld_version() -> Option<LldVersion> {
+    let output = Command::new("ld.lld").arg("--version").output().ok()?;
+    if !output.status.success() {
+        return None;
+    }
+    String::from_utf8_lossy(&output.stdout)
+        .split_whitespace()
+        .find_map(|tok| {
+            let mut splitted = tok.split('.');
+            let major = splitted.next()?.parse::<u32>().ok()?;
+            let minor = splitted.next()?.parse::<u32>().ok()?;
+
+            Some(LldVersion { major, minor })
+        })
+}
+
+/// TLSDESC is supported in LLD from version 18.1.
+const MIN_LLD_VERSION_FOR_TLSDESC: LldVersion = LldVersion {
+    major: 18,
+    minor: 1,
+};
+
+/// Validate that a suitable LLD is available for cross-language LTO.
 ///
-/// 1. Avoid the need of a system-wide LLD install
-/// 2. Pick a recent LLD, as opposed to e.g. CentOS 7' LLVM7 which is too old to handle TLSDESC
-///    relocations properly.
-fn find_rust_lld_dir() -> Option<PathBuf> {
-    let rustc = env::var("RUSTC").unwrap_or_else(|_| "rustc".into());
-    let target = env::var("TARGET").ok()?;
-
-    let output = Command::new(&rustc)
-        .arg("--print")
-        .arg("sysroot")
-        .output()
-        .ok()?;
-
-    let sysroot = std::str::from_utf8(&output.stdout).ok()?.trim();
-    let dir = PathBuf::from(sysroot)
-        .join("lib/rustlib")
-        .join(&target)
-        .join("bin/gcc-ld");
-
-    dir.join("ld.lld").exists().then_some(dir)
+/// Returns the rust-lld `gcc-ld/` directory if found; `None` means the system
+/// `ld.lld` will be used instead. Panics with a clear message when the
+/// requirements are not met.
+fn resolve_lld_for_inline(target_arch: &str) -> Option<PathBuf> {
+    if let Some(dir) = find_rust_lld_dir() {
+        return Some(dir);
+    }
+
+    match system_lld_version() {
+        Some(v) if target_arch != "x86_64" || v >= MIN_LLD_VERSION_FOR_TLSDESC => None,
+        Some(v) => panic!(
+            "LIBDD_OTEL_THREAD_CTX_INLINE requires LLD >= {MIN_LLD_VERSION_FOR_TLSDESC} on \
+             x86-64 (for -mllvm -enable-tlsdesc), but system ld.lld is version {v}. \
+             Install a newer LLD or use a Rust toolchain that bundles rust-lld."
+        ),
+        None => panic!(
+            "LIBDD_OTEL_THREAD_CTX_INLINE requires LLD for cross-language LTO, but neither \
+             rust-lld nor a system ld.lld was found."
+        ),
+    }
 }
 
 fn main() {
     generate_and_configure_header("otel-thread-ctx.h");
+
     let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap();
+    if target_os != "linux" {
+        return;
+    }
+
+    println!("cargo:rerun-if-env-changed=LIBDD_OTEL_THREAD_CTX_INLINE");
+
+    let inline_mode = env::var("LIBDD_OTEL_THREAD_CTX_INLINE").is_ok_and(|v| v == "1");
+    let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
+    let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap();
 
     // Export the TLSDESC thread-local variable to the dynamic symbol table so external readers
     // (e.g. the eBPF profiler) can discover it. Rust's cdylib linker applies a version script with
@@ -50,15 +93,34 @@ fn main() {
     // Merging multiple version scripts is not supported by GNU ld, so we need lld. We prefer the
     // toolchain's bundled rust-lld (LLD 19+ since Rust 1.84) over the system lld (if it even
     // exists). If rust-lld is not found we fall back to whatever `lld` the system provides.
-    if target_os == "linux" {
-        let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
 
+    // If `LIBDD_OTEL_THREAD_CTX_INLINE` is set to `1`, we try to inline the C shim. See the README
+    // for more details.
+    if inline_mode {
+        let rust_lld_dir = resolve_lld_for_inline(&target_arch);
+
+        // Emit link args for ALL link types (not just cdylib) so that test binaries also link
+        // correctly when RUSTFLAGS sets clang as the linker (in practice we should only build/care
+        // about the shared object file in inline mode).
+        if let Some(dir) = rust_lld_dir {
+            println!("cargo:rustc-link-arg=-B{}", dir.display());
+        }
+        println!("cargo:rustc-link-arg=-fuse-ld=lld");
+
+        // On x86-64, tell the LLVM backend to use TLSDESC during LTO codegen.
+        // On aarch64 TLSDESC is the default and the only model.
+        if target_arch == "x86_64" {
+            println!("cargo:rustc-link-arg=-Wl,-mllvm,-enable-tlsdesc");
+        }
+    } else {
+        // Default mode: only the cdylib needs lld (for the version script).
         if let Some(gcc_ld_dir) = find_rust_lld_dir() {
             println!("cargo:rustc-cdylib-link-arg=-B{}", gcc_ld_dir.display());
         }
         println!("cargo:rustc-cdylib-link-arg=-fuse-ld=lld");
-        println!(
-            "cargo:rustc-cdylib-link-arg=-Wl,--version-script={manifest_dir}/tls-dynamic-list.txt"
-        );
     }
+
+    println!(
+        "cargo:rustc-cdylib-link-arg=-Wl,--version-script={manifest_dir}/tls-dynamic-list.txt"
+    );
 }

libdd-otel-thread-ctx/Cargo.toml

@@ -17,4 +17,5 @@ crate-type = ["lib"]
 bench = false
 
 [build-dependencies]
+build_common = { path = "../build-common" }
 cc = "1.1.31"