Stress test for MADV_FREE_REUSABLE / MADV_FREE_REUSE

test_madvise.cpp

#include <pthread.h>
#include <sys/mman.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdatomic.h>
#include <errno.h>
#include <stdint.h>
#include <sched.h>

// Stress test spawning worker threads that do the following:
//
//   1. madvise(x, y, MADV_FREE_REUSABLE)
//   2. madvise(x, y, MADV_FREE_REUSE)
//   3. Write a MAGIC_BYTE to these pages
//   4. Sleep up to 30 seconds
//   5. Check the pages still contain the MAGIC_BYTE value
//
// It appears that due to a kernel bug on Tahoe, this check fails very
// intermittently but usually within 10-15 minutes on an M4 Pro with 48 GB RAM.
// The kernel somehow zeroes a memory page during step 4.
//
// This test also creates a pool of "pressure" threads that allocate/touch
// memory.
//
// Run like this:
//
//   $ clang++ -o test_madvise test_madvise.cpp -lpthread -O2
//   $ ./test_madvise
//
// Then wait 15-30 minutes.

static constexpr size_t PAGE_SIZE = 16384;
static constexpr size_t PAGES_PER_THREAD = 12288;  // 192MB per thread
static constexpr size_t SLOT_SIZE = PAGES_PER_THREAD * PAGE_SIZE;
static constexpr int THREADS_PER_CHUNK = 32;
static constexpr size_t CHUNK_SIZE = THREADS_PER_CHUNK * SLOT_SIZE;
static constexpr int NUM_CHUNKS = 4;  // 4 chunks, 32 threads each = 128 workers
static constexpr int NUM_ROUNDS = 200;
static constexpr uint8_t MAGIC_BYTE = 0xAF;

static atomic_int started_threads = 0;
static atomic_int go = 0;

struct worker_arg {
  void* region;
  size_t length;
  int id;
};

static void* worker(void* arg) {
  struct worker_arg* wa = (struct worker_arg*)arg;
  void* region = wa->region;
  size_t length = wa->length;
  int id = wa->id;

  atomic_fetch_add(&started_threads, 1);
  while (!atomic_load(&go)) {
    usleep(100);
  }

  unsigned seed = (unsigned)id;
  for (int round = 0; round < NUM_ROUNDS; round++) {
    // Use MADV_FREE_REUSABLE.
    int status;
    do {
      status = madvise(region, length, MADV_FREE_REUSABLE);
    } while (status == -1 && errno == EAGAIN);
    if (status != 0) {
      fprintf(stderr, "thread %d: REUSABLE failed: %s\n", id, strerror(errno));
      abort();
    }

    sched_yield();

    // Use MADV_FREE_REUSE.
    while (1) {
      int res = madvise(region, length, MADV_FREE_REUSE);
      if (res == 0) break;
      if (errno != EAGAIN) {
        fprintf(stderr, "thread %d: REUSE failed: %s\n", id, strerror(errno));
        abort();
      }
    }

    // Write magic bytes.
    volatile uint8_t* ptr = (volatile uint8_t*)region;
    for (size_t off = 0; off < length; off += 4096) {
      ptr[off] = MAGIC_BYTE;
    }

    // Sleep 1ms-30s.
    usleep(1000 + (rand_r(&seed) % 30000000));

    // Verify magic bytes are still intact.
    for (size_t off = 0; off < length; off += 4096) {
      if (ptr[off] != MAGIC_BYTE) {
        fprintf(stderr,
                "CORRUPTION thread %d round %d offset %zu: "
                "expected 0x%02x got 0x%02x\n",
                id, round, off, MAGIC_BYTE, ptr[off]);
        return nullptr;
      }
    }
  }

  return nullptr;
}

static constexpr int NUM_PRESSURE_THREADS = 32;
static constexpr size_t PRESSURE_CHUNK = 512 * 1024 * 1024UL;
static constexpr int MAX_PRESSURE_CHUNKS = 10;

