issue #557: support correct cpu_set_t size

ansible_mitogen/affinity.py

@@ -79,6 +79,7 @@
 import os
 import struct
 
+import mitogen.core
 import mitogen.parent
 
 
@@ -246,8 +247,19 @@ def assign_subprocess(self):
 
 
 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('<Q', mask & shiftmask))
+            mask >>= 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)
 
 

tests/ansible/tests/affinity_test.py

@@ -211,5 +211,16 @@ def test_clear_on_popen(self):
             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()