Allow run-time configuration of core pinning

#include <stdio.h>

#include <stdlib.h>

#include <arpa/inet.h>

#include "common.h"

    printf("\n");

}

#include "cpulist.h"

int pin_thread_to_core(pthread_attr_t *ap, int core)

{

    char *cpulist_conf = getenv("PRRT_CORES");

    if (cpulist_conf) {

        uint32_t cpulist = cpulist_parse(cpulist_conf);

        core = cpulist_get_cpu(cpulist, core);

    }

    cpu_set_t cpuset;

    CPU_ZERO(&cpuset);

    CPU_SET(core, &cpuset);

    return pthread_attr_setaffinity_np(ap, sizeof(cpu_set_t), &cpuset);

}

#ifndef PRRT_CPULIST_H

#define PRRT_CPULIST_H

#include <stdint.h>

#include <stdlib.h>

#include <stdbool.h>

#include <limits.h>

#define CPULIST_ERR ((uint32_t) 0)

#ifdef __GNUC__

#  define cpulist__inline inline __attribute__((__always_inline__,__unused__))

#else

#  define cpulist__inline inline

#endif

/**

 * Add a cpu number to a cpu list

 *

 * \param cpulist the cpu list

 * \param num the cpu number

 * \returns a new cpu list that contains all cpus from the given cpu list, plus

 * the new number. If the old cpu list already contains the given cpu number,

 * it is returned unmodified.

 */

static cpulist__inline

uint32_t cpulist_add_cpu(uint32_t cpulist, unsigned int bit) {

	if (bit >= 32)

		return cpulist;

	uint32_t one = 1;

	return cpulist | (one << bit);

}

/**

 * Check if a cpu number is included in a cpu list

 *

 * \param cpulist the cpu list

 * \param num a cpu number

 * \returns \c true if the cpu list included the given cpu number, \c false otherwise

 */

static cpulist__inline

bool cpulist_has_cpu(uint32_t cpulist, unsigned int num) {

	if (num >= 32)

		return false;

	uint32_t one = 1;

	return !!(cpulist & (one << num));

}

/**

 *

 * The purpose of this function is to assign arbitrary cpu numbers to cpu

 * numbers allowed by the given cpu list. This is a very simple way of work

 * distribution.

 *

 * \param cpulist the cpu list

 * \param num an arbitrary cpu number that is mapped to a cpu number allowed by the given cpu list

 * \returns a cpu number that is included in the cpu list, if possible

 * \returns UINT_MAX if the cpu list is empty

 */

static cpulist__inline

unsigned int cpulist_get_cpu(uint32_t cpulist, unsigned int num) {

	if (!cpulist)

		return 0;

	// TODO: find a better load distribution algorithm

	unsigned int bit = UINT_MAX;

	bool has = 0;

	do {

		num -= has;

		bit = (bit + 1) & 31;

		has = cpulist_has_cpu(cpulist, bit);

	} while (num || !has);

	return bit;

}

/**

 * Remove all large cpu numbers from a cpu list

 *

 * This function removes all large cpu numbers from a cpu list. The given limit

 * is removed as well. Only cpus lower than num remain in the cpu list.

 *

 * The purpose of this function is to limit a cpu list to cpu numbers that are

 * actually available in the system.

 *

 * \param cpulist the cpu list

 * \param num     the cpu limit

 * \returns the new cpu list with large numbers removed

 */

static cpulist__inline

uint32_t cpulist_cut(uint32_t cpulist, unsigned int num) {

	uint32_t one = 1;

	uint32_t all = (one << num) - 1;

	return cpulist & all;

}

/**

 * Parse a cpu list

 *

 * This function takes a string that describes a cpu list in an intuitive

 * format and converts it to a machine-readable cpu list.

 *

 * Example: "1-3,5" ==> [ 1, 2, 3, 5 ]

 *

 * Syntax:

 *  cpulist ::= <range> | <range> ',' <range> # comma-separated list of ranges

 *  range   ::= <num>   | <num> '-' <num>     # single number or actual range

 *  num     ::= <digit> | <num> <digit>       # at least one digit

 *  digit   ::= '0' .. '9'

 */

static cpulist__inline

uint32_t cpulist_parse(const char *str)

{

	uint32_t result = 0;

	char *pos = (char *) str;

	while (*pos) {

		char *end;

		long from = strtol(pos, &end, 10);

		long to;

		switch(*end) {

		case '\0':

		case ',':

			to = from;

			break;

		case '-':

			pos = end + 1;

			to = strtol(pos, &end, 10);

			break;

		default:

			return CPULIST_ERR;

		}

		pos = end;

		for (long l = from; l <= to; l++)

			result = cpulist_add_cpu(result, (int) l);

		switch (*pos) {

		case '\0':

			return result;

		case ',':

			pos++;

			break;

		default:

			return CPULIST_ERR;

		}

	}

	return result;

}

#endif /* PRRT_CPULIST_H */