[LibOS,Pal] Add sched_setaffinity/sched_getaffinity syscall support

This patch adds syscall support for setting/getting cpu affinity
of threads.

Co-authored-by: Gary <gordon.king@intel.com>

Documentation/devel/new-syscall.rst

@@ -10,15 +10,15 @@ For example, assume we are implementing :manpage:`sched_setaffinity(2)`. You
 must find the definition of ``sched_setaffinity`` in
 :file:`shim_syscalls.c`, which will be the following code::
 
-   SHIM_SYSCALL_RETURN_ENOSYS(sched_setaffinity, 3, int, pid_t, pid, size_t,
-                              len, __kernel_cpu_set_t*, user_mask_ptr)
+   SHIM_SYSCALL_RETURN_ENOSYS(sched_setaffinity, 3, long, pid_t, pid, unsigned int,
+                              len, unsigned long*, user_mask_ptr)
 
 Change this line to ``DEFINE_SHIM_SYSCALL(...)`` to name the function that
 implements this system call: ``shim_do_sched_setaffinity`` (this is the naming
 convention, please follow it)::
 
-   DEFINE_SHIM_SYSCALL(sched_setaffinity, 3, shim_do_sched_setaffinity, int, pid_t, pid, size_t, len,
-                       __kernel_cpu_set_t*, user_mask_ptr)
+   DEFINE_SHIM_SYSCALL(sched_setaffinity, 3, shim_do_sched_setaffinity, long, pid_t, pid,
+                       unsigned int, len, unsigned long*, user_mask_ptr)
 
 
 2. Add definitions to system call table
@@ -30,7 +30,7 @@ in :file:`shim_table.h`: ``__shim_sched_setaffinity`` and
 second in respect to the system call you are implementing, with the same
 prototype as defined in :file:`shim_syscalls.c`::
 
-   int shim_do_sched_setaffinity(pid_t pid, size_t len, __kernel_cpu_set_t* user_mask_ptr);
+   long shim_do_sched_setaffinity(pid_t pid, unsigned int len, unsigned long* user_mask_ptr);
 
 3. Implement system call
 ------------------------
@@ -41,7 +41,7 @@ earlier) in a new source file or any existing source file in
 
 For example, in :file:`LibOS/shim/src/sys/shim_sched.c`::
 
