<commit_before>#ifdef __linux__
#define _GNU_SOURCE
#endif
#include <platform.h>
#include <dtrace.h>

#if defined(PLATFORM_IS_LINUX) || defined(PLATFORM_IS_BSD)
  #include <sched.h>
  #include <stdlib.h>
  #include <unistd.h>
  #include <stdio.h>
#elif defined(PLATFORM_IS_MACOSX)
  #include <unistd.h>
  #include <mach/mach.h>
  #include <mach/thread_policy.h>
#elif defined(PLATFORM_IS_WINDOWS)
  #include <processtopologyapi.h>
#endif

#include "cpu.h"
#include "../mem/pool.h"

#if defined(PLATFORM_IS_LINUX)
static uint32_t avail_cpu_count;
static uint32_t* avail_cpu_list;
static uint32_t avail_cpu_size;
#endif

#if defined(PLATFORM_IS_MACOSX) || defined(PLATFORM_IS_BSD)

#include <sys/types.h>
#include <sys/sysctl.h>

static uint32_t property(const char* key)
{
  int value;
  size_t len = sizeof(int);
  sysctlbyname(key, &value, &len, NULL, 0);
  return value;
}
#endif

static uint32_t hw_cpu_count;

#if defined(PLATFORM_IS_LINUX)
static bool cpu_physical(uint32_t cpu)
{
  char file[FILENAME_MAX];
  snprintf(file, FILENAME_MAX,
    "/sys/devices/system/cpu/cpu%d/topology/thread_siblings_list", cpu);

  FILE* fp = fopen(file, "r");

  if(fp != NULL)
  {
    char name[16];
    size_t len = fread(name, 1, 15, fp);
    name[len] = '\0';
    fclose(fp);

    if(cpu != (uint32_t)atoi(name))
      return false;
  }

  return true;
}

static uint32_t cpu_add_mask_to_list(uint32_t i, cpu_set_t* mask)
{
  uint32_t count = CPU_COUNT(mask);
  uint32_t index = 0;
  uint32_t found = 0;

  while(found < count)
  {
    if(CPU_ISSET(index, mask))
    {
      avail_cpu_list[i++] = index;
      found++;
    }

    index++;
  }

  return i;
}
#endif

