Stopping linux threads example

server.cpp

// See <https://mazzo.li/posts/stopping-linux-threads.html>
// for blog post.
//
// Spawns a thread with a server listening on 55555 UDP, and
// then terminates it after 1 minute.
//
// I compile and run with
//
//     clang++ -Wall -std=c++20 server.cpp -lpthread -o server && ./server
//
// Then you can
//
//     echo "hello" | nc -u 127.0.0.1 55555
//
// to test.
#include <stdio.h>
#include <string.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>
#include <signal.h>
#include <stdlib.h>
#include <pthread.h>
#include <errno.h>
#include <atomic>
#include <vector>

static void die_error(const char* what, int err) {
  fprintf(stderr, "%s: %s\n", what, strerror(err));
  exit(1);
}

static void die_syscall(const char* what) {
  die_error(what, errno);
}

static void pthread_sigmask_or_die(int how, const sigset_t* set) {
  int ret = pthread_sigmask(SIG_SETMASK, set, nullptr);
  if (ret != 0) {
    die_error("pthread_sigmask", ret);
  }
}

static void* server(void*) {
  // normally USR1 is blocked
  sigset_t usr1_masked;
  sigemptyset(&usr1_masked);
  sigaddset(&usr1_masked, SIGUSR1);
  pthread_sigmask_or_die(SIG_SETMASK, &usr1_masked);

  // but while in syscalls, we want USR1 to reach it
  sigset_t usr1_unmasked;
  sigemptyset(&usr1_unmasked);

  // create socket
  int sock = socket(AF_INET, SOCK_DGRAM, 0);
  if (sock < 0) {
    die_syscall("socket");
  }

  // bind to 55555
  struct sockaddr_in servaddr = { 0 };
  servaddr.sin_family = AF_INET;
  servaddr.sin_addr.s_addr = INADDR_ANY;
  servaddr.sin_port = htons(55555);
  if (bind(sock, (struct sockaddr*)&servaddr, sizeof(servaddr)) < 0) {
    die_syscall("bind");
  }

  printf("listening on port 55555...\n");

  // process
  std::vector<char> buffer(9000); // jumbo frames anyone?
  for (;;) {
    printf("waiting for packet...\n");

    pthread_sigmask_or_die(SIG_SETMASK, &usr1_unmasked);
    ssize_t recvlen = recvfrom(sock, buffer.data(), buffer.size(), 0, nullptr, nullptr);
    pthread_sigmask_or_die(SIG_SETMASK, &usr1_masked);
    if (recvlen < 0 && errno == EINTR) {
      break;
    }

    if (recvlen < 0) {
      die_syscall("recvfrom");
    }

    printf("got packet.\n");
    size_t written = 0;
    while (written < recvlen) {
      ssize_t r = write(STDOUT_FILENO, &buffer[written], recvlen-written);
      if (r < 0) {
        die_syscall("write");
      }
      written += r;
    }
  }

  // close socket before exiting
  close(sock);

  return nullptr;
}

static void dummy_handler(int) {}

int main() {
  // install signal handler -- it just needs
  // to interrupt the syscall
  {
    struct sigaction act = {{ 0 }};
    act.sa_handler = &dummy_handler;
    if (sigaction(SIGUSR1, &act, nullptr) < 0) {
      die_syscall("sigaction");
    }
  }

  printf("starting server in different thread\n");
  pthread_t thr;
  {
    int res = pthread_create(&thr, nullptr, &server, nullptr);
    if (res != 0) {
      die_error("pthread_create", res);
    }
  }
  sleep(60);

  printf("stopping server\n");
  {
    int res = pthread_kill(thr, SIGUSR1);
    if (res != 0) {
      die_error("pthread_kill", res);
    }
    res = pthread_join(thr, nullptr);
    if (res != 0) {
      die_error("pthread_join", res);
    }
  }
  return 0;
}