/* Three ways to get a thread out of a blocking recvfrom, tested rather than
   assumed: close the socket, shut it down, or give it a receive timeout.

   This decides how udp workers can be stopped. They sit in recvfrom and are
   never joined today, and the sockets of a SO_REUSEPORT group all share one
   address, so a worker cannot be singled out by sending it a datagram.

     cc -O2 -pthread -o /tmp/wake tests/udp_wake_probe.c && /tmp/wake

   On Linux 6.1: close does not wake it, shutdown does, SO_RCVTIMEO does. */
#include <errno.h>
#include <netinet/in.h>
#include <pthread.h>
#include <stdatomic.h>
#include <stdio.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>

static _Atomic int returned;
static _Atomic int saved_errno;
static int         sock;

static void       *reader(void *args) {
  char buffer[64];
  (void)args;
  if (recvfrom(sock, buffer, sizeof(buffer), 0, NULL, NULL) < 0)
    atomic_store(&saved_errno, errno);
  atomic_store(&returned, 1);
  return NULL;
}

static int bind_udp(int timeout_ms) {
  struct sockaddr_in address;
  int                s = socket(AF_INET, SOCK_DGRAM, 0);

  memset(&address, 0, sizeof(address));
  address.sin_family      = AF_INET;
  address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
  address.sin_port        = 0;
  bind(s, (struct sockaddr *)&address, sizeof(address));

  if (timeout_ms) {
    struct timeval how_long;
    how_long.tv_sec  = timeout_ms / 1000;
    how_long.tv_usec = (timeout_ms % 1000) * 1000;
    setsockopt(s, SOL_SOCKET, SO_RCVTIMEO, &how_long, sizeof(how_long));
  }
  return s;
}

static void try_it(const char *name, int timeout_ms, int do_close, int do_shutdown) {
  pthread_t thread;
  int       waited;

  atomic_store(&returned, 0);
  atomic_store(&saved_errno, 0);
  sock = bind_udp(timeout_ms);
  pthread_create(&thread, NULL, reader, NULL);
  usleep(300000); /* let it reach recvfrom */

  if (do_shutdown)
    shutdown(sock, SHUT_RD);
  if (do_close)
    close(sock);

  for (waited = 0; waited < 2000 && !atomic_load(&returned); ++waited)
    usleep(1000);

  if (atomic_load(&returned)) {
    printf("  %-24s returned after %4d ms, errno %s\n", name, waited,
           atomic_load(&saved_errno) ? strerror(atomic_load(&saved_errno)) : "none");
    pthread_join(thread, NULL);
  } else {
    printf("  %-24s STILL BLOCKED after 2000 ms\n", name);
    /* unblock it so the probe can finish */
    struct sockaddr_in to;
    socklen_t          length = sizeof(to);
    int                poke   = socket(AF_INET, SOCK_DGRAM, 0);
    if (!getsockname(sock, (struct sockaddr *)&to, &length))
      sendto(poke, "x", 1, 0, (struct sockaddr *)&to, length);
    close(poke);
    usleep(200000);
    if (atomic_load(&returned))
      pthread_join(thread, NULL);
    else
      pthread_detach(thread);
  }
  if (!do_close)
    close(sock);
}

int main(void) {
  try_it("close() the socket", 0, 1, 0);
  try_it("shutdown(SHUT_RD)", 0, 0, 1);
  try_it("SO_RCVTIMEO 500 ms", 500, 0, 0);
  return 0;
}