[DO NOT MERGE] Reproduce flaky Rust parallel cargo test SIGSEGV

c/src/core/c_api.cpp

@@ -347,19 +347,24 @@ extern "C" cuvsError_t cuvsRMMMemoryResourceReset()
   });
 }
 
-thread_local std::unique_ptr<rmm::mr::pinned_host_memory_resource> pinned_mr;
+namespace {
+// One process-wide resource: pinned_host_memory_resource is stateless, so a buffer allocated on
+// one thread can be freed on another (a thread_local resource left it null on other threads).
+rmm::mr::pinned_host_memory_resource& pinned_mr()
+{
+  static rmm::mr::pinned_host_memory_resource mr;
+  return mr;
+}
+}  // namespace
 
 extern "C" cuvsError_t cuvsRMMHostAlloc(void** ptr, size_t bytes)
 {
-  return cuvs::core::translate_exceptions([=] {
-    if (pinned_mr == nullptr) { pinned_mr = std::make_unique<rmm::mr::pinned_host_memory_resource>(); }
-    *ptr = pinned_mr->allocate_sync(bytes);
-  });
+  return cuvs::core::translate_exceptions([=] { *ptr = pinned_mr().allocate_sync(bytes); });
 }
 
 extern "C" cuvsError_t cuvsRMMHostFree(void* ptr, size_t bytes)
 {
-  return cuvs::core::translate_exceptions([=] { pinned_mr->deallocate_sync(ptr, bytes); });
+  return cuvs::core::translate_exceptions([=] { pinned_mr().deallocate_sync(ptr, bytes); });
 }
 
 thread_local std::string last_error_text = "";

c/tests/CMakeLists.txt

@@ -109,6 +109,7 @@ ConfigureTest(NAME cuvs_c_headers PATH core/headers.c)
 ConfigureTest(NAME cuvs_c_test PATH core/c_api.c)
 ConfigureTest(NAME cuvs_c_resource_pool_test PATH core/resource_pool.cpp)
 ConfigureTest(NAME cuvs_c_resource_stream_pool_test PATH core/resource_stream_pool.cpp)
+ConfigureTest(NAME cuvs_c_host_alloc_test PATH core/host_alloc.cpp)
 target_link_libraries(cuvs_c_test PRIVATE CUDA::cudart_static)
 target_link_libraries(cuvs_c_resource_pool_test PRIVATE CUDA::cudart_static)
 target_link_libraries(cuvs_c_resource_stream_pool_test PRIVATE CUDA::cudart_static)

c/tests/core/host_alloc.cpp

@@ -0,0 +1,66 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cuvs/core/c_api.h>
+
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <cstddef>
+#include <thread>
+#include <vector>
+
+// Pinned host buffers may be freed on a different thread than the one that allocated them, e.g.
+// when a binding moves the owning object across threads.
+TEST(RMMHostAlloc, FreeOnAnotherThread)
+{
+  constexpr std::size_t kBytes = 1024;
+  void* ptr                    = nullptr;
+  cuvsError_t alloc_status     = CUVS_ERROR;
+  std::thread([&] { alloc_status = cuvsRMMHostAlloc(&ptr, kBytes); }).join();
+  ASSERT_EQ(alloc_status, CUVS_SUCCESS);
+  ASSERT_NE(ptr, nullptr);
+
+  cuvsError_t free_status = CUVS_ERROR;
+  std::thread([&] { free_status = cuvsRMMHostFree(ptr, kBytes); }).join();
+  EXPECT_EQ(free_status, CUVS_SUCCESS);
+}
+
+TEST(RMMHostAlloc, ConcurrentAllocAndFree)
+{
+  constexpr int kThreads       = 8;
+  constexpr int kIters         = 50;
+  constexpr std::size_t kBytes = 4096;
+  std::vector<void*> ptrs(kThreads * kIters, nullptr);
+  std::atomic<int> failures{0};
+
+  std::vector<std::thread> threads;
+  for (int t = 0; t < kThreads; ++t) {
+    threads.emplace_back([&, t] {
+      for (int i = 0; i < kIters; ++i) {
+        if (cuvsRMMHostAlloc(&ptrs[t * kIters + i], kBytes) != CUVS_SUCCESS) { failures++; }
+      }
+    });
+  }
+  for (auto& th : threads) {
+    th.join();
+  }
+  threads.clear();
+  ASSERT_EQ(failures.load(), 0);
+
+  // Free every buffer from a different thread than the one that allocated it.
+  for (int t = 0; t < kThreads; ++t) {
+    threads.emplace_back([&, t] {
+      int const owner = (t + 1) % kThreads;
+      for (int i = 0; i < kIters; ++i) {
+        if (cuvsRMMHostFree(ptrs[owner * kIters + i], kBytes) != CUVS_SUCCESS) { failures++; }
+      }
+    });
+  }
+  for (auto& th : threads) {
+    th.join();
+  }
+  EXPECT_EQ(failures.load(), 0);
+}

