[GR-38924] Improve multiprocessing.connection.wait.

PullRequest: graalpython/2286

Author: steve-s

graalpython/com.oracle.graal.python.test/src/tests/test_multiprocessing.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2022, 2022, Oracle and/or its affiliates. All rights reserved.
+# DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+#
+# The Universal Permissive License (UPL), Version 1.0
+#
+# Subject to the condition set forth below, permission is hereby granted to any
+# person obtaining a copy of this software, associated documentation and/or
+# data (collectively the "Software"), free of charge and under any and all
+# copyright rights in the Software, and any and all patent rights owned or
+# freely licensable by each licensor hereunder covering either (i) the
+# unmodified Software as contributed to or provided by such licensor, or (ii)
+# the Larger Works (as defined below), to deal in both
+#
+# (a) the Software, and
+#
+# (b) any piece of software and/or hardware listed in the lrgrwrks.txt file if
+# one is included with the Software each a "Larger Work" to which the Software
+# is contributed by such licensors),
+#
+# without restriction, including without limitation the rights to copy, create
+# derivative works of, display, perform, and distribute the Software and make,
+# use, sell, offer for sale, import, export, have made, and have sold the
+# Software and the Larger Work(s), and to sublicense the foregoing rights on
+# either these or other terms.
+#
+# This license is subject to the following condition:
+#
+# The above copyright notice and either this complete permission notice or at a
+# minimum a reference to the UPL must be included in all copies or substantial
+# portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+import multiprocessing
+import time
+from multiprocessing.connection import wait
+
+
+def test_wait_timeout():
+    timeout = 3
+    a, b = multiprocessing.Pipe()
+    x, y = multiprocessing.connection.Pipe(False)  # Truffle multiprocessing pipe
+    for fds in [[a, b], [x, y], [a, b, x, y]]:
+        start = time.monotonic()
+        res = wait(fds, timeout)
+        delta = time.monotonic() - start
+        assert not res
+        assert delta < timeout * 2
+        assert delta > timeout / 2
+
+
+def test_wait():
+    a, b = multiprocessing.Pipe()
+    x, y = multiprocessing.connection.Pipe(False)  # Truffle multiprocessing pipe
+    a.send(42)
+    res = wait([b, y], 3)
+    assert res == [b], "res1"
+    assert b.recv() == 42, "res2"
+    y.send(33)
+    res = wait([b, x], 3)
+    assert res == [x], "res3"
+    assert x.recv() == 33, "res4"
+    a.send(1)
+    y.send(2)
+    res = wait([b, x], 3)
+    assert set(res) == set([b, x])
+    assert b.recv() == 1
+    assert x.recv() == 2

graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_multiprocessing_spawn.txt

@@ -47,6 +47,7 @@
 *graalpython.lib-python.3.test.test_multiprocessing_spawn.TestWait.test_wait_slow
 *graalpython.lib-python.3.test.test_multiprocessing_spawn.TestWait.test_wait_socket
 *graalpython.lib-python.3.test.test_multiprocessing_spawn.TestWait.test_wait_socket_slow
+*graalpython.lib-python.3.test.test_multiprocessing_spawn.TestWait.test_wait_timeout
 *graalpython.lib-python.3.test.test_multiprocessing_spawn.WithProcessesTestBarrier.test_abort
 *graalpython.lib-python.3.test.test_multiprocessing_spawn.WithProcessesTestBarrier.test_abort_and_reset
 *graalpython.lib-python.3.test.test_multiprocessing_spawn.WithProcessesTestBarrier.test_action

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/MultiprocessingModuleBuiltins.java

@@ -54,6 +54,7 @@
 import com.oracle.graal.python.builtins.objects.PNone;
 import com.oracle.graal.python.builtins.objects.buffer.PythonBufferAccessLibrary;
 import com.oracle.graal.python.builtins.objects.bytes.PBytes;
+import com.oracle.graal.python.builtins.objects.common.SequenceNodes;
 import com.oracle.graal.python.builtins.objects.common.SequenceStorageNodes;
 import com.oracle.graal.python.builtins.objects.exception.OSErrorEnum;
 import com.oracle.graal.python.builtins.objects.ints.PInt;
@@ -71,10 +72,13 @@
 import com.oracle.graal.python.nodes.function.builtins.PythonBinaryBuiltinNode;
 import com.oracle.graal.python.nodes.function.builtins.PythonUnaryBuiltinNode;
 import com.oracle.graal.python.nodes.util.CannotCastException;
