SmackReactor.java
/**
*
* Copyright 2018-2023 Florian Schmaus
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.jivesoftware.smack;
import java.io.IOException;
import java.nio.channels.CancelledKeyException;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.SelectableChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Date;
import java.util.Iterator;
import java.util.List;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.DelayQueue;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.logging.Level;
import java.util.logging.Logger;
/**
* The SmackReactor for non-blocking I/O.
* <p>
* Highlights include:
* <ul>
* <li>Multiple reactor threads</li>
* <li>Scheduled actions</li>
* </ul>
*
* <pre>
*
* ) ) )
* ( ( (
* ) ) )
* (~~~~~~~~~)
* | Smack |
* |Reactor|
* I _._
* I /' `\
* I | |
* f | |~~~~~~~~~~~~~~|
* .' | | # # # # |
* '______|___|___________###|
* </pre>
*/
public class SmackReactor {
private static final Logger LOGGER = Logger.getLogger(SmackReactor.class.getName());
private static final int DEFAULT_REACTOR_THREAD_COUNT = 2;
private static final int PENDING_SET_INTEREST_OPS_MAX_BATCH_SIZE = 1024;
private static SmackReactor INSTANCE;
static synchronized SmackReactor getInstance() {
if (INSTANCE == null) {
INSTANCE = new SmackReactor("DefaultReactor");
}
return INSTANCE;
}
private final Selector selector;
private final String reactorName;
private final List<Reactor> reactorThreads = Collections.synchronizedList(new ArrayList<>());
private final DelayQueue<ScheduledAction> scheduledActions = new DelayQueue<>();
private final Lock registrationLock = new ReentrantLock();
/**
* The semaphore protecting the handling of the actions. Note that it is
* initialized with -1, which basically means that one thread will always do I/O using
* select().
*/
private final Semaphore actionsSemaphore = new Semaphore(-1, false);
private final Queue<SelectionKey> pendingSelectionKeys = new ConcurrentLinkedQueue<>();
private final Queue<SetInterestOps> pendingSetInterestOps = new ConcurrentLinkedQueue<>();
SmackReactor(String reactorName) {
this.reactorName = reactorName;
try {
selector = Selector.open();
}
catch (IOException e) {
throw new IllegalStateException(e);
}
setReactorThreadCount(DEFAULT_REACTOR_THREAD_COUNT);
}
public SelectionKey registerWithSelector(SelectableChannel channel, int ops, ChannelSelectedCallback callback)
throws ClosedChannelException {
SelectionKeyAttachment selectionKeyAttachment = new SelectionKeyAttachment(callback);
registrationLock.lock();
try {
selector.wakeup();
return channel.register(selector, ops, selectionKeyAttachment);
} finally {
registrationLock.unlock();
}
}
public void setInterestOps(SelectionKey selectionKey, int interestOps) {
SetInterestOps setInterestOps = new SetInterestOps(selectionKey, interestOps);
pendingSetInterestOps.add(setInterestOps);
selector.wakeup();
}
private static final class SetInterestOps {
private final SelectionKey selectionKey;
private final int interestOps;
private SetInterestOps(SelectionKey selectionKey, int interestOps) {
this.selectionKey = selectionKey;
this.interestOps = interestOps;
}
}
ScheduledAction schedule(Runnable runnable, long delay, TimeUnit unit, ScheduledAction.Kind scheduledActionKind) {
long releaseTimeEpoch = System.currentTimeMillis() + unit.toMillis(delay);
Date releaseTimeDate = new Date(releaseTimeEpoch);
ScheduledAction scheduledAction = new ScheduledAction(runnable, releaseTimeDate, this, scheduledActionKind);
scheduledActions.add(scheduledAction);
selector.wakeup();
return scheduledAction;
}
/**
* Cancels the scheduled action.
*
* @param scheduledAction the scheduled action to cancel.
* @return <code>true</code> if the scheduled action was still pending and got removed, <code>false</code> otherwise.
*/
boolean cancel(ScheduledAction scheduledAction) {
return scheduledActions.remove(scheduledAction);
}
private class Reactor extends Thread {
private volatile long shutdownRequestTimestamp = -1;
@Override
public void run() {
try {
reactorLoop();
} finally {
if (shutdownRequestTimestamp > 0) {
long shutDownDelay = System.currentTimeMillis() - shutdownRequestTimestamp;
LOGGER.info(this + " shut down after " + shutDownDelay + "ms");
} else {
boolean contained = reactorThreads.remove(this);
assert contained;
}
}
}
private void reactorLoop() {
// Loop until reactor shutdown was requested.
while (shutdownRequestTimestamp < 0) {
handleScheduledActionsOrPerformSelect();
handlePendingSelectionKeys();
}
}
@SuppressWarnings("LockNotBeforeTry")
private void handleScheduledActionsOrPerformSelect() {
ScheduledAction dueScheduledAction = null;
boolean permitToHandleScheduledActions = actionsSemaphore.tryAcquire();
if (permitToHandleScheduledActions) {
try {
dueScheduledAction = scheduledActions.poll();
} finally {
actionsSemaphore.release();
}
}
if (dueScheduledAction != null) {
dueScheduledAction.run();
return;
}
int newSelectedKeysCount = 0;
List<SelectionKey> selectedKeys;
synchronized (selector) {
ScheduledAction nextScheduledAction = scheduledActions.peek();
long selectWait;
if (nextScheduledAction == null) {
// There is no next scheduled action, wait indefinitely in select() or until another thread invokes
// selector.wakeup().
selectWait = 0;
} else {
selectWait = nextScheduledAction.getTimeToDueMillis();
if (selectWait <= 0) {
// A scheduled action was just released and became ready to execute.
return;
}
}
// Before we call select, we handle the pending the interest Ops. This will not block since no other
// thread is currently in select() at this time.
// Note: This was put deliberately before the registration lock. It may cause more synchronization but
// allows for more parallelism.
// Hopefully that assumption is right.
int myHandledPendingSetInterestOps = 0;
for (SetInterestOps setInterestOps; (setInterestOps = pendingSetInterestOps.poll()) != null;) {
setInterestOpsCancelledKeySafe(setInterestOps.selectionKey, setInterestOps.interestOps);
if (myHandledPendingSetInterestOps++ >= PENDING_SET_INTEREST_OPS_MAX_BATCH_SIZE) {
// This thread has handled enough "set pending interest ops" requests. Wakeup another one to
// handle the remaining (if any).
selector.wakeup();
break;
}
}
// Ensure that a wakeup() in registerWithSelector() gives the corresponding
// register() in the same method the chance to actually register the channel. In
// other words: This construct ensures that there is never another select()
// between a corresponding wakeup() and register() calls.
// See also https://stackoverflow.com/a/1112809/194894
registrationLock.lock();
registrationLock.unlock();
try {
newSelectedKeysCount = selector.select(selectWait);
} catch (IOException e) {
LOGGER.log(Level.SEVERE, "IOException while using select()", e);
return;
}
if (newSelectedKeysCount == 0) {
return;
}
// Copy the selected-key set over to selectedKeys, remove the keys from the
// selected key set and loose interest of the key OPs for the time being.
// Note that we perform this operation in two steps in order to maximize the
// timespan setRacing() is set.
Set<SelectionKey> selectedKeySet = selector.selectedKeys();
for (SelectionKey selectionKey : selectedKeySet) {
SelectionKeyAttachment selectionKeyAttachment = (SelectionKeyAttachment) selectionKey.attachment();
selectionKeyAttachment.setRacing();
}
for (SelectionKey selectionKey : selectedKeySet) {
setInterestOpsCancelledKeySafe(selectionKey, 0);
}
selectedKeys = new ArrayList<>(selectedKeySet);
selectedKeySet.clear();
}
int selectedKeysCount = selectedKeys.size();
int currentReactorThreadCount = reactorThreads.size();
int myKeyCount;
if (selectedKeysCount > currentReactorThreadCount) {
myKeyCount = selectedKeysCount / currentReactorThreadCount;
} else {
myKeyCount = selectedKeysCount;
}
final Level reactorSelectStatsLogLevel = Level.FINE;
if (LOGGER.isLoggable(reactorSelectStatsLogLevel)) {
LOGGER.log(reactorSelectStatsLogLevel,
"New selected key count: " + newSelectedKeysCount
+ ". Total selected key count " + selectedKeysCount
+ ". My key count: " + myKeyCount
+ ". Current reactor thread count: " + currentReactorThreadCount);
}
Collection<SelectionKey> mySelectedKeys = new ArrayList<>(myKeyCount);
Iterator<SelectionKey> it = selectedKeys.iterator();
for (int i = 0; i < myKeyCount; i++) {
SelectionKey selectionKey = it.next();
mySelectedKeys.add(selectionKey);
}
while (it.hasNext()) {
// Drain to pendingSelectionKeys.
SelectionKey selectionKey = it.next();
pendingSelectionKeys.add(selectionKey);
}
if (selectedKeysCount - myKeyCount > 0) {
// There where pending selection keys: Wakeup another reactor thread to handle them.
selector.wakeup();
}
handleSelectedKeys(mySelectedKeys);
}
private void handlePendingSelectionKeys() {
final int pendingSelectionKeysSize = pendingSelectionKeys.size();
if (pendingSelectionKeysSize == 0) {
return;
}
int currentReactorThreadCount = reactorThreads.size();
int myKeyCount = pendingSelectionKeysSize / currentReactorThreadCount;
// The division could result in myKeyCount being zero, even though there are pending selection keys.
// Therefore, ensure that this thread tries to get at least one pending selection key by invoking poll().
// Otherwise, it could happen that we end up in a busy loop, where myKeyCount is zero and this thread invokes
// selector.wakeup() below because pendingSelectionsKeys is not empty, but the woken up reactor thread wil
// end up with myKeyCount being zero again, restarting the busy-loop cycle.
if (myKeyCount == 0) myKeyCount = 1;
Collection<SelectionKey> selectedKeys = new ArrayList<>(myKeyCount);
for (int i = 0; i < myKeyCount; i++) {
SelectionKey selectionKey = pendingSelectionKeys.poll();
if (selectionKey == null) {
// We lost a race and can abort here since the pendingSelectionKeys queue is empty.
break;
}
selectedKeys.add(selectionKey);
}
if (!pendingSelectionKeys.isEmpty()) {
// There are more pending selection keys, wakeup a thread blocked in select() to handle them.
selector.wakeup();
}
handleSelectedKeys(selectedKeys);
}
private void setInterestOpsCancelledKeySafe(SelectionKey selectionKey, int interestOps) {
try {
selectionKey.interestOps(interestOps);
}
catch (CancelledKeyException e) {
final Level keyCancelledLogLevel = Level.FINER;
if (LOGGER.isLoggable(keyCancelledLogLevel)) {
LOGGER.log(keyCancelledLogLevel, "Key '" + selectionKey + "' has been cancelled", e);
}
}
}
void requestShutdown() {
shutdownRequestTimestamp = System.currentTimeMillis();
}
}
private static void handleSelectedKeys(Collection<SelectionKey> selectedKeys) {
for (SelectionKey selectionKey : selectedKeys) {
SelectableChannel channel = selectionKey.channel();
SelectionKeyAttachment selectionKeyAttachment = (SelectionKeyAttachment) selectionKey.attachment();
ChannelSelectedCallback channelSelectedCallback = selectionKeyAttachment.channelSelectedCallback;
channelSelectedCallback.onChannelSelected(channel, selectionKey);
}
}
public interface ChannelSelectedCallback {
void onChannelSelected(SelectableChannel channel, SelectionKey selectionKey);
}
public void setReactorThreadCount(int reactorThreadCount) {
if (reactorThreadCount < 2) {
throw new IllegalArgumentException("Must have at least two reactor threads, but you requested " + reactorThreadCount);
}
synchronized (reactorThreads) {
int deltaThreads = reactorThreadCount - reactorThreads.size();
if (deltaThreads > 0) {
// Start new reactor thread. Note that we start the threads before we increase the permits of the
// actionsSemaphore.
for (int i = 0; i < deltaThreads; i++) {
Reactor reactor = new Reactor();
reactor.setDaemon(true);
reactor.setName("Smack " + reactorName + " Thread #" + i);
reactorThreads.add(reactor);
reactor.start();
}
actionsSemaphore.release(deltaThreads);
} else {
// Stop existing reactor threads. First we change the sign of deltaThreads, then we decrease the permits
// of the actionsSemaphore *before* we signal the selected reactor threads that they should shut down.
deltaThreads -= deltaThreads;
for (int i = deltaThreads - 1; i > 0; i--) {
// Note that this could potentially block forever, starving on the unfair semaphore.
actionsSemaphore.acquireUninterruptibly();
}
for (int i = deltaThreads - 1; i > 0; i--) {
Reactor reactor = reactorThreads.remove(i);
reactor.requestShutdown();
}
selector.wakeup();
}
}
}
public static final class SelectionKeyAttachment {
private final ChannelSelectedCallback channelSelectedCallback;
private final AtomicBoolean reactorThreadRacing = new AtomicBoolean();
private SelectionKeyAttachment(ChannelSelectedCallback channelSelectedCallback) {
this.channelSelectedCallback = channelSelectedCallback;
}
private void setRacing() {
// We use lazySet here since it is sufficient if the value does not become visible immediately.
reactorThreadRacing.lazySet(true);
}
public void resetReactorThreadRacing() {
reactorThreadRacing.set(false);
}
public boolean isReactorThreadRacing() {
return reactorThreadRacing.get();
}
}
}