numa-aware EventLoopGroup (#220)

also, latest submodules

Author: graebm

awscrt/common.py

@@ -0,0 +1,21 @@
+"""
+Cross-platform library for `awscrt`.
+"""
+import _awscrt
+from typing import List
+
+
+def get_cpu_group_count() -> int:
+    """
+    Returns number of processor groups on the system.
+
+    Useful for working with non-uniform memory access (NUMA) nodes.
+    """
+    return _awscrt.get_cpu_group_count()
+
+
+def get_cpu_count_for_group(group_idx: int) -> int:
+    """
+    Returns number of processors in a given group.
+    """
+    return _awscrt.get_cpu_count_for_group(group_idx)

awscrt/io.py

@@ -46,8 +46,13 @@ class EventLoopGroup(NativeResource):
     need to do async work will ask the EventLoopGroup for an event-loop to use.
 
     Args:
-        num_threads (int): Number of event-loops to create.
-            Pass 0 to create one for each processor on the machine.
+        num_threads (Optional[int]): Maximum number of event-loops to create.
+            If unspecified, one is created for each processor on the machine.
+
+        cpu_group (Optional[int]): Optional processor group to which all
+            threads will be pinned. Useful for systems with non-uniform
+            memory access (NUMA) nodes. If specified, the number of threads
+            will be capped at the number of processors in the group.
 
     Attributes:
         shutdown_event (threading.Event): Signals when EventLoopGroup's threads
@@ -57,16 +62,26 @@ class EventLoopGroup(NativeResource):
 
     __slots__ = ('shutdown_event')
 
-    def __init__(self, num_threads=0):
+    def __init__(self, num_threads=None, cpu_group=None):
         super().__init__()
 
+        if num_threads is None:
+            # C uses 0 to indicate defaults
+            num_threads = 0
+
+        if cpu_group is None:
+            is_pinned = False
+            cpu_group = 0
+        else:
+            is_pinned = True
+
         shutdown_event = threading.Event()
 
         def on_shutdown():
             shutdown_event.set()
 
         self.shutdown_event = shutdown_event
-        self._binding = _awscrt.event_loop_group_new(num_threads, on_shutdown)
+        self._binding = _awscrt.event_loop_group_new(num_threads, is_pinned, cpu_group, on_shutdown)
 
 
 class HostResolverBase(NativeResource):

crt/aws-c-auth

@@ -0,0 +1,26 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+#include "common.h"
+
+#include <aws/common/system_info.h>
+
+PyObject *aws_py_get_cpu_group_count(PyObject *self, PyObject *args) {
+    (void)self;
+    (void)args;
+    uint16_t count = aws_get_cpu_group_count();
+    return PyLong_FromUnsignedLong(count);
+}
+
+PyObject *aws_py_get_cpu_count_for_group(PyObject *self, PyObject *args) {
+    (void)self;
+
+    uint16_t group_idx;
+    if (!PyArg_ParseTuple(args, "H", &group_idx)) {
+        return NULL;
+    }
+
+    size_t count = aws_get_cpu_count_for_group(group_idx);
+    return PyLong_FromSize_t(count);
+}

crt/aws-c-common

@@ -0,0 +1,17 @@
+#ifndef AWS_CRT_PYTHON_COMMON_H
+#define AWS_CRT_PYTHON_COMMON_H
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+/**
+ * This file includes definitions for common aws-c-common functions.
+ */
+
+#include "module.h"
+
+PyObject *aws_py_get_cpu_group_count(PyObject *self, PyObject *args);
+PyObject *aws_py_get_cpu_count_for_group(PyObject *self, PyObject *args);
+
+#endif /* AWS_CRT_PYTHON_COMMON_H */

crt/aws-c-io

@@ -131,8 +131,10 @@ PyObject *aws_py_event_loop_group_new(PyObject *self, PyObject *args) {
     struct aws_allocator *allocator = aws_py_get_allocator();
 
     uint16_t num_threads;
+    int is_pinned;
+    uint16_t cpu_group;
     PyObject *shutdown_complete_py;
-    if (!PyArg_ParseTuple(args, "HO", &num_threads, &shutdown_complete_py)) {
+    if (!PyArg_ParseTuple(args, "HpHO", &num_threads, &is_pinned, &cpu_group, &shutdown_complete_py)) {
         return NULL;
     }
 
@@ -146,7 +148,12 @@ PyObject *aws_py_event_loop_group_new(PyObject *self, PyObject *args) {
         .shutdown_callback_user_data = binding,
     };
 
-    binding->native = aws_event_loop_group_new_default(allocator, num_threads, &shutdown_options);
+    if (is_pinned) {
+        binding->native =
+            aws_event_loop_group_new_default_pinned_to_cpu_group(allocator, num_threads, cpu_group, &shutdown_options);
+    } else {
+        binding->native = aws_event_loop_group_new_default(allocator, num_threads, &shutdown_options);
+    }
     if (binding->native == NULL) {
         PyErr_SetAwsLastError();
         goto elg_init_failed;

source/common.c

@@ -5,6 +5,7 @@
 #include "module.h"
 
 #include "auth.h"
+#include "common.h"
 #include "crypto.h"
 #include "event_stream.h"
 #include "http.h"
@@ -477,6 +478,8 @@ static PyMethodDef s_module_methods[] = {
     AWS_PY_METHOD_DEF(get_error_name, METH_VARARGS),
     AWS_PY_METHOD_DEF(get_error_message, METH_VARARGS),
     AWS_PY_METHOD_DEF(get_corresponding_builtin_exception, METH_VARARGS),
+    AWS_PY_METHOD_DEF(get_cpu_group_count, METH_VARARGS),
+    AWS_PY_METHOD_DEF(get_cpu_count_for_group, METH_VARARGS),
 
     /* IO */
     AWS_PY_METHOD_DEF(is_alpn_available, METH_NOARGS),

source/common.h

@@ -0,0 +1,15 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0.
+
+from test import NativeResourceTest
+from awscrt.common import *
+
+
+class TestSystemInfo(NativeResourceTest):
+    def test_get_cpu_group_count(self):
+        self.assertGreater(get_cpu_group_count(), 0)
+
+    def test_get_cpu_count_for_group(self):
+        group_count = get_cpu_group_count()
+        for group_i in range(group_count):
+            self.assertGreater(get_cpu_count_for_group(group_i), 0)