t_fault_sigaltstack.cc

// clang++ -Wall -Werror -std=c++11 t_fault_sigaltstack.cc -o t_fault_sigaltstack -lpthread

#include <err.h>
#include <errno.h>
#include <inttypes.h>
#include <libgen.h>
#include <pthread.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>
#include <time.h>
#include <unistd.h>

#include <limits>
#include <string>
#include <vector>

#if defined(__APPLE__)
#include <dispatch/dispatch.h>
#else
#include <semaphore.h>
#endif

#define HANDLE_EINTR(x)                                     \
  ({                                                        \
    decltype(x) eintr_wrapper_result;                       \
    do {                                                    \
      eintr_wrapper_result = (x);                           \
    } while (eintr_wrapper_result == -1 && errno == EINTR); \
    eintr_wrapper_result;                                   \
  })

#define IGNORE_EINTR(x)                                   \
  ({                                                      \
    decltype(x) eintr_wrapper_result;                     \
    do {                                                  \
      eintr_wrapper_result = (x);                         \
      if (eintr_wrapper_result == -1 && errno == EINTR) { \
        eintr_wrapper_result = 0;                         \
      }                                                   \
    } while (0);                                          \
    eintr_wrapper_result;                                 \
  })

namespace {

void HandleSignal(int sig, siginfo_t* info, void* ucontext) {}

class Semaphore {
 public:
#if defined(__APPLE__)
  Semaphore() : semaphore_(dispatch_semaphore_create(0)) {
    if (!semaphore_) {
      err(EXIT_FAILURE, "dispatch_semaphore_create: failed");
    }
  }

  ~Semaphore() { dispatch_release(semaphore_); }

  void Signal() { dispatch_semaphore_signal(semaphore_); }
  void Wait() { dispatch_semaphore_wait(semaphore_, DISPATCH_TIME_FOREVER); }
#else
  Semaphore() {
    if (sem_init(&semaphore_, 0, 0) != 0) {
      err(EXIT_FAILURE, "sem_init");
    }
  }

  ~Semaphore() {
    if (sem_destroy(&semaphore_) != 0) {
      err(EXIT_FAILURE, "sem_destroy");
    }
  }

  void Signal() {
    if (sem_post(&semaphore_) != 0) {
      err(EXIT_FAILURE, "sem_post");
    }
  }

  void Wait() {
    if (HANDLE_EINTR(sem_wait(&semaphore_)) != 0) {
      err(EXIT_FAILURE, "sem_wait");
    }
  }
#endif

 private:
#if defined(__APPLE__)
  dispatch_semaphore_t semaphore_;
#else
  sem_t semaphore_;
#endif