-   int shim_do_sched_setaffinity(pid_t pid, size_t len, __kernel_cpu_set_t* user_mask_ptr) {
+   long shim_do_sched_setaffinity(pid_t pid, unsigned int len, unsigned long* user_mask_ptr) {
       /* code for implementing the semantics of sched_setaffinity */
    }
 

LibOS/shim/include/shim_table.h

@@ -489,8 +489,8 @@ pid_t shim_do_gettid(void);
 int shim_do_tkill(int pid, int sig);
 time_t shim_do_time(time_t* tloc);
 int shim_do_futex(int* uaddr, int op, int val, void* utime, int* uaddr2, int val3);
-int shim_do_sched_setaffinity(pid_t pid, size_t len, __kernel_cpu_set_t* user_mask_ptr);
-int shim_do_sched_getaffinity(pid_t pid, size_t len, __kernel_cpu_set_t* user_mask_ptr);
+long shim_do_sched_setaffinity(pid_t pid, unsigned int cpumask_size, unsigned long* user_mask_ptr);
+long shim_do_sched_getaffinity(pid_t pid, unsigned int cpumask_size, unsigned long* user_mask_ptr);
 int shim_do_set_tid_address(int* tidptr);
 int shim_do_semtimedop(int semid, struct sembuf* sops, unsigned int nsops,
                        const struct timespec* timeout);

LibOS/shim/include/shim_types.h

@@ -289,19 +289,6 @@ struct iovec {
     size_t iov_len; /* Length of data. */
 };
 
-/* bits/sched.h */
-/* Type for array elements in 'cpu_set_t'.  */
-typedef unsigned long int __kernel_cpu_mask;
-
-/* Size definition for CPU sets.  */
-#define __CPU_SETSIZE 1024
-#define __NCPUBITS    (8 * sizeof(__kernel_cpu_mask))
-
-/* Data structure to describe CPU mask.  */
-typedef struct {
-    __kernel_cpu_mask __bits[__CPU_SETSIZE / __NCPUBITS];
-} __kernel_cpu_set_t;
-
 struct getcpu_cache {
     unsigned long blob[128 / sizeof(long)];
 };

LibOS/shim/src/shim_syscalls.c

@@ -601,11 +601,11 @@ DEFINE_SHIM_SYSCALL(time, 1, shim_do_time, time_t, time_t*, tloc)
 DEFINE_SHIM_SYSCALL(futex, 6, shim_do_futex, int, int*, uaddr, int, op, int, val, void*, utime,
                     int*, uaddr2, int, val3)
 
-DEFINE_SHIM_SYSCALL(sched_setaffinity, 3, shim_do_sched_setaffinity, int, pid_t, pid, size_t, len,
-                    __kernel_cpu_set_t*, user_mask_ptr)
+DEFINE_SHIM_SYSCALL(sched_setaffinity, 3, shim_do_sched_setaffinity, long, pid_t, pid,
+                    unsigned int, len, unsigned long*, user_mask_ptr)
 
-DEFINE_SHIM_SYSCALL(sched_getaffinity, 3, shim_do_sched_getaffinity, int, pid_t, pid, size_t, len,
-                    __kernel_cpu_set_t*, user_mask_ptr)
+DEFINE_SHIM_SYSCALL(sched_getaffinity, 3, shim_do_sched_getaffinity, int, pid_t, pid,
+                    unsigned int, len, unsigned long*, user_mask_ptr)
 
 #if defined(__i386__) || defined(__x86_64__)
 SHIM_SYSCALL_RETURN_ENOSYS(set_thread_area, 1, int, struct user_desc*, u_info)

LibOS/shim/src/sys/shim_sched.c

@@ -16,6 +16,7 @@
 #include "pal.h"
 #include "shim_internal.h"
 #include "shim_table.h"
+#include "shim_thread.h"
 
 int shim_do_sched_yield(void) {
     DkThreadYieldExecution();
@@ -139,51 +140,85 @@ int shim_do_sched_rr_get_interval(pid_t pid, struct timespec* interval) {
     return 0;
 }
 
-static int check_affinity_params(int ncpus, size_t len, __kernel_cpu_set_t* user_mask_ptr) {
-    /* Check that user_mask_ptr is valid; if not, should return -EFAULT */
-    if (test_user_memory(user_mask_ptr, len, true))
+long shim_do_sched_setaffinity(pid_t pid, unsigned int cpumask_size, unsigned long* user_mask_ptr) {
+    int ret;
+
+    /* check if user_mask_ptr is valid */
+    if (test_user_memory(user_mask_ptr, cpumask_size, /*write=*/false))
         return -EFAULT;
 
-    /* Linux kernel bitmap is based on long. So according to its
-     * implementation, round up the result to sizeof(long) */
-    size_t bitmask_long_count    = (ncpus + sizeof(long) * 8 - 1) / (sizeof(long) * 8);
-    size_t bitmask_size_in_bytes = bitmask_long_count * sizeof(long);
-    if (len < bitmask_size_in_bytes)
-        return -EINVAL;
-    /* Linux kernel also rejects non-natural size */
-    if (len & (sizeof(long) - 1))
-        return -EINVAL;
+    struct shim_thread* thread = pid ? lookup_thread(pid) : get_cur_thread();
+    if (!thread)
+        return -ESRCH;
 
-    return bitmask_size_in_bytes;
-}
+    /* lookup_thread() internally increments thread count; do the same in case of
+       get_cur_thread(). */
+    if (pid == 0)
+        get_thread(thread);
 
-/* dummy implementation: ignore user-supplied mask and return success */
-int shim_do_sched_setaffinity(pid_t pid, size_t len, __kernel_cpu_set_t* user_mask_ptr) {
-    __UNUSED(pid);
-    int ncpus = PAL_CB(cpu_info.online_logical_cores);
+    /* Internal graphene threads are not affinitized; if we hit an internal thread here, this is
+       some bug in user app. */
+    if (is_internal(thread)) {
+        put_thread(thread);
+        return -ESRCH;
+    }
 
-    int bitmask_size_in_bytes = check_affinity_params(ncpus, len, user_mask_ptr);
-    if (bitmask_size_in_bytes < 0)
-        return bitmask_size_in_bytes;
+    ret = DkThreadSetCpuAffinity(thread->pal_handle, cpumask_size, user_mask_ptr);
+    if (!ret) {
+        put_thread(thread);
+        return -PAL_ERRNO();
+    }
 
+    put_thread(thread);
     return 0;
 }
 
-/* dummy implementation: always return all-ones (as many as there are host CPUs)  */
-int shim_do_sched_getaffinity(pid_t pid, size_t len, __kernel_cpu_set_t* user_mask_ptr) {
-    __UNUSED(pid);
-    int ncpus = PAL_CB(cpu_info.online_logical_cores);
+long shim_do_sched_getaffinity(pid_t pid, unsigned int cpumask_size, unsigned long* user_mask_ptr) {
+    int ret;
+    size_t cpu_cnt = PAL_CB(cpu_info.online_logical_cores);
 
-    int bitmask_size_in_bytes = check_affinity_params(ncpus, len, user_mask_ptr);
-    if (bitmask_size_in_bytes < 0)
-        return bitmask_size_in_bytes;
+    /* Check if user_mask_ptr is valid */
+    if (test_user_memory(user_mask_ptr, cpumask_size, /*write=*/true))
+        return -EFAULT;
 
-    memset(user_mask_ptr, 0, len);
-    for (int i = 0; i < ncpus; i++) {
-        ((uint8_t*)user_mask_ptr)[i / 8] |= 1 << (i % 8);
+    /* Linux kernel bitmap is based on long. So according to its implementation, round up the result
+     * to sizeof(long) */
+    size_t bitmask_size_in_bytes = BITS_TO_LONGS(cpu_cnt) * sizeof(long);
+    if (cpumask_size < bitmask_size_in_bytes) {
+        debug("size of cpumask must be at least %lu but supplied cpumask is %u\n",
+               bitmask_size_in_bytes, cpumask_size);
+        return -EINVAL;
     }
-    /* imitate the Linux kernel implementation
-     * See SYSCALL_DEFINE3(sched_getaffinity) */
+
+    /* Linux kernel also rejects non-natural size */
+    if (cpumask_size & (sizeof(long) - 1))
+        return -EINVAL;
+
+    struct shim_thread* thread = pid ? lookup_thread(pid) : get_cur_thread();
+    if (!thread)
+        return -ESRCH;
+
+    /* lookup_thread() internally increments thread count; do the same in case of
+       get_cur_thread(). */
+    if (pid == 0)
+        get_thread(thread);
+
+    /* Internal graphene threads are not affinitized; if we hit an internal thread here, this is
+       some bug in user app. */
+    if (is_internal(thread)) {
+        put_thread(thread);
+        return -ESRCH;
+    }
+
+    memset(user_mask_ptr, 0, cpumask_size);
+    ret = DkThreadGetCpuAffinity(thread->pal_handle, bitmask_size_in_bytes, user_mask_ptr);
+    if (!ret) {
+        put_thread(thread);
+        return -PAL_ERRNO();
+    }
+
+    put_thread(thread);
+    /* on success, imitate Linux kernel implementation: see SYSCALL_DEFINE3(sched_getaffinity) */
     return bitmask_size_in_bytes;
 }
 

LibOS/shim/test/regression/Makefile

@@ -62,8 +62,10 @@ c_executables = \
 	proc_cpuinfo \
 	proc_path \
 	pselect \
+	pthread_set_get_affinity \
 	readdir \
 	sched \
+	sched_set_get_affinity \
 	select \
 	shared_object \
 	sigaction_per_process \
@@ -161,6 +163,7 @@ CFLAGS-proc_common = -pthread
 CFLAGS-spinlock += -I$(PALDIR)/../include/lib -I$(PALDIR)/../include/arch/$(ARCH) -pthread
 CFLAGS-sigaction_per_process += -pthread
 CFLAGS-signal_multithread += -pthread
+CFLAGS-pthread_set_get_affinity += -pthread
 
 CFLAGS-attestation += -I$(PALDIR)/../lib/crypto/mbedtls/crypto/include \
                       -I$(PALDIR)/host/Linux-SGX \

LibOS/shim/test/regression/manifest.template

@@ -33,6 +33,8 @@ sgx.trusted_children.victim = file:exec_victim.sig
 
 sgx.allowed_files.tmp_dir = file:tmp/
 
-sgx.thread_num = 6
+# Set this value to at least 4 (1 for the main thread, 2 for Graphene internal threads and 1 for
+# helper threads some tests use)
+sgx.thread_num = 8
 
 sgx.static_address = 1

LibOS/shim/test/regression/pthread_set_get_affinity.c

@@ -0,0 +1,148 @@
+/* SPDX-License-Identifier: LGPL-3.0-or-later */
+/* Copyright (C) 2020 Intel Corporation */
+
+/*
+ * Test to set/get cpu affinity by parent process on behalf of its child threads.
+ */
+
+#define _GNU_SOURCE
+#include <err.h>
+#include <errno.h>
+#include <pthread.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/syscall.h>
+#include <unistd.h>
+
+#define min(a, b)               (((a) < (b)) ? (a) : (b))
+#define MAIN_THREAD_CNT         1
+#define INTERNAL_THREAD_CNT     2
+#define MANIFEST_SGX_THREAD_CNT 8 /* corresponds to sgx.thread_num in the manifest template */
+
+/* barrier to synchronize between parent and children */
+pthread_barrier_t barrier;
+
+/* Run a busy loop for some iterations, so that we can verify affinity with htop manually */
+static void* dowork(void* args) {
+    uint64_t* iterations = (uint64_t*)args;
+    __asm__ volatile (
+                      "movq %0, %%rax\n"
+                      "loop:\n"
+                      "dec %%rax\n"
+                      "cmp $0, %%rax\n"
+                      "jne loop\n"
+                      : /*no outs*/ : "m"(*iterations)  : "rax", "cc");
+
+    int ret = pthread_barrier_wait(&barrier);
+    if (ret != 0 && ret != PTHREAD_BARRIER_SERIAL_THREAD) {
+        errx(EXIT_FAILURE, "Child did not wait on barrier!");
+    }
+    return NULL;
+}
+
+int main(int argc, const char** argv) {
+    int ret;
+    long numprocs = sysconf(_SC_NPROCESSORS_ONLN);
+    if (numprocs < 0) {
+        err(EXIT_FAILURE, "Failed to retrieve the number of logical processors!");
+    }
+
+    /* If you want to run on all cores then increase sgx.thread_num in the manifest.template and
+     * also set MANIFEST_SGX_THREAD_CNT to the same value.
+     */
+    numprocs = min(numprocs, (MANIFEST_SGX_THREAD_CNT - (INTERNAL_THREAD_CNT + MAIN_THREAD_CNT)));
+
+    /* Affinitize threads to alternate logical processors to do a quick check from htop manually */
+    numprocs = (numprocs >= 2) ? numprocs/2 : 1;
+
+    pthread_t* threads = (pthread_t*)malloc(numprocs * sizeof(pthread_t));
+    if (!threads) {
+         errx(EXIT_FAILURE, "memory allocation failed");
+    }
+
+    if (pthread_barrier_init(&barrier, NULL, numprocs + 1)) {
+        free(threads);
+        errx(EXIT_FAILURE, "pthread barrier init failed");
+    }
+
+    cpu_set_t cpus, get_cpus;
+    uint64_t iterations = argc > 1 ? atol(argv[1]) : 10000000000;
+
+    /* Validate parent set/get affinity for child */
+    for (long i = 0; i < numprocs; i++) {
+        CPU_ZERO(&cpus);
+        CPU_ZERO(&get_cpus);
+        CPU_SET(i*2, &cpus);
+
+        ret = pthread_create(&threads[i], NULL, dowork, (void*)&iterations);
+        if (ret != 0) {
+            free(threads);
+            errx(EXIT_FAILURE, "pthread_create failed!");
+        }
+
+        ret = pthread_setaffinity_np(threads[i], sizeof(cpus), &cpus);
+        if (ret != 0) {
+            free(threads);
+            errx(EXIT_FAILURE, "pthread_setaffinity_np failed for child!");
+        }
+
+        ret = pthread_getaffinity_np(threads[i], sizeof(get_cpus), &get_cpus);
+        if (ret != 0) {
+            free(threads);
+            errx(EXIT_FAILURE, "pthread_getaffinity_np failed for child!");
+        }
+
+        if (!CPU_EQUAL_S(sizeof(cpus), &cpus, &get_cpus)) {
+            free(threads);
+            errx(EXIT_FAILURE, "get cpuset is not equal to set cpuset on proc: %ld", i);
+        }
+    }
+
+    /* unblock the child threads */
+    ret = pthread_barrier_wait(&barrier);
+    if (ret != 0 && ret != PTHREAD_BARRIER_SERIAL_THREAD) {
+        free(threads);
+        errx(EXIT_FAILURE, "Parent did not wait on barrier!");
+    }
+
+    for (int i = 0; i < numprocs; i++) {
+        ret = pthread_join(threads[i], NULL);
+        if (ret != 0) {
+            free(threads);
+            errx(EXIT_FAILURE, "pthread_join failed!");
+        }
+    }
+
+    /* Validating parent set/get affinity for children done. Free resources */
+    pthread_barrier_destroy(&barrier);
+    free(threads);
+
+    /* Validate parent set/get affinity for itself */
+    CPU_ZERO(&cpus);
+    CPU_SET(0, &cpus);
+    ret = pthread_setaffinity_np(pthread_self(), sizeof(cpus), &cpus);
+    if (ret != 0) {
+        errx(EXIT_FAILURE, "pthread_setaffinity_np failed for parent!");
+    }
+
+    CPU_ZERO(&get_cpus);
+    ret = pthread_getaffinity_np(pthread_self(), sizeof(get_cpus), &get_cpus);
+    if (ret != 0) {
+        errx(EXIT_FAILURE, "pthread_getaffinity_np failed for parent!");
+    }
+
+    if (!CPU_EQUAL_S(sizeof(cpus), &cpus, &get_cpus)) {
+        errx(EXIT_FAILURE, "get cpuset is not equal to set cpuset on proc 0");
+    }
+
+    /* Negative test case with empty cpumask */
+    CPU_ZERO(&cpus);
+    ret = pthread_setaffinity_np(pthread_self(), sizeof(cpus), &cpus);
+    if (ret != EINVAL) {
+        errx(EXIT_FAILURE, "pthread_setaffinity_np with empty cpumask did not return EINVAL!");
+    }
+
+    printf("TEST OK\n");
+    return 0;
+}

LibOS/shim/test/regression/sched.c

@@ -5,8 +5,8 @@
 #include <sys/resource.h>
 #include <sys/time.h>
 
-/* This test checks that our dummy implementations work correctly. None of the
- * below syscalls are actually propagated to the host OS or change anything.
+/* This test checks that our dummy implementations work correctly. None of the below syscalls except
+ * sched_setaffinity and sched_getaffinity are actually propagated to the host OS or change anything
  * NOTE: This test works correctly only on Graphene (not on Linux). */
 
 int main(int argc, char** argv) {
@@ -29,6 +29,7 @@ int main(int argc, char** argv) {
 
     cpu_set_t my_set;
     CPU_ZERO(&my_set);
+    CPU_SET(0, &my_set);
     if (sched_setaffinity(0, sizeof(cpu_set_t), &my_set) == -1) {
         perror("Error setting affinity");
         return 1;

LibOS/shim/test/regression/sched_set_get_affinity.c

@@ -0,0 +1,71 @@
+/* SPDX-License-Identifier: LGPL-3.0-or-later */
+/* Copyright (C) 2020 Intel Corporation */
+
+/*
+ * Test setting/getting cpu affinity on a single processor or multiple processors.
+ */
+
+#define _GNU_SOURCE
+#include <assert.h>
+#include <err.h>
+#include <errno.h>
+#include <fcntl.h>
+#include <sched.h>
+#include <stdbool.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+int main(int argc, const char** argv) {
+    int ret;
+    cpu_set_t cpus, get_cpus;
+    long numprocs = sysconf(_SC_NPROCESSORS_ONLN);
+    if (numprocs < 0) {
+        err(EXIT_FAILURE, "Failed to retrieve the number of logical processors!");
+    }
+
+    for (long i = 0; i < numprocs; i++) {
+        printf("Testing processor id: %ld\n", i);
+        CPU_ZERO(&cpus);
+        CPU_ZERO(&get_cpus);
+        CPU_SET(i, &cpus);
+        ret = sched_setaffinity(0, sizeof(cpus), &cpus);
+        if (ret < 0) {
+            errx(EXIT_FAILURE, "Failed to set affinity for current thread, core id: %ld", i);
+        }
+        ret = sched_getaffinity(0, sizeof(get_cpus), &get_cpus);
+        if (ret < 0) {
+            errx(EXIT_FAILURE, "Failed to get affinity for current thread, core id: %ld", i);
+        }
+        if (!CPU_EQUAL_S(sizeof(cpus), &cpus, &get_cpus)) {
+            errx(EXIT_FAILURE, "The get cpu set is not equal to set on core id: %ld", i);
+        }
+    }
+
+    if (numprocs >= 2) {
+        /* test for multiple cpu affinity */
+        CPU_ZERO(&cpus);
+        CPU_ZERO(&get_cpus);
+        CPU_SET(0, &cpus);
+        CPU_SET(1, &cpus);
+        ret = sched_setaffinity(0, sizeof(cpus), &cpus);
+        if (ret < 0) {
+            err(EXIT_FAILURE, "Failed to set multiple affinity for current thread");
+        }
+        ret = sched_getaffinity(0, sizeof(get_cpus), &get_cpus);
+        if (ret < 0) {
+            err(EXIT_FAILURE, "Failed to get multiple affinity for current thread");
+        }
+        if (!CPU_EQUAL_S(sizeof(cpus), &cpus, &get_cpus)) {
+            errx(EXIT_FAILURE, "The get cpu set is not equal to set on core id: 0 & 1");
+        }
+    } else {
+        printf("Multiple CPU affinity test skipped since only one core was identified\n");
+    }
+
+    printf("TEST OK\n");
+    return 0;
+}

LibOS/shim/test/regression/test_libos.py

@@ -518,6 +518,13 @@ class TC_30_Syscall(RegressionTestCase):
         self.assertIn('child OK', stdout);
         self.assertIn('parent OK', stdout);
 
+    def test_101_sched_set_get_cpuaffinity(self):
+        stdout, _ = self.run_binary(['sched_set_get_affinity'])
+        self.assertIn('TEST OK', stdout)
+
+    def test_102_pthread_set_get_affinity(self):
+        stdout, _ = self.run_binary(['pthread_set_get_affinity', '1000'])
+        self.assertIn('TEST OK', stdout)
 
 @unittest.skipUnless(HAS_SGX,
     'This test is only meaningful on SGX PAL because only SGX catches raw '

Pal/include/arch/x86_64/Linux/sysdep-arch.h

@@ -155,6 +155,9 @@
 #undef INTERNAL_SYSCALL_ERRNO_P
 #define INTERNAL_SYSCALL_ERRNO_P(val) (-((long)val))
 
+#undef INTERNAL_SYSCALL_ERRNO_RANGE
+#define INTERNAL_SYSCALL_ERRNO_RANGE(val) ((val) >= -133 /* EHWPOISON */ && (val) <= -1 /* EPERM */)
+
 #define LOAD_ARGS_0()
 #define LOAD_REGS_0
 #define ASM_ARGS_0

Pal/include/lib/api.h

@@ -64,6 +64,14 @@ typedef ptrdiff_t ssize_t;
         (((x) & ((x) - 1)) == 0); \
     })
 
+#define DIV_ROUND_UP(n,d)               (((n) + (d) - 1) / (d))
+
+#define BITS_IN_BYTE                    8
+#define BITS_IN_TYPE(type)              (sizeof(type) * BITS_IN_BYTE)
+#define BITS_TO_LONGS(nr)               DIV_ROUND_UP(nr, BITS_IN_TYPE(long))
+/* Note: This macro is not intended for use when nbits == BITS_IN_TYPE(type) */
+#define SET_HIGHEST_N_BITS(type, nbits) (~(((uint64_t)1 << (BITS_IN_TYPE(type) - (nbits))) - 1))
+
 #define IS_ALIGNED(val, alignment)     ((val) % (alignment) == 0)
 #define ALIGN_DOWN(val, alignment)     ((val) - (val) % (alignment))
 #define ALIGN_UP(val, alignment)       ALIGN_DOWN((val) + (alignment) - 1, alignment)

Pal/include/pal/pal.h

@@ -511,6 +511,35 @@ noreturn void DkThreadExit(PAL_PTR clear_child_tid);
  */
 PAL_BOL DkThreadResume(PAL_HANDLE thread);
 
+/*!
+ * \brief Sets the CPU affinity of a thread.
+ *
+ * All bit positions exceeding the count of host CPUs are ignored. Returns an error if no CPUs were
+ * selected.
+ *
+ * \param thread PAL thread for which to set the CPU affinity.
+ * \param cpumask_size size in bytes of the bitmask pointed by \a cpu_mask.
+ * \param cpu_mask pointer to the new CPU mask.
+ *
+ * \return Returns 1 on success, 0 on failure. Use PAL_ERRNO() to get the actual error code.
+ */
+PAL_BOL DkThreadSetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask);
+
+/*!
+ * \brief Gets the CPU affinity of a thread.
+ *
+ * This function assumes that \a cpumask_size is valid and greater than 0. Also, \a cpumask_size
+ * must be able to fit all the processors in the host and must be aligned by sizeof(long). For
+ * example, if the host supports 4 CPUs, \a cpumask_size should be 8 bytes.
+ *
+ * \param thread PAL thread for which to get the CPU affinity.
+ * \param cpumask_size size in bytes of the bitmask pointed by \a cpu_mask.
+ * \param cpu_mask pointer to hold the current CPU mask.
+ *
+ * \return Returns 1 on success, 0 on failure. Use PAL_ERRNO() to get the actual error code.
+ */
+PAL_BOL DkThreadGetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask);
+
 /*
  * Exception Handling
  */

Pal/src/db_threading.c

@@ -74,3 +74,29 @@ PAL_BOL DkThreadResume(PAL_HANDLE threadHandle) {
 
     LEAVE_PAL_CALL_RETURN(PAL_TRUE);
 }
+
+PAL_BOL DkThreadSetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask) {
+    ENTER_PAL_CALL(DkThreadSetCpuAffinity);
+
+    int ret = _DkThreadSetCpuAffinity(thread, cpumask_size, cpu_mask);
+
+    if (ret < 0) {
+        _DkRaiseFailure(-ret);
+        LEAVE_PAL_CALL_RETURN(PAL_FALSE);
+    }
+
+    LEAVE_PAL_CALL_RETURN(PAL_TRUE);
+}
+
+PAL_BOL DkThreadGetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask) {
+    ENTER_PAL_CALL(DkThreadGetCpuAffinity);
+
+    int ret = _DkThreadGetCpuAffinity(thread, cpumask_size, cpu_mask);
+
+    if (ret < 0) {
+        _DkRaiseFailure(-ret);
+        LEAVE_PAL_CALL_RETURN(PAL_FALSE);
+    }
+
+    LEAVE_PAL_CALL_RETURN(PAL_TRUE);
+}

Pal/src/host/Linux-SGX/db_main.c

@@ -451,6 +451,8 @@ noreturn void pal_linux_main(char* uptr_libpal_uri, size_t libpal_uri_len, char*
     PAL_HANDLE first_thread = malloc(HANDLE_SIZE(thread));
     SET_HANDLE_TYPE(first_thread, thread);
     first_thread->thread.tcs = g_enclave_base + GET_ENCLAVE_TLS(tcs_offset);
+    /* child threads are assigned TIDs 2,3,...; see pal_start_thread() */
+    first_thread->thread.tid = 1;
     g_pal_control.first_thread = first_thread;
     SET_ENCLAVE_TLS(thread, &first_thread->thread);
 

Pal/src/host/Linux-SGX/db_misc.c

@@ -311,7 +311,7 @@ int _DkCpuIdRetrieve(unsigned int leaf, unsigned int subleaf, unsigned int value
 
     /* the cpu core info cannot be cached due to its data varying depending on the calling thread */
     if (leaf == CPUID_EXT_TOPOLOGY_ENUMERATION_LEAF ||
-        leaf == CPUID_V2EXT_TOPOLOGY_ENUMERATION_LEAF) {
+            leaf == CPUID_V2EXT_TOPOLOGY_ENUMERATION_LEAF) {
         skip_cache = true;
     }
 

Pal/src/host/Linux-SGX/db_threading.c

@@ -45,7 +45,8 @@ extern void* g_enclave_base;
  * ensure uniqueness if needed in the future
  */
 static PAL_IDX pal_assign_tid(void) {
-    static struct atomic_int tid = ATOMIC_INIT(0);
+    /* tid 1 is assigned to the first thread; see pal_linux_main() */
+    static struct atomic_int tid = ATOMIC_INIT(1);
     return __atomic_add_fetch(&tid.counter, 1, __ATOMIC_SEQ_CST);
 }
 
@@ -57,7 +58,8 @@ void pal_start_thread(void) {
         if (!tmp->tcs) {
             new_thread = tmp;
             new_thread->tid = pal_assign_tid();
-            new_thread->tcs = g_enclave_base + GET_ENCLAVE_TLS(tcs_offset);
+            __atomic_store_n(&new_thread->tcs, (g_enclave_base + GET_ENCLAVE_TLS(tcs_offset)),
+                             __ATOMIC_RELEASE);
             break;
         }
     _DkInternalUnlock(&g_thread_list_lock);
@@ -105,6 +107,11 @@ int _DkThreadCreate(PAL_HANDLE* handle, int (*callback)(void*), const void* para
     if (IS_ERR(ret))
         return unix_to_pal_error(ERRNO(ret));
 
+    /* There can be subtle race between the parent and child so hold the parent until child updates
+       its tcs. */
+    while (!__atomic_load_n(&new_thread->thread.tcs, __ATOMIC_ACQUIRE))
+        CPU_RELAX();
+
     *handle = new_thread;
     return 0;
 }
@@ -145,6 +152,16 @@ int _DkThreadResume(PAL_HANDLE threadHandle) {
     return IS_ERR(ret) ? unix_to_pal_error(ERRNO(ret)) : ret;
 }
 
+int _DkThreadSetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask) {
+    int ret = ocall_sched_setaffinity(thread->thread.tcs, cpumask_size, cpu_mask);
+    return IS_ERR(ret) ? unix_to_pal_error(ERRNO(ret)) : ret;
+}
+
+int _DkThreadGetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask) {
+    int ret = ocall_sched_getaffinity(thread->thread.tcs, cpumask_size, cpu_mask);
+    return IS_ERR(ret) ? unix_to_pal_error(ERRNO(ret)) : ret;
+}
+
 struct handle_ops g_thread_ops = {
     /* nothing */
 };

Pal/src/host/Linux-SGX/enclave_ocalls.c

@@ -1470,3 +1470,83 @@ out:
     sgx_reset_ustack(old_ustack);
     return retval;
 }
+
+int ocall_sched_setaffinity(void* tcs, size_t cpumask_size, void* cpu_mask) {
+    int retval = 0;
+    ms_ocall_sched_setaffinity_t* ms;
+
+    void* old_ustack = sgx_prepare_ustack();
+    ms = sgx_alloc_on_ustack_aligned(sizeof(*ms), alignof(*ms));
+    if (!ms) {
+        sgx_reset_ustack(old_ustack);
+        return -EPERM;
+    }
+
+    WRITE_ONCE(ms->ms_tcs, tcs);
+    WRITE_ONCE(ms->ms_cpumask_size, cpumask_size);
+    void* untrusted_cpu_mask = sgx_copy_to_ustack(cpu_mask, cpumask_size);
+    if (!untrusted_cpu_mask) {
+        sgx_reset_ustack(old_ustack);
+        return -EPERM;
+    }
+    WRITE_ONCE(ms->ms_cpu_mask, untrusted_cpu_mask);
+
+    retval = sgx_exitless_ocall(OCALL_SCHED_SETAFFINITY, ms);
+    if (IS_ERR(retval) && !IS_UNIX_ERR(retval))
+        retval = -EPERM;
+
+    sgx_reset_ustack(old_ustack);
+    return retval;
+}
+
+static bool is_cpumask_valid(void* cpu_mask, size_t cpumask_size) {
+    size_t max_cpumask_bits = cpumask_size * BITS_IN_BYTE;
+    size_t valid_cpumask_bits = g_pal_control.cpu_info.online_logical_cores;
+    size_t invalid_bits = max_cpumask_bits - valid_cpumask_bits;
+
+    if (invalid_bits == 0)
+        return true;
+
+    /* create an invalid cpu_mask bits */
+    unsigned long invalid_cpumask = SET_HIGHEST_N_BITS(unsigned long, invalid_bits);
+
+    /* Extract last 64bits to check if any invalid cpu bits are set */
+    int idx = (cpumask_size / sizeof(unsigned long)) - 1;
+    unsigned long cpumask = ((unsigned long*)cpu_mask)[idx];
+
+    return !(cpumask & invalid_cpumask);
+}
+
+int ocall_sched_getaffinity(void* tcs, size_t cpumask_size, void* cpu_mask) {
+    int retval = 0;
+    ms_ocall_sched_getaffinity_t* ms;
+
+    void* old_ustack = sgx_prepare_ustack();
+    ms = sgx_alloc_on_ustack_aligned(sizeof(*ms), alignof(*ms));
+    if (!ms) {
+        sgx_reset_ustack(old_ustack);
+        return -EPERM;
+    }
+
+    WRITE_ONCE(ms->ms_tcs, tcs);
+    WRITE_ONCE(ms->ms_cpumask_size, cpumask_size);
+    void* untrusted_cpu_mask = sgx_copy_to_ustack(cpu_mask, cpumask_size);
+    if (!untrusted_cpu_mask) {
+        sgx_reset_ustack(old_ustack);
+        return -EPERM;
+    }
+    WRITE_ONCE(ms->ms_cpu_mask, untrusted_cpu_mask);
+
+    retval = sgx_exitless_ocall(OCALL_SCHED_GETAFFINITY, ms);
+    if (IS_ERR(retval) && !IS_UNIX_ERR(retval))
+        retval = -EPERM;
+
+    if (retval > 0 && !sgx_copy_to_enclave(cpu_mask, cpumask_size, untrusted_cpu_mask, retval))
+        retval = -EPERM;
+
+    if (retval > 0 && !is_cpumask_valid(cpu_mask, cpumask_size))
+        retval = -EPERM;
+
+    sgx_reset_ustack(old_ustack);
+    return retval;
+}

Pal/src/host/Linux-SGX/enclave_ocalls.h

@@ -69,6 +69,10 @@ int ocall_shutdown(int sockfd, int how);
 
 int ocall_resume_thread(void* tcs);
 
+int ocall_sched_setaffinity(void* tcs, size_t cpumask_size, void* cpu_mask);
+
+int ocall_sched_getaffinity(void* tcs, size_t cpumask_size, void* cpu_mask);
+
 int ocall_clone_thread(void);
 
 int ocall_create_process(const char* uri, size_t nargs, const char** args, int* stream_fd,

Pal/src/host/Linux-SGX/ocall_types.h

@@ -42,6 +42,8 @@ enum {
     OCALL_MKDIR,
     OCALL_GETDENTS,
     OCALL_RESUME_THREAD,
+    OCALL_SCHED_SETAFFINITY,
+    OCALL_SCHED_GETAFFINITY,
     OCALL_CLONE_THREAD,
     OCALL_CREATE_PROCESS,
     OCALL_FUTEX,
@@ -171,6 +173,18 @@ typedef struct {
     const char* ms_args[];
 } ms_ocall_create_process_t;
 
+typedef struct {
+    void* ms_tcs;
+    size_t ms_cpumask_size;
+    void* ms_cpu_mask;
+} ms_ocall_sched_setaffinity_t;
+
+typedef struct {
+    void* ms_tcs;
+    size_t ms_cpumask_size;
+    void* ms_cpu_mask;
+} ms_ocall_sched_getaffinity_t;
+
 typedef struct {
     uint32_t* ms_futex;
     int ms_op, ms_val;

Pal/src/host/Linux-SGX/pal_linux.h

@@ -23,10 +23,11 @@
 #include "sysdep-arch.h"
 #include "uthash.h"
 
-#define IS_ERR   INTERNAL_SYSCALL_ERROR
-#define IS_ERR_P INTERNAL_SYSCALL_ERROR_P
-#define ERRNO    INTERNAL_SYSCALL_ERRNO
-#define ERRNO_P  INTERNAL_SYSCALL_ERRNO_P
+#define IS_ERR      INTERNAL_SYSCALL_ERROR
+#define IS_ERR_P    INTERNAL_SYSCALL_ERROR_P
+#define ERRNO       INTERNAL_SYSCALL_ERRNO
+#define ERRNO_P     INTERNAL_SYSCALL_ERRNO_P
+#define IS_UNIX_ERR INTERNAL_SYSCALL_ERRNO_RANGE
 
 extern struct pal_linux_state {
     PAL_NUM parent_process_id;

Pal/src/host/Linux-SGX/sgx_enclave.c

@@ -251,7 +251,34 @@ static long sgx_ocall_getdents(void* pms) {
 
 static long sgx_ocall_resume_thread(void* pms) {
     ODEBUG(OCALL_RESUME_THREAD, pms);
-    return interrupt_thread(pms);
+    int tid = get_tid_from_tcs(pms);
+    if (tid < 0)
+        return tid;
+
+    long ret = INLINE_SYSCALL(tgkill, 3, g_pal_enclave.pal_sec.pid, tid, SIGCONT);
+    return ret;
+}
+
+static long sgx_ocall_sched_setaffinity(void* pms) {
+    ms_ocall_sched_setaffinity_t* ms = (ms_ocall_sched_setaffinity_t*)pms;
+    ODEBUG(OCALL_SCHED_SETAFFINITY, ms);
+    int tid = get_tid_from_tcs(ms->ms_tcs);
+    if (tid < 0)
+        return tid;
+
+    long ret = INLINE_SYSCALL(sched_setaffinity, 3, tid, ms->ms_cpumask_size, ms->ms_cpu_mask);
+    return ret;
+}
+
+static long sgx_ocall_sched_getaffinity(void* pms) {
+    ms_ocall_sched_getaffinity_t* ms = (ms_ocall_sched_getaffinity_t*)pms;
+    ODEBUG(OCALL_SCHED_GETAFFINITY, ms);
+    int tid = get_tid_from_tcs(ms->ms_tcs);
+    if (tid < 0)
+        return tid;
+
+    long ret = INLINE_SYSCALL(sched_getaffinity, 3, tid, ms->ms_cpumask_size, ms->ms_cpu_mask);
+    return ret;
 }
 
 static long sgx_ocall_clone_thread(void* pms) {
@@ -669,6 +696,8 @@ sgx_ocall_fn_t ocall_table[OCALL_NR] = {
     [OCALL_MKDIR]            = sgx_ocall_mkdir,
     [OCALL_GETDENTS]         = sgx_ocall_getdents,
     [OCALL_RESUME_THREAD]    = sgx_ocall_resume_thread,
+    [OCALL_SCHED_SETAFFINITY]= sgx_ocall_sched_setaffinity,
+    [OCALL_SCHED_GETAFFINITY]= sgx_ocall_sched_getaffinity,
     [OCALL_CLONE_THREAD]     = sgx_ocall_clone_thread,
     [OCALL_CREATE_PROCESS]   = sgx_ocall_create_process,
     [OCALL_FUTEX]            = sgx_ocall_futex,

Pal/src/host/Linux-SGX/sgx_internal.h

@@ -126,7 +126,7 @@ void async_exit_pointer(void);
 void eresume_pointer(void);
 void async_exit_pointer_end(void);
 
-int interrupt_thread(void* tcs);
+int get_tid_from_tcs(void* tcs);
 int clone_thread(void);
 
 void create_tcs_mapper(void* tcs_base, unsigned int thread_num);

Pal/src/host/Linux-SGX/sgx_thread.c

@@ -284,13 +284,13 @@ int clone_thread(void) {
     return 0;
 }
 
-int interrupt_thread(void* tcs) {
+int get_tid_from_tcs(void* tcs) {
     int index = (sgx_arch_tcs_t*)tcs - g_enclave_tcs;
     struct thread_map* map = &g_enclave_thread_map[index];
     if (index >= g_enclave_thread_num)
         return -EINVAL;
     if (!map->tid)
         return -EINVAL;
-    INLINE_SYSCALL(tgkill, 3, g_pal_enclave.pal_sec.pid, map->tid, SIGCONT);
-    return 0;
+
+    return map->tid;
 }

Pal/src/host/Linux/db_threading.c

@@ -279,6 +279,18 @@ int _DkThreadResume(PAL_HANDLE threadHandle) {
     return 0;
 }
 
+int _DkThreadSetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask) {
+    int ret = INLINE_SYSCALL(sched_setaffinity, 3, thread->thread.tid, cpumask_size, cpu_mask);
+
+    return IS_ERR(ret) ? unix_to_pal_error(ERRNO(ret)) : ret;
+}
+
+int _DkThreadGetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask) {
+    int ret = INLINE_SYSCALL(sched_getaffinity, 3, thread->thread.tid, cpumask_size, cpu_mask);
+
+    return IS_ERR(ret) ? unix_to_pal_error(ERRNO(ret)) : ret;
+}
+
 struct handle_ops g_thread_ops = {
     /* nothing */
 };

Pal/src/host/Skeleton/db_threading.c

@@ -19,7 +19,7 @@ int _DkThreadCreate(PAL_HANDLE* handle, int (*callback)(void*), const void* para
     return -PAL_ERROR_NOTIMPLEMENTED;
 }
 
-int _DkThreadDelayExecution(unsigned long* duration) {
+int _DkThreadDelayExecution(uint64_t* duration) {
     return -PAL_ERROR_NOTIMPLEMENTED;
 }
 
@@ -41,6 +41,14 @@ int _DkThreadResume(PAL_HANDLE threadHandle) {
     return -PAL_ERROR_NOTIMPLEMENTED;
 }
 
+int _DkThreadSetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask) {
+    return -PAL_ERROR_NOTIMPLEMENTED;
+}
+
+int _DkThreadGetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask) {
+    return -PAL_ERROR_NOTIMPLEMENTED;
+}
+
 struct handle_ops g_thread_ops = {
     /* nothing */
 };

Pal/src/pal-symbols

@@ -6,6 +6,8 @@ DkThreadDelayExecution
 DkThreadYieldExecution
 DkThreadExit
 DkThreadResume
+DkThreadSetCpuAffinity
+DkThreadGetCpuAffinity
 DkMutexCreate
 DkNotificationEventCreate
 DkSynchronizationEventCreate

Pal/src/pal_internal.h

@@ -233,6 +233,8 @@ void _DkThreadYieldExecution(void);
 int _DkThreadResume(PAL_HANDLE threadHandle);
 int _DkProcessCreate(PAL_HANDLE* handle, const char* uri, const char** args);
 noreturn void _DkProcessExit(int exitCode);
+int _DkThreadSetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask);
+int _DkThreadGetCpuAffinity(PAL_HANDLE thread, PAL_NUM cpumask_size, PAL_PTR cpu_mask);
 
 /* DkMutex calls */
 int _DkMutexCreate(PAL_HANDLE* handle, int initialCount);