TSAN verification of CachingHostAllocator snapshot data races (pytorch/pytorch#182407)

mock_host_alloc_race.cpp

// Disposable TSAN harness for the CachingHostAllocator context/allocated_ race.
//
// It instantiates the real device-agnostic CachingHostAllocatorImpl template
// with a mock backend (malloc-backed, no CUDA), then hammers alloc/free from
// several threads while another thread repeatedly calls getSegments() and a
// controller toggles recording on/off.  The pre-fix code writes
// context_when_allocated_ (a shared_ptr) unlocked in the allocate fast path and
// races the locked read in getSegments; it also writes allocated_ under a
// different mutex than getSegments reads it.  Build with -fsanitize=thread.

#include <ATen/core/CachingHostAllocator.h>

#include <atomic>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <thread>
#include <vector>

namespace {

// A trivial "stream" and "event" type so the template can be instantiated
// without any CUDA dependency.  We never enter a capture, never record events,
// so most of these are only needed to satisfy the template's type parameters.
struct MockStream {
  MockStream() = default;
  explicit MockStream(c10::Stream) {}
  bool operator==(const MockStream&) const { return true; }
  // Only needed so the (never-taken) capture-predicate path in the template
  // type-checks; no capture is ever active in this harness.
  operator c10::Stream() const {
    return c10::Stream(c10::Stream::UNSAFE, c10::Device(c10::DeviceType::CPU, 0),
                       0);
  }
};
struct MockEvent {};

}  // namespace

template <>
struct std::hash<MockStream> {
  size_t operator()(const MockStream&) const noexcept { return 0; }
};

namespace {

class MockHostAllocator
    : public at::CachingHostAllocatorImpl<MockStream, MockEvent> {
 public:
  // Deterministic: no background event-processing thread.
  bool pinned_use_background_threads() override { return false; }

 private:
  void allocate_host_memory(size_t size, void** ptr) override {
    *ptr = std::malloc(size);
  }
  void free_block(at::HostBlock<MockStream>* block) override {
    std::free(block->ptr_);
  }
  // We never record streams (no stream ever attached to a block), so free()
  // always takes the maybe_cache_block path and blocks get reused.
  void record_stream(std::optional<std::vector<MockEvent>>&,
                     MockStream) override {}
  bool query_event(MockEvent&) override { return true; }
  MockStream get_current_stream() const override { return MockStream{}; }
  bool stream_is_capturing(MockStream) const override { return false; }
};

std::shared_ptr<c10::GatheredContext> makeContext() {
  // Any non-null GatheredContext; the base struct is enough to exercise the
  // shared_ptr refcount traffic that the race corrupts.
  return std::make_shared<c10::GatheredContext>();
}

}  // namespace

int main() {
  MockHostAllocator alloc;
  std::atomic<bool> stop{false};

  // Recording ON with ALLOC context so maybeGatherContext returns non-null and
  // context_when_allocated_ is actually populated.
  alloc.recordHistory(
      /*enabled=*/true, makeContext, /*max_entries=*/100000,
      c10::CachingDeviceAllocator::RecordContext::ALL, /*clearHistory=*/false);

  // Coverage counters (verified after the run to prove the intended paths ran).
  std::atomic<uint64_t> total_allocs{0};
  std::atomic<uint64_t> off_phase_allocs{0};  // allocs while recording is off
  std::atomic<uint64_t> reader_nonnull_ctx{0};  // segments seen with a context

  // Allocator worker threads: alloc a fixed-size block then free it, so blocks
  // churn through the free list and get reused via the fast path.
  constexpr size_t kSize = 512;
  std::vector<std::thread> workers;
  for (int t = 0; t < 4; ++t) {
    workers.emplace_back([&] {
      while (!stop.load(std::memory_order_relaxed)) {
        for (int i = 0; i < 200; ++i) {
          bool history_off = !alloc.isHistoryEnabled();
          auto [ptr, ctx] = alloc.allocate(kSize);
          total_allocs.fetch_add(1, std::memory_order_relaxed);
          if (history_off) {
            off_phase_allocs.fetch_add(1, std::memory_order_relaxed);
          }
          if (ctx) {
            alloc.free(ctx);
          }
        }
      }
    });
  }

  // Reader thread: repeatedly snapshot segments (reads context_when_allocated_
  // and allocated_ under block->mutex_).
  std::thread reader([&] {
    while (!stop.load(std::memory_order_relaxed)) {
      auto segs = alloc.getSegments();
      for (const auto& s : segs) {
        if (s.context_when_allocated) {
          reader_nonnull_ctx.fetch_add(1, std::memory_order_relaxed);
        }
      }
    }
  });

  // Controller thread: toggle recording on/off for the whole run, with a short
  // sleep so disable windows are spread across the entire run (not just a
  // startup burst). This is what exercises the fixed code's disable-mid-
  // snapshot else-branch concurrently with getSegments.
  std::thread toggler([&] {
    bool on = true;
    while (!stop.load(std::memory_order_relaxed)) {
      on = !on;
      if (on) {
        alloc.recordHistory(
            true, makeContext, 100000,
            c10::CachingDeviceAllocator::RecordContext::ALL, false);
      } else {
        alloc.recordHistory(
            false, makeContext, 100000,
            c10::CachingDeviceAllocator::RecordContext::NEVER, false);
      }
      std::this_thread::sleep_for(std::chrono::microseconds(200));
    }
  });

  std::this_thread::sleep_for(std::chrono::seconds(3));
  stop.store(true, std::memory_order_relaxed);
  for (auto& w : workers) w.join();
  reader.join();
  toggler.join();

  // Positive-coverage checks: the run is only meaningful if the intended paths
  // actually executed. A clean TSAN result with these all satisfied is strong
  // evidence; a clean result with these failing would be a false negative.
  auto stats = alloc.getStats();
  std::fprintf(stderr,
               "coverage: total_allocs=%llu off_phase_allocs=%llu "
               "num_host_alloc=%lld reader_nonnull_ctx=%llu\n",
               (unsigned long long)total_allocs.load(),
               (unsigned long long)off_phase_allocs.load(),
               (long long)stats.num_host_alloc,
               (unsigned long long)reader_nonnull_ctx.load());
  bool ok = true;
  // Fast path (get_free_block hit) must dominate: few real host allocations
  // versus millions of iterations.
  if (stats.num_host_alloc == 0 || (uint64_t)stats.num_host_alloc >=
                                        total_allocs.load()) {
    std::fprintf(stderr, "FAIL: fast-path reuse did not dominate\n");
    ok = false;
  }
  // The disable-mid-snapshot else-branch requires allocations while off.
  if (off_phase_allocs.load() == 0) {
    std::fprintf(stderr, "FAIL: no allocations occurred while recording off\n");
    ok = false;
  }
  // The reader must have observed populated contexts (else it never read the
  // racy field in a meaningful state).
  if (reader_nonnull_ctx.load() == 0) {
    std::fprintf(stderr, "FAIL: reader never observed a non-null context\n");
    ok = false;
  }
  std::fprintf(stderr, "%s\n", ok ? "coverage OK" : "coverage INSUFFICIENT");
  return ok ? 0 : 2;
}

REPORT.md

TSAN verification of the CachingHostAllocator snapshot data races

This report documents a ThreadSanitizer (TSAN) reproduction of two data races in aten/src/ATen/core/CachingHostAllocator.h that were introduced/exposed by the "Support memory snapshot for CPU pinned memory" change, and confirms the committed fixes eliminate them.

Background

The host (pinned-memory) caching allocator uses fine-grained per-block locking (block->mutex_), unlike the single-lock CUDACachingAllocator. Adding getSegments() (reached only via _snapshot() when history recording is enabled) created a concurrent reader of two per-block fields:

1. context_when_allocated_ (a std::shared_ptr<GatheredContext>) — the allocate fast path (block reused from the free list) wrote it without block->mutex_, while getSegments() reads it under block->mutex_. The block stays in pool.blocks_ the whole time it is cached, so a concurrent snapshot can read the shared_ptr while the allocating thread overwrites it — a shared_ptr control-block refcount race (corruption / potential use-after-free, not a benign torn read).

2. allocated_ (a plain bool) — written in get_free_block() under the free-list mutex, read in getSegments() under the block mutex. Two different locks guarding the same field is a data race (benign on x86 in practice, but UB and a guaranteed TSAN report).

Why not TSAN the real allocator directly

PyTorch's build makes CUDA and TSAN mutually exclusive (CMakeLists.txt: cmake_dependent_option(USE_CUDA "Use CUDA" ON "NOT USE_TSAN" OFF)), and the existing cuda_caching_host_allocator_test.cpp tests are single-threaded and gated on at::cuda::is_available(). So the real CUDA pinned-memory path cannot be TSAN'd, and the existing tests would never trip the race even with CUDA.

The racing code, however, lives entirely in the device-agnostic template CachingHostAllocatorImpl<S, E> in the header. The CUDA-specific parts are backend virtuals (allocate_host_memory, record_stream, query_event, stream_is_capturing, ...). So the race is reproducible with zero CUDA by instantiating the real template with a malloc-backed mock backend.

Harness

A standalone gtest-free program (mock_host_alloc_race.cpp) that:

• Subclasses the real CachingHostAllocatorImpl with a mock backend (malloc/free, no-op stream/event ops, no background threads).
• Runs 4 worker threads doing allocate(512) -> free() in a loop, so blocks churn through the free list and get reused via the fast path.
• Runs 1 reader thread calling getSegments() in a loop, counting how many segments it observed with a non-null context_when_allocated.
• Runs 1 controller thread toggling recordHistory(on/off) continuously for the whole run (with a 200us sleep between toggles) so disable windows are spread across the entire run, exercising the fixed code's disable-mid-snapshot else-branch concurrently with getSegments.

After the run it asserts positive coverage, so that a clean TSAN result is not a false negative from the racy paths never executing:

• fast-path reuse dominated (num_host_alloc tiny vs. total allocations),
• allocations occurred while recording was off (the else-branch was reachable),
• the reader observed non-null contexts (it read the racy field populated).

Built with g++ -std=c++17 -fsanitize=thread -O1 -g, linked against the (uninstrumented) libc10.so. This does not undermine the result, for two precise reasons: (1) both sides of both races are in the instrumented TU -- allocate, get_free_block, getSegments, maybe_cache_block are all header template code compiled into the instrumented object, and the shared_ptr refcount/pointer ops are libstdc++ header code also compiled into it (which is why TSAN sees the refcount race at shared_ptr_base.h); (2) happens-before comes from pthread interceptors, not instrumentation -- TSAN interposes pthread_mutex_lock/unlock via the dynamic linker, so std::mutex establishes HB even through uninstrumented libc10. The only per-iteration libc10 surface is config accessors and RingBuffer/clock plumbing, none of which touches the two fields under test.

Results

Each configuration was run multiple times; results were stable (not flaky).

Version	TSAN races	What TSAN reported
Pre-fix: both fixes reverted	7	the shared_ptr context race (size-8 reads in getSegments holding the block mutex vs. the unlocked fast-path write holding none, plus refcount atomics) and the allocated_ bool race (size-1 read in getSegments under the block mutex vs. the get_free_block write under the free-list mutex -- a different lock)
Fixed (committed HEAD, both fixes)	0	clean, reproduced across repeated runs, with positive coverage satisfied

The pre-fix reports show the exact locking asymmetry: the reader thread holds the block mutex while the writer holds none (context race) or a different mutex (bool race). The fixed run is clean while its coverage assertions confirm the racy paths actually executed -- in a representative fixed run: ~282k total allocations, only 4 real host allocations (so fast-path reuse dominated), ~206k allocations performed while recording was off (so the disable-window else-branch was heavily exercised concurrently with the reader), and ~72k segments observed by the reader with a non-null context. So the clean result is a true negative, not the harness failing to schedule the race.

Note: line numbers in the TSAN excerpts below are from the reverted (pre-fix) working tree used for the positive control, and differ slightly from committed HEAD (e.g. the getSegments read is at :1143/:1145 at HEAD). They identify the racing statements, not exact HEAD locations.

Representative pre-fix TSAN output (context race)

WARNING: ThreadSanitizer: data race
  Read of size 8 ... by thread T5 (mutexes: write M0, write M1):
    #0 std::__shared_ptr<c10::GatheredContext, ...>::operator=(... const&) shared_ptr_base.h:1523
    ...
    getSegments  aten/src/ATen/core/CachingHostAllocator.h:1133
  Previous write of size 8 ... by thread T1 (no mutexes):
    #0 std::__shared_ptr<c10::GatheredContext, ...>::operator=(...&&)
    ...
    allocate     aten/src/ATen/core/CachingHostAllocator.h:354

Representative pre-fix TSAN output (allocated_ bool race)

WARNING: ThreadSanitizer: data race
  Read of size 1 ... by thread T5 (mutexes: write M0, write M1):
    getSegments      aten/src/ATen/core/CachingHostAllocator.h:1135
  Previous write of size 1 ... by thread T1 (mutexes: write M2):
    get_free_block   aten/src/ATen/core/CachingHostAllocator.h:685

(M0/M1 = block mutex; M2 = the free-list mutex — different locks, same field.)

The fixes (committed)

1. Fast-path context write — gate the block->mutex_ lock behind record_history_, and in the recording-off branch only take the lock when there is actually something to write (else if (context || block->context_when_allocated_)). This closes the disable-mid-snapshot window (a cached block can carry a stale non-null context that maybe_cache_block only resets while recording) while keeping the steady- state hot path lock-free and write-free.

2. allocated_ flag — changed to std::atomic<bool> with relaxed load/store at all sites. Relaxed ordering suffices: it is an independent flag, and every non-atomic invariant around it is still enforced by the existing mutexes.

Note: the fast-path lock in fix (1) is only taken when recording is enabled, so there is no added overhead on the common cached-alloc/free hot path when memory history recording is off.