diff --git a/ansible_mitogen/affinity.py b/ansible_mitogen/affinity.py index 57926516..09a6acee 100644 --- a/ansible_mitogen/affinity.py +++ b/ansible_mitogen/affinity.py @@ -79,6 +79,7 @@ import multiprocessing import os import struct +import mitogen.core import mitogen.parent @@ -246,8 +247,19 @@ class FixedPolicy(Policy): class LinuxPolicy(FixedPolicy): + def _mask_to_bytes(self, mask): + """ + Convert the (type long) mask to a cpu_set_t. + """ + chunks = [] + shiftmask = (2 ** 64) - 1 + for x in range(16): + chunks.append(struct.pack('>= 64 + return mitogen.core.b('').join(chunks) + def _set_cpu_mask(self, mask): - s = struct.pack('L', mask) + s = self._mask_to_bytes(mask) _sched_setaffinity(os.getpid(), len(s), s) diff --git a/tests/ansible/tests/affinity_test.py b/tests/ansible/tests/affinity_test.py index 102608d4..8ee96085 100644 --- a/tests/ansible/tests/affinity_test.py +++ b/tests/ansible/tests/affinity_test.py @@ -211,5 +211,16 @@ class LinuxPolicyTest(testlib.TestCase): tf.close() +class MockLinuxPolicyTest(testlib.TestCase): + klass = ansible_mitogen.affinity.LinuxPolicy + + # Test struct.pack() in _set_cpu_mask(). + + def test_high_cpus(self): + policy = self.klass(cpu_count=4096) + for x in range(1, 4096, 32): + policy.assign_subprocess() + + if __name__ == '__main__': unittest2.main()