src/java.base/share/classes/sun/nio/ch/SelectorImpl.java

rev 50580 : [mq]: select-with-consumer

@@ -34,12 +34,14 @@
 import java.nio.channels.spi.AbstractSelector;
 import java.nio.channels.spi.SelectorProvider;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Iterator;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Consumer;
 
 
 /**
  * Base Selector implementation class.
  */

@@ -49,16 +51,19 @@
 {
     // The set of keys registered with this Selector
     private final Set<SelectionKey> keys;
 
     // The set of keys with data ready for an operation
-    protected final Set<SelectionKey> selectedKeys;
+    private final Set<SelectionKey> selectedKeys;
 
     // Public views of the key sets
     private final Set<SelectionKey> publicKeys;             // Immutable
     private final Set<SelectionKey> publicSelectedKeys;     // Removal allowed, but not addition
 
+    // used to check for reentrancy
+    private boolean inSelect;
+    
     protected SelectorImpl(SelectorProvider sp) {
         super(sp);
         keys = ConcurrentHashMap.newKeySet();
         selectedKeys = new HashSet<>();
         publicKeys = Collections.unmodifiableSet(keys);

@@ -81,17 +86,10 @@
         ensureOpen();
         return publicSelectedKeys;
     }
 
     /**
-     * Returns the public view of the selected-key set
-     */
-    protected final Set<SelectionKey> nioSelectedKeys() {
-        return publicSelectedKeys;
-    }
-
-    /**
      * Marks the beginning of a select operation that might block
      */
     protected final void begin(boolean blocking) {
         if (blocking) begin();
     }

@@ -104,39 +102,72 @@
     }
 
     /**
      * Selects the keys for channels that are ready for I/O operations.
      *
+     * @param action  the action to perform, can be null
      * @param timeout timeout in milliseconds to wait, 0 to not wait, -1 to
      *                wait indefinitely
      */
-    protected abstract int doSelect(long timeout) throws IOException;
+    protected abstract int doSelect(Consumer<SelectionKey> action, long timeout)
+        throws IOException;
 
-    private int lockAndDoSelect(long timeout) throws IOException {
+    private int lockAndDoSelect(Consumer<SelectionKey> action, long timeout)
+        throws IOException
+    {
         synchronized (this) {
             ensureOpen();
+            if (inSelect)
+                throw new IllegalStateException("select in progress");
+            inSelect = true;
+            try {
             synchronized (publicSelectedKeys) {
-                return doSelect(timeout);
+                    return doSelect(action, timeout);
+                }
+            } finally {
+                inSelect = false;
             }
         }
     }
 
     @Override
     public final int select(long timeout) throws IOException {
         if (timeout < 0)
             throw new IllegalArgumentException("Negative timeout");
-        return lockAndDoSelect((timeout == 0) ? -1 : timeout);
+        return lockAndDoSelect(null, (timeout == 0) ? -1 : timeout);
     }
 
     @Override
     public final int select() throws IOException {
-        return lockAndDoSelect(-1);
+        return lockAndDoSelect(null, -1);
     }
 
     @Override
     public final int selectNow() throws IOException {
-        return lockAndDoSelect(0);
+        return lockAndDoSelect(null, 0);
+    }
+
+    @Override
+    public final int select(Consumer<SelectionKey> action, long timeout)
+        throws IOException
+    {
+        Objects.requireNonNull(action);
+        if (timeout < 0)
+            throw new IllegalArgumentException("Negative timeout");
+        return lockAndDoSelect(action, (timeout == 0) ? -1 : timeout);
+    }
+
+    @Override
+    public final int select(Consumer<SelectionKey> action) throws IOException {
+        Objects.requireNonNull(action);
+        return lockAndDoSelect(action, -1);
+    }
+
+    @Override
+    public final int selectNow(Consumer<SelectionKey> action) throws IOException {
+        Objects.requireNonNull(action);
+        return lockAndDoSelect(action, 0);
     }
 
     /**
      * Invoked by implCloseSelector to close the selector.
      */

@@ -238,10 +269,43 @@
             }
         }
     }
 
     /**
+     * Invoked by selection operations to handle ready events. If an action
+     * is specified then it is invoked to handle the key, otherwise the key
+     * is added to the selected-key set (or updated when it is already in the
+     * set).
+     */
+    protected final int processReadyEvents(int rOps,
+                                           SelectionKeyImpl ski,
+                                           Consumer<SelectionKey> action) {
+        if (action != null) {
+            ski.translateAndSetReadyOps(rOps);
+            if ((ski.nioReadyOps() & ski.nioInterestOps()) != 0) {
+                action.accept(ski);
+                ensureOpen();
+                return 1;
+            }
+        } else {
+            assert Thread.holdsLock(publicSelectedKeys);
+            if (selectedKeys.contains(ski)) {
+                if (ski.translateAndUpdateReadyOps(rOps)) {
+                    return 1;
+                }
+            } else {
+                ski.translateAndSetReadyOps(rOps);
+                if ((ski.nioReadyOps() & ski.nioInterestOps()) != 0) {
+                    selectedKeys.add(ski);
+                    return 1;
+                }
+            }
+        }
+        return 0;
+    }
+
+    /**
      * Invoked by interestOps to ensure the interest ops are updated at the
      * next selection operation.
      */
     protected abstract void setEventOps(SelectionKeyImpl ski);
 }