  Semaphore(const Semaphore&) = delete;
  void operator=(const Semaphore&) = delete;
};

struct Semaphores {
  Semaphore start;
  Semaphore stop;
};

void* ThreadMain(void* arg) {
  Semaphores* semaphores = static_cast<Semaphores*>(arg);

#if defined(__APPLE__)
  uint64_t tid;
  errno = pthread_threadid_np(pthread_self(), &tid);
  if (errno != 0) {
    err(EXIT_FAILURE, "pthread_threadid_np");
  }
#else
  uint64_t tid = gettid();
#endif

  char path[65];
  snprintf(path, sizeof(path), "/tmp/t_fault_sigaltstack.%d.%" PRIu64 ".XXXXXX",
           getpid(), tid);
  int fd = HANDLE_EINTR(mkstemp(path));
  if (fd < 0) {
    err(EXIT_FAILURE, "mkstemp");
  }

  if (unlink(path) != 0) {
    err(EXIT_FAILURE, "unlink");
  }

  stack_t ss;
  ss.ss_flags = 0;
  ss.ss_size = SIGSTKSZ;

  if (ftruncate(fd, ss.ss_size) != 0) {
    err(EXIT_FAILURE, "ftruncate");
  }

  ss.ss_sp =
      mmap(nullptr, ss.ss_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
  if (ss.ss_sp == MAP_FAILED) {
    err(EXIT_FAILURE, "mmap");
  }

  if (IGNORE_EINTR(close(fd)) != 0) {
    err(EXIT_FAILURE, "close");
  }

  if (sigaltstack(&ss, nullptr) != 0) {
    err(EXIT_FAILURE, "sigaltstack");
  }

  semaphores->start.Signal();

  semaphores->stop.Wait();

  ss.ss_flags = SS_DISABLE;
  if (sigaltstack(&ss, nullptr) != 0) {
    err(EXIT_FAILURE, "sigaltstack");
  }

  if (munmap(ss.ss_sp, ss.ss_size) != 0) {
    err(EXIT_FAILURE, "munmap");
  }

  return nullptr;
}

#if defined(CLOCK_UPTIME_RAW_APPROX)
constexpr clockid_t kClockId = CLOCK_UPTIME_RAW_APPROX;
#elif defined(CLOCK_MONOTONIC_COARSE)
constexpr clockid_t kClockId = CLOCK_MONOTONIC_COARSE;
#else
constexpr clockid_t kClockId = CLOCK_MONOTONIC;
#endif

bool TimerRunning(const timespec& end_time) {
  timespec now_time;
  if (clock_gettime(kClockId, &now_time) != 0) {
    err(EXIT_FAILURE, "clock_gettime");
  }

  return (now_time.tv_sec == end_time.tv_sec)
             ? (now_time.tv_nsec < end_time.tv_nsec)
             : (now_time.tv_sec < end_time.tv_sec);
}

std::string* g_me;
[[noreturn]] void Usage() {
  fprintf(stderr, "usage: %s [-n threads] [-t seconds]\n", g_me->c_str());
  exit(EXIT_FAILURE);
}

template <typename T>
T StringToUnsigned(const char* string) {
  char* end;
  errno = 0;
  unsigned long value = strtoul(string, &end, 0);
  if (errno != 0 || *end != '\0' || value > std::numeric_limits<T>::max()) {
    Usage();
  }
  return value;
}

}  // namespace

int main(int argc, char* argv[]) {
  g_me = new std::string(basename(argv[0]));

  size_t thread_count = 8;
  time_t timeout_seconds = 60;

  int option;
  while ((option = getopt(argc, argv, "n:t:")) != -1) {
    switch (option) {
      case 'n': {
        thread_count = StringToUnsigned<size_t>(optarg);
        if (thread_count < 1) {
          Usage();
        }
        break;
      }
      case 't': {
        timeout_seconds = StringToUnsigned<time_t>(optarg);
        break;
      }
      default: {
        Usage();
      }
    }
  }
  argc -= optind;
  argv += optind;
  if (argc > 0) {
    Usage();
  }

  struct sigaction sa;
  sa.sa_sigaction = HandleSignal;
  sigemptyset(&sa.sa_mask);
  sa.sa_flags = SA_SIGINFO | SA_ONSTACK;

  if (sigaction(SIGURG, &sa, nullptr) != 0) {
    err(EXIT_FAILURE, "sigaction");
  }

  Semaphores semaphores;

  std::vector<pthread_t> pthreads;
  for (size_t i = 0; i < thread_count; ++i) {
    pthread_t pthread;
    errno = pthread_create(&pthread, nullptr, ThreadMain, &semaphores);
    if (errno != 0) {
      err(EXIT_FAILURE, "pthread_create");
    }
    pthreads.push_back(pthread);

    semaphores.start.Wait();
  }

  timespec end_time;
  if (clock_gettime(kClockId, &end_time) != 0) {
    err(EXIT_FAILURE, "clock_gettime");
  }
  end_time.tv_sec += timeout_seconds;

  while (timeout_seconds == 0 || TimerRunning(end_time)) {
    for (const pthread_t pthread : pthreads) {
      errno = pthread_kill(pthread, SIGURG);
      if (errno != 0) {
        err(EXIT_FAILURE, "pthread_kill");
      }
    }
  }

  for (size_t i = 0; i < pthreads.size(); ++i) {
    semaphores.stop.Signal();
  }

  for (const pthread_t pthread : pthreads) {
    errno = pthread_join(pthread, nullptr);
    if (errno != 0) {
      err(EXIT_FAILURE, "pthread_join");
    }
  }

  return EXIT_SUCCESS;
}