incoming_cpu.c

#define _GNU_SOURCE
#include <assert.h>
#include <limits.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <pthread.h>
#include <stdio.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

static int open_listener(unsigned cpuid)
{
    int sockfd, ret;
    struct sockaddr_in sin;

    sockfd = socket(PF_INET, SOCK_STREAM, 0);
    assert(sockfd >= 0);

    unsigned on = 1;
    ret = setsockopt(sockfd, SOL_SOCKET, SO_REUSEPORT, &on, sizeof(on));
    assert(ret == 0);

    on = 1;
    ret = setsockopt(sockfd, SOL_TCP, TCP_NODELAY, &on, sizeof(on));
    assert(ret == 0);

    ret = setsockopt(sockfd, SOL_SOCKET, SO_INCOMING_CPU, &cpuid, sizeof(cpuid));
    assert(ret == 0);

    memset(&sin, 0, sizeof(sin));
    sin.sin_addr.s_addr = htonl(0);
    sin.sin_port = htons(12345);
    ret = bind(sockfd, (struct sockaddr *)&sin, sizeof(sin));
    assert(ret == 0);

    ret = listen(sockfd, SOMAXCONN);
    assert(ret == 0);

    return sockfd;
}

static void *thread_main(void *_cpuid)
{
    unsigned cpuid = (unsigned)_cpuid;
    int listen_fd, conn_fd, ret;

    // set CPU affinity to the cpuid being specified
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpuid, &cpuset);
    ret = pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
    assert(ret == 0);

    // open a listener with SO_REUSEPORT and SO_INCOMING_CPU set
    listen_fd = open_listener(cpuid);
    assert(listen_fd >= 0);

    // accept a connection, report desired CPUID ind and actual CPUID
    while (1) {
        while ((conn_fd = accept(listen_fd, NULL, NULL)) < 0)
            ;
        unsigned actual_cpuid = UINT_MAX;
        socklen_t actual_cpuidlen = sizeof(actual_cpuid);
        ret = getsockopt(conn_fd, SOL_SOCKET, SO_INCOMING_CPU, &actual_cpuid, &actual_cpuidlen);
        assert(ret == 0);
        assert(actual_cpuidlen == sizeof(actual_cpuid));
        printf("accepted socket; desired cpuid=%u, actual cpuid=%u\n", cpuid, actual_cpuid);
        close(conn_fd);
    }

    return NULL;
}

int main(int argc, char **argv)
{
    for (unsigned cpuid = 1; cpuid < 4; ++cpuid) {
        pthread_t tid;
        pthread_create(&tid, NULL, thread_main, (void *)cpuid);
    }

    thread_main((void *)0);

    return 0;
}