node-sched/binding.cc at main · nxtedition/node-sched · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
#define NAPI_VERSION 8

#include <assert.h>
#include <napi-macros.h>
#include <node_api.h>

#include <errno.h>
#include <string.h>
#include <stdio.h>
#include <unistd.h>

#ifdef __linux__
  #include <sched.h>
#endif

static napi_value throw_notsup(napi_env env) {
  napi_throw_error(env, NULL, "CPU affinity is not supported on this platform");
  return NULL;
}

static napi_value throw_errno(napi_env env, const char *syscall) {
  char msg[256];
  snprintf(msg, sizeof(msg), "%s failed: %s (errno %d)", syscall, strerror(errno), errno);
  napi_throw_error(env, NULL, msg);
  return NULL;
}

NAPI_METHOD(sched_getaffinity) {
  NAPI_ARGV(1);

#ifndef __linux__
  return throw_notsup(env);
#else
  // pid
  int64_t pid_i64 = 0;
  if (napi_get_value_int64(env, argv[0], &pid_i64) != napi_ok) {
    napi_throw_type_error(env, NULL, "sched_getaffinity(pid): pid must be a number");
    return NULL;
  }
  pid_t pid = (pid_t)pid_i64;

  cpu_set_t set;
  CPU_ZERO(&set);

  if (sched_getaffinity(pid, sizeof(set), &set) != 0) {
    return throw_errno(env, "sched_getaffinity");
  }

  napi_value result;
  napi_create_array(env, &result);

  // Note: limited by CPU_SETSIZE (musl typically 128)
  uint32_t j = 0;
  for (int cpu = 0; cpu < (int)CPU_SETSIZE; cpu++) {
    if (CPU_ISSET(cpu, &set)) {
      napi_value v;
      napi_create_int32(env, cpu, &v);
      napi_set_element(env, result, j++, v);
    }
  }
  return result;
#endif
}

NAPI_METHOD(sched_setaffinity) {
  NAPI_ARGV(2);

#ifndef __linux__
  return throw_notsup(env);
#else
  // pid
  int64_t pid_i64 = 0;
  if (napi_get_value_int64(env, argv[0], &pid_i64) != napi_ok) {
    napi_throw_type_error(env, NULL, "sched_setaffinity(pid, cpus): pid must be a number");
    return NULL;
  }
  pid_t pid = (pid_t)pid_i64;

  // cpus array
  bool is_array = false;
  if (napi_is_array(env, argv[1], &is_array) != napi_ok || !is_array) {
    napi_throw_type_error(env, NULL, "sched_setaffinity(pid, cpus): cpus must be an array");
    return NULL;
  }

  uint32_t len = 0;
  napi_get_array_length(env, argv[1], &len);

  cpu_set_t set;
  CPU_ZERO(&set);

  for (uint32_t i = 0; i < len; i++) {
    napi_value elem;
    if (napi_get_element(env, argv[1], i, &elem) != napi_ok) {
      napi_throw_type_error(env, NULL, "Failed reading cpus array element");
      return NULL;
    }
    int32_t cpu = -1;
    if (napi_get_value_int32(env, elem, &cpu) != napi_ok) {
      napi_throw_type_error(env, NULL, "CPU index must be an integer");
      return NULL;
    }
    if (cpu < 0 || cpu >= (int)CPU_SETSIZE) {
      napi_throw_range_error(env, NULL, "CPU index out of range (limited by CPU_SETSIZE)");
      return NULL;
    }
    CPU_SET(cpu, &set);
  }

  if (sched_setaffinity(pid, sizeof(set), &set) != 0) {
    return throw_errno(env, "sched_setaffinity");
  }

  napi_value undef;
  napi_get_undefined(env, &undef);
  return undef;
#endif
}

NAPI_INIT() {
  NAPI_EXPORT_FUNCTION(sched_getaffinity);
  NAPI_EXPORT_FUNCTION(sched_setaffinity);
}