ci/build_rust.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
@@ -42,12 +42,11 @@ rapids-logger "Begin Rust build"
 
 sccache --stop-server 2>/dev/null || true
 
-bash ./build.sh rust
+INITIAL_RC=0
+bash ./build.sh rust || INITIAL_RC=$?
+echo "INITIAL build.sh rust (cargo test, default parallelism) rc=${INITIAL_RC}"
 
-sccache --show-adv-stats
+sccache --show-adv-stats || true
 sccache --stop-server >/dev/null 2>&1 || true
 
-# Also test out that we can publish cuvs-sys via a dry-run
-pushd ./rust/cuvs-sys
-cargo publish --dry-run
-popd
+bash ci/stress_rust.sh

ci/stress_rust.sh

@@ -0,0 +1,113 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# CI-only: repeat the Rust test suite with parallel test threads (issue #2233).
+set -uo pipefail
+
+STRESS_MINUTES=${STRESS_MINUTES:-110}
+ENV_NAME=${ENV_NAME:-rust}
+RUST_DIR=${RUST_DIR:-rust}
+
+. /opt/conda/etc/profile.d/conda.sh
+set +u
+conda activate "${ENV_NAME}"
+set -u
+
+rapids-logger "Installing gdb and compute-sanitizer"
+rapids-mamba-retry install -y -n "${ENV_NAME}" gdb "cuda-sanitizer-api=${RAPIDS_CUDA_VERSION%.*}" || \
+  rapids-mamba-retry install -y -n "${ENV_NAME}" gdb || true
+which gdb compute-sanitizer || true
+
+nvidia-smi || true
+nproc
+cat /proc/sys/kernel/core_pattern || true
+ulimit -c unlimited || true
+ulimit -c
+
+cd "${RUST_DIR}" || exit 1
+export RUST_BACKTRACE=full
+cargo test -p cuvs --lib --no-run 2>&1 | tee norun.log
+TEST_BIN=$(sed -n 's/.*Executable unittests src\/lib.rs (\(.*\))/\1/p' norun.log | tail -1)
+TEST_BIN=$(realpath "${TEST_BIN}" 2>/dev/null || true)
+echo "TEST_BIN=${TEST_BIN}"
+if [ ! -x "${TEST_BIN}" ]; then echo "Failed to build the test binary"; exit 1; fi
+"${TEST_BIN}" --list 2>/dev/null | tail -3
+HAVE_STRESS=0
+if "${TEST_BIN}" --list --ignored 2>/dev/null | grep -q "stress_"; then HAVE_STRESS=1; fi
+
+LOGDIR=$(mktemp -d)
+mkdir -p "${LOGDIR}/cores"
+GDB_CMDS=(-q -batch -ex "set pagination off" -ex "set print thread-events off"
+  -ex "handle SIGPIPE nostop noprint pass" -ex "handle SIG32 nostop noprint pass"
+  -ex "handle SIG33 nostop noprint pass" -ex "handle SIG34 nostop noprint pass"
+  -ex "handle SIGUSR1 nostop noprint pass" -ex "handle SIGUSR2 nostop noprint pass"
+  -ex run -ex "info threads" -ex "thread apply all bt" -ex "info sharedlibrary")
+
+THREADS_LIST=("" 8 16 32 64)
+deadline=$(( $(date +%s) + STRESS_MINUTES * 60 ))
+iter=0
+fails=0
+crashes=0
+declare -A per_mode_runs per_mode_fails
+while [ "$(date +%s)" -lt "${deadline}" ]; do
+  iter=$((iter + 1))
+  t=${THREADS_LIST[$((iter % ${#THREADS_LIST[@]}))]}
+  targs=()
+  [ -n "${t}" ] && targs=(--test-threads="${t}")
+  mode="native"
+  if (( iter % 4 == 0 )); then mode="gdb"; fi
+  if (( iter % 6 == 0 )) && (( HAVE_STRESS )); then mode="stress"; fi
+  key="${mode}-t${t:-default}"
+  per_mode_runs[$key]=$(( ${per_mode_runs[$key]:-0} + 1 ))
+  log="${LOGDIR}/iter${iter}.log"
+  start=$(date +%s)
+  case "${mode}" in
+    native)
+      (cd "${LOGDIR}/cores" && timeout 900 "${TEST_BIN}" "${targs[@]}" > "${log}" 2>&1); rc=$? ;;
+    gdb)
+      timeout 1800 gdb "${GDB_CMDS[@]}" --args "${TEST_BIN}" "${targs[@]}" > "${log}" 2>&1; rc=$?
+      # gdb returns 0 even when the inferior crashed; detect the signal in the log.
+      if grep -q "received signal SIG" "${log}"; then rc=139; fi ;;
+    stress)
+      CUVS_STRESS_THREADS=16 CUVS_STRESS_ITERS=20 timeout 1200 "${TEST_BIN}" --ignored --test-threads=2 stress_ > "${log}" 2>&1; rc=$? ;;
+  esac
+  dur=$(( $(date +%s) - start ))
+  summary=$(grep -E "^test result:" "${log}" | tail -1)
+  echo "ITER ${iter} mode=${mode} threads=${t:-default} rc=${rc} dur=${dur}s ${summary}"
+  if [ "${rc}" -ne 0 ]; then
+    fails=$((fails + 1))
+    per_mode_fails[$key]=$(( ${per_mode_fails[$key]:-0} + 1 ))
+    if [ "${rc}" -ge 128 ] && [ "${rc}" -ne 124 ]; then crashes=$((crashes + 1)); fi
+    echo "::group::FAILURE iter ${iter} (${key}) rc=${rc}"
+    if [ "${mode}" = "gdb" ]; then
+      sed -n '/received signal SIG/,$p' "${log}" | head -2000
+      grep -E "^test .*(FAILED|ok)$" "${log}" | tail -40
+    else
+      grep -vE "^test .* \.\.\. ok$" "${log}" | tail -300
+    fi
+    echo "::endgroup::"
+    ls -la "${LOGDIR}/cores"
+    for core in "${LOGDIR}"/cores/core*; do
+      [ -e "${core}" ] || continue
+      gdb -q -batch -ex "info threads" -ex "thread apply all bt" "${TEST_BIN}" "${core}" 2>&1 | head -3000
+      rm -f "${core}"
+    done
+  fi
+done
+
+echo "STRESS SUMMARY: iterations=${iter} failures=${fails} crashes=${crashes}"
+for k in "${!per_mode_runs[@]}"; do
+  echo "  ${k}: runs=${per_mode_runs[$k]} fails=${per_mode_fails[$k]:-0}"
+done
+
+if command -v compute-sanitizer > /dev/null; then
+  rapids-logger "compute-sanitizer memcheck, --test-threads=8"
+  timeout 2400 compute-sanitizer --tool memcheck --error-exitcode 1 "${TEST_BIN}" --test-threads=8 2>&1 | tail -200
+  echo "memcheck rc=${PIPESTATUS[0]}"
+  rapids-logger "compute-sanitizer racecheck-free: synccheck skipped; initcheck on ivf_flat"
+  timeout 1200 compute-sanitizer --tool memcheck --error-exitcode 1 "${TEST_BIN}" --test-threads=16 ivf_flat 2>&1 | tail -100
+  echo "memcheck ivf_flat rc=${PIPESTATUS[0]}"
+fi
+
+if [ "${fails}" -ne 0 ]; then exit 1; fi

ci/stress_rust_old.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# CI-only: run the stress loop on the Rust crate as it was when #2233 was filed (before the
+# DLPack/lifetime rewrite in #2245), against released libcuvs 26.06 (the reporter's setup).
+set -euo pipefail
+
+OLD_SHA=78135be4af3c7b458fd18b7b570c9b6978796daa
+
+. /opt/conda/etc/profile.d/conda.sh
+conda config --set channel_priority flexible
+
+rapids-dependency-file-generator \
+  --output conda \
+  --file-key rust \
+  --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION}" \
+  | sed -e 's/==26\.12\.\*,>=0\.0\.0a0/==26.06.*/' -e 's/==26\.12\.\*/==26.06.*/' \
+        -e 's/^- rapidsai-nightly$/- rapidsai/' -e 's/^  - rapidsai-nightly$/  - rapidsai/' \
+  | tee env_old.yaml
+
+rapids-mamba-retry env create --yes -f env_old.yaml -n rustold
+set +u
+conda activate rustold
+set -u
+conda list | grep -E "libcuvs|libraft|librmm|cuda-version|^rust " || true
+
+git fetch --depth 1 https://github.com/NVIDIA/cuvs.git "${OLD_SHA}"
+mkdir -p old
+git archive "${OLD_SHA}" rust | tar -x -C old
+ls old/rust
+# The old crate is versioned 26.8.0 and requires a same-minor libcuvs; the reporter used 26.06.
+sed -i 's/26\.8\.0/26.6.0/g' old/rust/Cargo.toml
+grep -n version old/rust/Cargo.toml
+
+ENV_NAME=rustold RUST_DIR=old/rust bash ci/stress_rust.sh

ci/test_c_host_alloc.sh

@@ -0,0 +1,50 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# CI-only: run c/tests/core/host_alloc.cpp (pinned host buffer freed on another thread) against
+# the nightly libcuvs (without the fix) and against this PR's libcuvs (with the fix).
+set -uo pipefail
+
+. /opt/conda/etc/profile.d/conda.sh
+CUDA_MM=${RAPIDS_CUDA_VERSION%.*}
+CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")")
+nvidia-smi || true
+
+build_and_run() {
+  local label=$1
+  local cxx
+  cxx=$(ls "${CONDA_PREFIX}"/bin/*-linux-gnu-g++ | head -1)
+  "${cxx}" -std=c++20 -O1 -g c/tests/core/host_alloc.cpp -o "host_alloc_${label}" \
+    -I"${CONDA_PREFIX}/include" -I"${CONDA_PREFIX}/targets/x86_64-linux/include" \
+    -L"${CONDA_PREFIX}/lib" -Wl,-rpath,"${CONDA_PREFIX}/lib" \
+    -lcuvs_c -lgtest -lgtest_main -pthread || { echo "BUILD FAILED (${label})"; return; }
+  local crashes=0 fails=0
+  for i in $(seq 1 10); do
+    "./host_alloc_${label}" > "run_${label}.log" 2>&1
+    rc=$?
+    [ "${i}" -eq 1 ] && tail -25 "run_${label}.log"
+    if [ "${rc}" -ge 128 ]; then crashes=$((crashes + 1)); elif [ "${rc}" -ne 0 ]; then fails=$((fails + 1)); fi
+    echo "${label} run ${i}: rc=${rc}"
+  done
+  echo "RESULT ${label}: 10 runs, ${crashes} crashes, ${fails} other failures"
+}
+
+PKGS=("cuda-version=${CUDA_MM}" cuda-cudart-dev dlpack gtest "gxx_linux-64=14.*" "sysroot_linux-64=2.28")
+
+rapids-logger "Baseline: nightly libcuvs without the fix"
+rapids-mamba-retry create -y -n cbase -c rapidsai-nightly -c conda-forge "libcuvs=26.12" "${PKGS[@]}"
+set +u; conda activate cbase; set -u
+conda list | grep -E "^libcuvs |^librmm "
+build_and_run baseline
+set +u; conda deactivate; set -u
+
+rapids-logger "Fixed: this PR's libcuvs"
+rapids-mamba-retry create -y -n cfix -c "${CPP_CHANNEL}" -c rapidsai-nightly -c conda-forge \
+  "libcuvs=26.12" "libcuvs-tests=26.12" "${PKGS[@]}"
+set +u; conda activate cfix; set -u
+conda list | grep -E "^libcuvs |^libcuvs-tests |^librmm "
+build_and_run fixed
+GT="${CONDA_PREFIX}/bin/gtests/libcuvs/cuvs_c_host_alloc_test"
+ls -la "${GT}" || true
+"${GT}" --gtest_repeat=100 2>&1 | tail -15
+echo "installed cuvs_c_host_alloc_test rc=${PIPESTATUS[0]}"

cpp/src/cluster/detail/kmeans_balanced.cuh

@@ -42,6 +42,7 @@
 #include <thrust/transform.h>
 
 #include <algorithm>
+#include <atomic>
 #include <limits>
 #include <optional>
 #include <tuple>
@@ -802,7 +803,8 @@ auto adjust_centers(const raft::resources& handle,
                                       601,  659,  733,  809,  863,  941,  1013, 1069, 1151, 1223,
                                       1291, 1373, 1451, 1511, 1583, 1657, 1733, 1811, 1889, 1987,
                                       2053, 2129, 2213, 2287, 2357, 2423, 2531, 2617, 2687, 2741};
-  static IdxT i_primes = 0;
+  // Shared by all threads that build balanced k-means concurrently, so it must be atomic.
+  static std::atomic<std::size_t> i_primes{0};
 
   auto average         = static_cast<MathT>(n_rows) / static_cast<MathT>(n_clusters);
   auto lower_threshold = average * balance_lower_tolerance;
@@ -839,8 +841,7 @@ auto adjust_centers(const raft::resources& handle,
 
   IdxT ofst;
   do {
-    i_primes = (i_primes + 1) % kPrimes.size();
-    ofst     = kPrimes[i_primes];
+    ofst = kPrimes[(i_primes.fetch_add(1, std::memory_order_relaxed) + 1) % kPrimes.size()];
   } while (n_rows % ofst == 0);
 
   rmm::device_uvector<IdxT> receiver_clusters(n_pairs, stream, device_memory);

rust/cuvs/src/lib.rs

@@ -22,6 +22,9 @@ pub mod version;
 #[cfg(test)]
 pub(crate) mod test_utils;
 
+#[cfg(test)]
+mod thread_stress;
+
 pub use dlpack::{AsDlTensor, AsDlTensorMut, DLPackError, DLTensorView, DLTensorViewMut, DType};
 pub use error::LibraryError;
 pub use resources::Resources;