void ponyint_cpu_init()
{
#if defined(PLATFORM_IS_LINUX)
  cpu_set_t all_cpus;
  cpu_set_t hw_cpus;
  cpu_set_t ht_cpus;

  sched_getaffinity(0, sizeof(cpu_set_t), &all_cpus);
  CPU_ZERO(&hw_cpus);
  CPU_ZERO(&ht_cpus);

  avail_cpu_count = CPU_COUNT(&all_cpus);
  uint32_t index = 0;
  uint32_t found = 0;

  while(found < avail_cpu_count)
  {
    if(CPU_ISSET(index, &all_cpus))
    {
      if(cpu_physical(index))
        CPU_SET(index, &hw_cpus);
      else
        CPU_SET(index, &ht_cpus);

      found++;
    }

    index++;
  }

  hw_cpu_count = CPU_COUNT(&hw_cpus);

  if(ponyint_numa_init())
  {
    uint32_t numa_count = ponyint_numa_cores();

    if(avail_cpu_count > numa_count)
      avail_cpu_size = numa_count;

    avail_cpu_list =
      (uint32_t*)ponyint_pool_alloc_size(avail_cpu_size * sizeof(uint32_t));
    avail_cpu_count =
      ponyint_numa_core_list(&hw_cpus, &ht_cpus, avail_cpu_list);
  } else {
    avail_cpu_size = avail_cpu_count;
    avail_cpu_list =
      (uint32_t*)ponyint_pool_alloc_size(avail_cpu_size * sizeof(uint32_t));

    uint32_t i = 0;
    i = cpu_add_mask_to_list(i, &hw_cpus);
    i = cpu_add_mask_to_list(i, &ht_cpus);
  }
#elif defined(PLATFORM_IS_BSD)
  hw_cpu_count = property("hw.ncpu");
#elif defined(PLATFORM_IS_MACOSX)
  hw_cpu_count = property("hw.physicalcpu");
#elif defined(PLATFORM_IS_WINDOWS)
  PSYSTEM_LOGICAL_PROCESSOR_INFORMATION info = NULL;
  PSYSTEM_LOGICAL_PROCESSOR_INFORMATION ptr = NULL;
  DWORD len = 0;<commit_msg>Fix memory overflow when allocating CPUs when numa is enabled<commit_after>#ifdef __linux__
#define _GNU_SOURCE
#endif
#include <platform.h>
#include <dtrace.h>

#if defined(PLATFORM_IS_LINUX) || defined(PLATFORM_IS_BSD)
  #include <sched.h>
  #include <stdlib.h>
  #include <unistd.h>
  #include <stdio.h>
#elif defined(PLATFORM_IS_MACOSX)
  #include <unistd.h>
  #include <mach/mach.h>
  #include <mach/thread_policy.h>
#elif defined(PLATFORM_IS_WINDOWS)
  #include <processtopologyapi.h>
#endif

#include "cpu.h"
#include "../mem/pool.h"

#if defined(PLATFORM_IS_LINUX)
static uint32_t avail_cpu_count;
static uint32_t* avail_cpu_list;
static uint32_t avail_cpu_size;
#endif

#if defined(PLATFORM_IS_MACOSX) || defined(PLATFORM_IS_BSD)

#include <sys/types.h>
#include <sys/sysctl.h>

static uint32_t property(const char* key)
{
  int value;
  size_t len = sizeof(int);
  sysctlbyname(key, &value, &len, NULL, 0);
  return value;
}
#endif

static uint32_t hw_cpu_count;

#if defined(PLATFORM_IS_LINUX)
static bool cpu_physical(uint32_t cpu)
{
  char file[FILENAME_MAX];
  snprintf(file, FILENAME_MAX,
    "/sys/devices/system/cpu/cpu%d/topology/thread_siblings_list", cpu);

  FILE* fp = fopen(file, "r");

  if(fp != NULL)
  {
    char name[16];
    size_t len = fread(name, 1, 15, fp);
    name[len] = '\0';
    fclose(fp);

    if(cpu != (uint32_t)atoi(name))
      return false;
  }

  return true;
}

static uint32_t cpu_add_mask_to_list(uint32_t i, cpu_set_t* mask)
{
  uint32_t count = CPU_COUNT(mask);
  uint32_t index = 0;
  uint32_t found = 0;

  while(found < count)
  {
    if(CPU_ISSET(index, mask))
    {
      avail_cpu_list[i++] = index;
      found++;
    }

    index++;
  }

  return i;
}
#endif

void ponyint_cpu_init()
{
#if defined(PLATFORM_IS_LINUX)
  cpu_set_t all_cpus;
  cpu_set_t hw_cpus;
  cpu_set_t ht_cpus;

  sched_getaffinity(0, sizeof(cpu_set_t), &all_cpus);
  CPU_ZERO(&hw_cpus);
  CPU_ZERO(&ht_cpus);

  avail_cpu_size = avail_cpu_count = CPU_COUNT(&all_cpus);
  uint32_t index = 0;
  uint32_t found = 0;

  while(found < avail_cpu_count)
  {
    if(CPU_ISSET(index, &all_cpus))
    {
      if(cpu_physical(index))
        CPU_SET(index, &hw_cpus);
      else
        CPU_SET(index, &ht_cpus);

      found++;
    }

    index++;
  }

  hw_cpu_count = CPU_COUNT(&hw_cpus);

  if(ponyint_numa_init())
  {
    uint32_t numa_count = ponyint_numa_cores();

    if(avail_cpu_count > numa_count)
      avail_cpu_size = numa_count;

    avail_cpu_list =
      (uint32_t*)ponyint_pool_alloc_size(avail_cpu_size * sizeof(uint32_t));
    avail_cpu_count =
      ponyint_numa_core_list(&hw_cpus, &ht_cpus, avail_cpu_list);
  } else {
    avail_cpu_list =
      (uint32_t*)ponyint_pool_alloc_size(avail_cpu_size * sizeof(uint32_t));

    uint32_t i = 0;
    i = cpu_add_mask_to_list(i, &hw_cpus);
    i = cpu_add_mask_to_list(i, &ht_cpus);
  }
#elif defined(PLATFORM_IS_BSD)
  hw_cpu_count = property("hw.ncpu");
#elif defined(PLATFORM_IS_MACOSX)
  hw_cpu_count = property("hw.physicalcpu");
#elif defined(PLATFORM_IS_WINDOWS)
  PSYSTEM_LOGICAL_PROCESSOR_INFORMATION info = NULL;
  PSYSTEM_LOGICAL_PROCESSOR_INFORMATION ptr = NULL;
  DWORD len = 0;
{'fraction_non_alphanumeric': 0.12418921056794811, 'fraction_numerical': 0.0170859041290935, 'mean_word_length': 4.229114971050455, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4521744', 'n_tokens_mistral': 2802, 'n_tokens_neox': 2666, 'n_words': 455}