+import com.oracle.graal.python.nodes.util.CastToJavaDoubleNode;
 import com.oracle.graal.python.nodes.util.CastToJavaIntExactNode;
 import com.oracle.graal.python.nodes.util.CastToJavaIntLossyNode;
 import com.oracle.graal.python.nodes.util.CastToTruffleStringNode;
 import com.oracle.graal.python.runtime.GilNode;
+import com.oracle.graal.python.runtime.PosixSupportLibrary;
+import com.oracle.graal.python.runtime.PosixSupportLibrary.Timeval;
 import com.oracle.graal.python.runtime.PythonContext;
 import com.oracle.graal.python.runtime.PythonContext.SharedMultiprocessingData;
 import com.oracle.graal.python.runtime.sequence.PSequence;
@@ -356,32 +360,81 @@ PNone close(@SuppressWarnings("unused") long fd) {
         }
     }
 
-    @Builtin(name = "_select", minNumOfPositionalArgs = 1, parameterNames = {"rlist"})
+    @Builtin(name = "_select", minNumOfPositionalArgs = 4)
     @GenerateNodeFactory
     abstract static class SelectNode extends PythonBuiltinNode {
+        /*
+         * We would like to poll two different things with a timeout: the actual file descriptors
+         * and the Java managed LinkedBlockingQueues.
+         * 
+         * The LinkedBlockingQueue does not expose anything that would allow us to wait on multiple
+         * LinkedBlockingQueues at once, so we'd have to spawn a thread for each or roll out our own
+         * synchronization of take/offer to allow that.
+         * 
+         * The actual file descriptors could be backed by Java POSIX emulation layer, or by the
+         * native POSIX implementation -- the `select` can run actual native select, which we cannot
+         * easily interrupt from Java if one of the LinkedBlockingQueue is unblocked earlier than
+         * the native select returns.
+         * 
+         * Given all these complexities, for the time being, we do active waiting here, but at least
+         * without holding the GIL, and we also yield in every iteration.
+         */
+
         @Specialization
-        Object doGeneric(VirtualFrame frame, Object rlist,
+        Object doGeneric(VirtualFrame frame, Object multiprocessingFdsList, Object multiprocessingObjsList, Object posixFileObjsList, Object timeoutObj,
+                        @Cached PosixModuleBuiltins.FileDescriptorConversionNode fdConvertor,
                         @Cached PyObjectSizeNode sizeNode,
                         @Cached PyObjectGetItem getItem,
+                        @Cached SequenceNodes.GetObjectArrayNode getObjectArrayNode,
                         @Cached ListNodes.FastConstructListNode constructListNode,
                         @Cached CastToJavaIntLossyNode castToJava,
+                        @Cached CastToJavaDoubleNode castToDouble,
                         @Cached GilNode gil) {
-            ArrayBuilder<Integer> notEmpty = new ArrayBuilder<>();
-            SharedMultiprocessingData sharedData = getContext().getSharedMultiprocessingData();
-            PSequence pSequence = constructListNode.execute(frame, rlist);
-            for (int i = 0; i < sizeNode.execute(frame, pSequence); i++) {
+            PythonContext context = getContext();
+            SharedMultiprocessingData sharedData = context.getSharedMultiprocessingData();
+
+            PSequence pSequence = constructListNode.execute(frame, multiprocessingFdsList);
+            int size = sizeNode.execute(frame, pSequence);
+            int[] multiprocessingFds = new int[size];
+            for (int i = 0; i < size; i++) {
                 Object pythonObject = getItem.execute(frame, pSequence, i);
-                int fd = toInt(castToJava, pythonObject);
-                gil.release(true);
-                try {
-                    if (!sharedData.isBlocking(fd)) {
-                        notEmpty.add(fd);
+                multiprocessingFds[i] = toInt(castToJava, pythonObject);
+            }
+
+            Object[] posixFileObjs = getObjectArrayNode.execute(posixFileObjsList);
+            int[] posixFds = new int[posixFileObjs.length];
+            for (int i = 0; i < posixFileObjs.length; i++) {
+                posixFds[i] = toInt(castToJava, fdConvertor.execute(frame, posixFileObjs[i]));
+            }
+
+            double timeout = castToDouble.execute(timeoutObj);
+
+            Object[] multiprocessingObjs = getObjectArrayNode.execute(multiprocessingObjsList);
+            gil.release(true);
+            try {
+                boolean[] selectedMultiprocessingFds = new boolean[multiprocessingFds.length];
+                boolean[] selectedPosixFds = new boolean[posixFds.length];
+
+                doSelect(context.getPosixSupport(), sharedData, posixFds, selectedPosixFds, multiprocessingFds, selectedMultiprocessingFds, timeout);
+
+                ArrayBuilder<Object> result = new ArrayBuilder<>(4);
+                for (int i = 0; i < selectedMultiprocessingFds.length; i++) {
+                    if (selectedMultiprocessingFds[i]) {
+                        result.add(multiprocessingObjs[i]);
+                    }
+                }
+                for (int i = 0; i < selectedPosixFds.length; i++) {
+                    if (selectedPosixFds[i]) {
+                        result.add(posixFileObjs[i]);
                     }
-                } finally {
-                    gil.acquire();
                 }
+
+                return factory().createList(result.toArray(new Object[0]));
+            } catch (PosixSupportLibrary.PosixException e) {
+                throw raiseOSErrorFromPosixException(frame, e);
+            } finally {
+                gil.acquire();
             }
-            return factory().createList(notEmpty.toObjectArray(new Object[0]));
         }
 
         private static int toInt(CastToJavaIntLossyNode castToJava, Object pythonObject) {
@@ -391,6 +444,48 @@ private static int toInt(CastToJavaIntLossyNode castToJava, Object pythonObject)
                 throw CompilerDirectives.shouldNotReachHere();
             }
         }
+
+        @TruffleBoundary
+        private static void doSelect(Object posix, SharedMultiprocessingData sharedData,
+                        int[] posixFds, boolean[] selectedPosixFds,
+                        int[] multiprocessingFds, boolean[] selectedMultiprocessingFds,
+                        double timeoutInS) throws PosixSupportLibrary.PosixException {
+            PosixSupportLibrary posixLib = PosixSupportLibrary.getUncached();
+            boolean blocking = timeoutInS >= 0;
+            boolean untilReady = timeoutInS == 0;
+            long deadline = 0;
+            if (blocking && !untilReady) {
+                long timeout = (long) (timeoutInS * 1000_000_000.0);
+                deadline = System.nanoTime() + timeout;
+            }
+            while (true) {
+                boolean selected = false;
+                if (posixFds.length > 0) {
+                    PosixSupportLibrary.SelectResult selectResult = posixLib.select(posix, posixFds,
+                                    PythonUtils.EMPTY_INT_ARRAY, PythonUtils.EMPTY_INT_ARRAY, Timeval.SELECT_TIMEOUT_NOW);
+                    System.arraycopy(selectResult.getReadFds(), 0, selectedPosixFds, 0, selectedPosixFds.length);
+                    if (blocking) {
+                        for (boolean b : selectedPosixFds) {
+                            selected |= b;
+                        }
+                    }
+                }
+                for (int i = 0; i < multiprocessingFds.length; i++) {
+                    int fd = multiprocessingFds[i];
+                    selectedMultiprocessingFds[i] = !sharedData.isBlocking(fd);
+                    if (selectedMultiprocessingFds[i]) {
+                        selected = true;
+                    }
+                }
+                if (!blocking || selected) {
+                    return;
+                }
+                if (deadline != 0 && deadline - System.nanoTime() < 0) {
+                    return;
+                }
+                Thread.yield();
+            }
+        }
     }
 
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/SelectModuleBuiltins.java

@@ -86,6 +86,11 @@
 @CoreFunctions(defineModule = "select")
 public class SelectModuleBuiltins extends PythonBuiltins {
 
+    /*
+     * ATTENTION: if we ever add "poll" support, update the code in
+     * MultiprocessingModuleBuilins#SelectNode to use it if available
+     */
+
     public SelectModuleBuiltins() {
         addBuiltinConstant("error", PythonErrorType.OSError);
     }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/PosixSupportLibrary.java

@@ -983,6 +983,8 @@ public ByteBuffer getByteBuffer() {
      */
     @ValueType
     public static final class Timeval {
+        public static final Timeval SELECT_TIMEOUT_NOW = new Timeval(0, 0);
+
         private final long seconds;
         private final long microseconds;