static void* pressure_thread(void*) {
  while (!atomic_load(&go)) {
    usleep(100);
  }

  void* chunks[MAX_PRESSURE_CHUNKS];
  for (int i = 0; i < MAX_PRESSURE_CHUNKS; i++) {
    chunks[i] = mmap(nullptr, PRESSURE_CHUNK, PROT_READ | PROT_WRITE,
                     MAP_PRIVATE | MAP_ANON, -1, 0);
    if (chunks[i] == MAP_FAILED) {
      fprintf(stderr, "pressure thread: mmap failed: %s\n", strerror(errno));
      abort();
    }
    memset(chunks[i], 0x55, PRESSURE_CHUNK);
  }

  for (int round = 0; round < 5000; round++) {
    for (int i = 0; i < MAX_PRESSURE_CHUNKS; i++) {
      volatile uint8_t* p = (volatile uint8_t*)chunks[i];
      for (size_t off = 0; off < PRESSURE_CHUNK; off += 4096) {
        p[off] = (uint8_t)round;
      }
    }
    if ((round % 20) == 0) {
      int slot = (round / 20) % MAX_PRESSURE_CHUNKS;
      munmap(chunks[slot], PRESSURE_CHUNK);
      chunks[slot] = mmap(nullptr, PRESSURE_CHUNK, PROT_READ | PROT_WRITE,
                          MAP_PRIVATE | MAP_ANON, -1, 0);
      if (chunks[slot] == MAP_FAILED) {
        fprintf(stderr, "pressure thread: mmap failed: %s\n", strerror(errno));
        abort();
      }
      memset(chunks[slot], 0x55, PRESSURE_CHUNK);
    }
  }

  for (int i = 0; i < MAX_PRESSURE_CHUNKS; i++) {
    munmap(chunks[i], PRESSURE_CHUNK);
  }
  return nullptr;
}

int main() {
  int total_threads = NUM_CHUNKS * THREADS_PER_CHUNK;
  printf("Starting %d worker threads across %d shared chunks (%zu MB each)\n"
         "  %d threads per chunk, %zu MB per thread, %d rounds\n"
         "  + %d pressure threads\n",
         total_threads, NUM_CHUNKS, CHUNK_SIZE / (1024 * 1024),
         THREADS_PER_CHUNK, SLOT_SIZE / (1024 * 1024), NUM_ROUNDS,
         NUM_PRESSURE_THREADS);

  pthread_t pressure[NUM_PRESSURE_THREADS];
  for (int i = 0; i < NUM_PRESSURE_THREADS; i++) {
    pthread_create(&pressure[i], nullptr, pressure_thread, nullptr);
  }

  // Allocate shared chunks. Each chunk is one mmap, with THREADS_PER_CHUNK
  // threads operating on non-overlapping slots within it.
  void* chunk_bases[NUM_CHUNKS];
  for (int i = 0; i < NUM_CHUNKS; i++) {
    chunk_bases[i] = mmap(nullptr, CHUNK_SIZE, PROT_READ | PROT_WRITE,
                          MAP_PRIVATE | MAP_ANON, -1, 0);
    if (chunk_bases[i] == MAP_FAILED) {
      fprintf(stderr, "chunk mmap failed at %d: %s\n", i, strerror(errno));
      return 1;
    }
  }

  pthread_t* threads = (pthread_t*)malloc(sizeof(pthread_t) * total_threads);
  struct worker_arg* args =
      (struct worker_arg*)malloc(sizeof(struct worker_arg) * total_threads);
  int created = 0;

  for (int c = 0; c < NUM_CHUNKS; c++) {
    for (int t = 0; t < THREADS_PER_CHUNK; t++) {
      int idx = c * THREADS_PER_CHUNK + t;
      args[idx].region = (uint8_t*)chunk_bases[c] + (size_t)t * SLOT_SIZE;
      args[idx].length = SLOT_SIZE;
      args[idx].id = idx;
      int err = pthread_create(&threads[idx], nullptr, worker, &args[idx]);
      if (err) {
        fprintf(stderr, "pthread_create failed at %d: %s\n", idx, strerror(err));
        abort();
      }
      created++;
    }
  }

  printf("Created %d worker threads, waiting for ready...\n", created);

  while (atomic_load(&started_threads) < created) {
    usleep(1000);
  }
  printf("All threads ready, starting test.\n");
  atomic_store(&go, 1);

  for (int i = 0; i < created; i++) {
    pthread_join(threads[i], nullptr);
  }
  for (int i = 0; i < NUM_PRESSURE_THREADS; i++) {
    pthread_join(pressure[i], nullptr);
  }

  for (int i = 0; i < NUM_CHUNKS; i++) {
    munmap(chunk_bases[i], CHUNK_SIZE);
  }
  free(threads);
  free(args);
  return 0;
}