#include "lib.h"

#include "cpu-count.h"

#ifdef HAVE_SCHED_H

#  define __USE_GNU

/* _GNU_SOURCE: for musl's sched.h to expose cpuset/CPU_* */

#  define _GNU_SOURCE

#  include <sched.h>

#  ifdef HAVE_SYS_CPUSET_H

#    include <sys/cpuset.h>

#  endif

#endif

int cpu_count_get(int *cpu_count_r, const char **error_r)

{

  int result;

#if defined(HAVE_SCHED_GETAFFINITY)

  cpu_set_t cs;

  CPU_ZERO(&cs);

  if (sched_getaffinity(0, sizeof(cs), &cs) < 0) {

    *error_r = t_strdup_printf("sched_getaffinity() failed: %m");

    return -1;

  }

  result = CPU_COUNT(&cs);

#elif defined(HAVE_CPUSET_GETAFFINITY)

  cpuset_t cs;

  CPU_CLR(sizeof(cs), &cs);

  if (cpuset_getaffinity(CPU_LEVEL_WHICH, CPU_WHICH_PID, -1,

             sizeof(cs), &cs) < 0) {

    *error_r = t_strdup_printf("cpuset_getaffinity() failed: %m");

    return -1;

  }

  result = CPU_COUNT(&cs);

#else

  *cpu_count_r = 0;

  *error_r = "Cannot get CPU count";

  return -1;

#endif

  *cpu_count_r = result;

  return 0;

}