001/**
002 *
003 * Copyright 2018-2020 Florian Schmaus
004 *
005 * Licensed under the Apache License, Version 2.0 (the "License");
006 * you may not use this file except in compliance with the License.
007 * You may obtain a copy of the License at
008 *
009 *     http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.jivesoftware.smack;
018
019import java.io.IOException;
020import java.nio.channels.CancelledKeyException;
021import java.nio.channels.ClosedChannelException;
022import java.nio.channels.SelectableChannel;
023import java.nio.channels.SelectionKey;
024import java.nio.channels.Selector;
025import java.util.ArrayList;
026import java.util.Collection;
027import java.util.Collections;
028import java.util.Date;
029import java.util.Iterator;
030import java.util.List;
031import java.util.Queue;
032import java.util.Set;
033import java.util.concurrent.ConcurrentLinkedQueue;
034import java.util.concurrent.DelayQueue;
035import java.util.concurrent.Semaphore;
036import java.util.concurrent.TimeUnit;
037import java.util.concurrent.atomic.AtomicBoolean;
038import java.util.concurrent.locks.Lock;
039import java.util.concurrent.locks.ReentrantLock;
040import java.util.logging.Level;
041import java.util.logging.Logger;
042
043/**
044 * The SmackReactor for non-blocking I/O.
045 * <p>
046 * Highlights include:
047 * <ul>
048 * <li>Multiple reactor threads</li>
049 * <li>Scheduled actions</li>
050 * </ul>
051 *
052 * <pre>
053 *
054 *           ) ) )
055 *        ( ( (
056 *      ) ) )
057 *   (~~~~~~~~~)
058 *    | Smack |
059 *    |Reactor|
060 *    I      _._
061 *    I    /'   `\
062 *    I   |       |
063 *    f   |   |~~~~~~~~~~~~~~|
064 *  .'    |   | #   #   #  # |
065 * '______|___|___________###|
066 * </pre>
067 */
068public class SmackReactor {
069
070    private static final Logger LOGGER = Logger.getLogger(SmackReactor.class.getName());
071
072    private static final int DEFAULT_REACTOR_THREAD_COUNT = 2;
073
074    private static final int PENDING_SET_INTEREST_OPS_MAX_BATCH_SIZE = 1024;
075
076    private static SmackReactor INSTANCE;
077
078    static synchronized SmackReactor getInstance() {
079        if (INSTANCE == null) {
080            INSTANCE = new SmackReactor("DefaultReactor");
081        }
082        return INSTANCE;
083    }
084
085    private final Selector selector;
086    private final String reactorName;
087
088    private final List<Reactor> reactorThreads = Collections.synchronizedList(new ArrayList<>());
089
090    private final DelayQueue<ScheduledAction> scheduledActions = new DelayQueue<>();
091
092    private final Lock registrationLock = new ReentrantLock();
093
094    /**
095     * The semaphore protecting the handling of the actions. Note that it is
096     * initialized with -1, which basically means that one thread will always do I/O using
097     * select().
098     */
099    private final Semaphore actionsSemaphore = new Semaphore(-1, false);
100
101    private final Queue<SelectionKey> pendingSelectionKeys = new ConcurrentLinkedQueue<>();
102
103    private final Queue<SetInterestOps> pendingSetInterestOps = new ConcurrentLinkedQueue<>();
104
105    SmackReactor(String reactorName) {
106        this.reactorName = reactorName;
107
108        try {
109            selector = Selector.open();
110        }
111        catch (IOException e) {
112            throw new IllegalStateException(e);
113        }
114
115        setReactorThreadCount(DEFAULT_REACTOR_THREAD_COUNT);
116    }
117
118    public SelectionKey registerWithSelector(SelectableChannel channel, int ops, ChannelSelectedCallback callback)
119            throws ClosedChannelException {
120        SelectionKeyAttachment selectionKeyAttachment = new SelectionKeyAttachment(callback);
121
122        registrationLock.lock();
123        try {
124            selector.wakeup();
125            return channel.register(selector, ops, selectionKeyAttachment);
126        } finally {
127            registrationLock.unlock();
128        }
129    }
130
131    public void setInterestOps(SelectionKey selectionKey, int interestOps) {
132        SetInterestOps setInterestOps = new SetInterestOps(selectionKey, interestOps);
133        pendingSetInterestOps.add(setInterestOps);
134        selector.wakeup();
135    }
136
137    private static final class SetInterestOps {
138        private final SelectionKey selectionKey;
139        private final int interestOps;
140
141        private SetInterestOps(SelectionKey selectionKey, int interestOps) {
142            this.selectionKey = selectionKey;
143            this.interestOps = interestOps;
144        }
145    }
146
147    ScheduledAction schedule(Runnable runnable, long delay, TimeUnit unit, ScheduledAction.Kind scheduledActionKind) {
148        long releaseTimeEpoch = System.currentTimeMillis() + unit.toMillis(delay);
149        Date releaseTimeDate = new Date(releaseTimeEpoch);
150        ScheduledAction scheduledAction = new ScheduledAction(runnable, releaseTimeDate, this, scheduledActionKind);
151        scheduledActions.add(scheduledAction);
152        selector.wakeup();
153        return scheduledAction;
154    }
155
156    /**
157     * Cancels the scheduled action.
158     *
159     * @param scheduledAction the scheduled action to cancel.
160     * @return <code>true</code> if the scheduled action was still pending and got removed, <code>false</code> otherwise.
161     */
162    boolean cancel(ScheduledAction scheduledAction) {
163        return scheduledActions.remove(scheduledAction);
164    }
165
166    private class Reactor extends Thread {
167
168        private volatile long shutdownRequestTimestamp = -1;
169
170        @Override
171        public void run() {
172            try {
173                reactorLoop();
174            } finally {
175                if (shutdownRequestTimestamp > 0) {
176                    long shutDownDelay = System.currentTimeMillis() - shutdownRequestTimestamp;
177                    LOGGER.info(this + " shut down after " + shutDownDelay + "ms");
178                } else {
179                    boolean contained = reactorThreads.remove(this);
180                    assert contained;
181                }
182            }
183        }
184
185        private void reactorLoop() {
186            // Loop until reactor shutdown was requested.
187            while (shutdownRequestTimestamp < 0) {
188                handleScheduledActionsOrPerformSelect();
189
190                handlePendingSelectionKeys();
191            }
192        }
193
194        @SuppressWarnings("LockNotBeforeTry")
195        private void handleScheduledActionsOrPerformSelect() {
196            ScheduledAction dueScheduledAction = null;
197
198            boolean permitToHandleScheduledActions = actionsSemaphore.tryAcquire();
199            if (permitToHandleScheduledActions) {
200                try {
201                    dueScheduledAction = scheduledActions.poll();
202                } finally {
203                    actionsSemaphore.release();
204                }
205            }
206
207            if (dueScheduledAction != null) {
208                dueScheduledAction.run();
209                return;
210            }
211
212            int newSelectedKeysCount = 0;
213            List<SelectionKey> selectedKeys;
214            synchronized (selector) {
215                ScheduledAction nextScheduledAction = scheduledActions.peek();
216
217                long selectWait;
218                if (nextScheduledAction == null) {
219                    // There is no next scheduled action, wait indefinitely in select() or until another thread invokes
220                    // selector.wakeup().
221                    selectWait = 0;
222                } else {
223                    selectWait = nextScheduledAction.getTimeToDueMillis();
224                    if (selectWait <= 0) {
225                        // A scheduled action was just released and became ready to execute.
226                        return;
227                    }
228                }
229
230                // Before we call select, we handle the pending the interest Ops. This will not block since no other
231                // thread is currently in select() at this time.
232                // Note: This was put deliberately before the registration lock. It may cause more synchronization but
233                // allows for more parallelism.
234                // Hopefully that assumption is right.
235                int myHandledPendingSetInterestOps = 0;
236                for (SetInterestOps setInterestOps; (setInterestOps = pendingSetInterestOps.poll()) != null;) {
237                    setInterestOpsCancelledKeySafe(setInterestOps.selectionKey, setInterestOps.interestOps);
238
239                    if (myHandledPendingSetInterestOps++ >= PENDING_SET_INTEREST_OPS_MAX_BATCH_SIZE) {
240                        // This thread has handled enough "set pending interest ops" requests. Wakeup another one to
241                        // handle the remaining (if any).
242                        selector.wakeup();
243                        break;
244                    }
245                }
246
247                // Ensure that a wakeup() in registerWithSelector() gives the corresponding
248                // register() in the same method the chance to actually register the channel. In
249                // other words: This construct ensures that there is never another select()
250                // between a corresponding wakeup() and register() calls.
251                // See also https://stackoverflow.com/a/1112809/194894
252                registrationLock.lock();
253                registrationLock.unlock();
254
255                try {
256                    newSelectedKeysCount = selector.select(selectWait);
257                } catch (IOException e) {
258                    LOGGER.log(Level.SEVERE, "IOException while using select()", e);
259                    return;
260                }
261
262                if (newSelectedKeysCount == 0) {
263                    return;
264                }
265
266                // Copy the selected-key set over to selectedKeys, remove the keys from the
267                // selected key set and loose interest of the key OPs for the time being.
268                // Note that we perform this operation in two steps in order to maximize the
269                // timespan setRacing() is set.
270                Set<SelectionKey> selectedKeySet = selector.selectedKeys();
271                for (SelectionKey selectionKey : selectedKeySet) {
272                    SelectionKeyAttachment selectionKeyAttachment = (SelectionKeyAttachment) selectionKey.attachment();
273                    selectionKeyAttachment.setRacing();
274                }
275                for (SelectionKey selectionKey : selectedKeySet) {
276                    setInterestOpsCancelledKeySafe(selectionKey, 0);
277                }
278
279                selectedKeys = new ArrayList<>(selectedKeySet.size());
280                selectedKeys.addAll(selectedKeySet);
281                selectedKeySet.clear();
282            }
283
284            int selectedKeysCount = selectedKeys.size();
285            int currentReactorThreadCount = reactorThreads.size();
286            int myKeyCount;
287            if (selectedKeysCount > currentReactorThreadCount) {
288                myKeyCount = selectedKeysCount / currentReactorThreadCount;
289            } else {
290                myKeyCount = selectedKeysCount;
291            }
292
293            final Level reactorSelectStatsLogLevel = Level.FINE;
294            if (LOGGER.isLoggable(reactorSelectStatsLogLevel)) {
295                LOGGER.log(reactorSelectStatsLogLevel,
296                                "New selected key count: " + newSelectedKeysCount
297                                + ". Total selected key count " + selectedKeysCount
298                                + ". My key count: " + myKeyCount
299                                + ". Current reactor thread count: " + currentReactorThreadCount);
300            }
301
302            Collection<SelectionKey> mySelectedKeys = new ArrayList<>(myKeyCount);
303            Iterator<SelectionKey> it = selectedKeys.iterator();
304            for (int i = 0; i < myKeyCount; i++) {
305                SelectionKey selectionKey = it.next();
306                mySelectedKeys.add(selectionKey);
307            }
308            while (it.hasNext()) {
309                // Drain to pendingSelectionKeys.
310                SelectionKey selectionKey = it.next();
311                pendingSelectionKeys.add(selectionKey);
312            }
313
314            if (selectedKeysCount - myKeyCount > 0) {
315                // There where pending selection keys: Wakeup another reactor thread to handle them.
316                selector.wakeup();
317            }
318
319            handleSelectedKeys(mySelectedKeys);
320        }
321
322        private void handlePendingSelectionKeys() {
323            final int pendingSelectionKeysSize = pendingSelectionKeys.size();
324            if (pendingSelectionKeysSize == 0) {
325                return;
326            }
327
328            int currentReactorThreadCount = reactorThreads.size();
329            int myKeyCount = pendingSelectionKeysSize / currentReactorThreadCount;
330            Collection<SelectionKey> selectedKeys = new ArrayList<>(myKeyCount);
331            for (int i = 0; i < myKeyCount; i++) {
332                SelectionKey selectionKey = pendingSelectionKeys.poll();
333                if (selectionKey == null) {
334                    // We lost a race and can abort here since the pendingSelectionKeys queue is empty.
335                    break;
336                }
337                selectedKeys.add(selectionKey);
338            }
339
340            if (!pendingSelectionKeys.isEmpty()) {
341                // There are more pending selection keys, wakeup a thread blocked in select() to handle them.
342                selector.wakeup();
343            }
344
345            handleSelectedKeys(selectedKeys);
346        }
347
348        private void setInterestOpsCancelledKeySafe(SelectionKey selectionKey, int interestOps) {
349            try {
350                selectionKey.interestOps(interestOps);
351            }
352            catch (CancelledKeyException e) {
353                final Level keyCancelledLogLevel = Level.FINER;
354                if (LOGGER.isLoggable(keyCancelledLogLevel)) {
355                    LOGGER.log(keyCancelledLogLevel, "Key '" + selectionKey + "' has been cancelled", e);
356                }
357            }
358        }
359
360        void requestShutdown() {
361            shutdownRequestTimestamp = System.currentTimeMillis();
362        }
363    }
364
365    private static void handleSelectedKeys(Collection<SelectionKey> selectedKeys) {
366        for (SelectionKey selectionKey : selectedKeys) {
367            SelectableChannel channel = selectionKey.channel();
368            SelectionKeyAttachment selectionKeyAttachment = (SelectionKeyAttachment) selectionKey.attachment();
369            ChannelSelectedCallback channelSelectedCallback = selectionKeyAttachment.channelSelectedCallback;
370            channelSelectedCallback.onChannelSelected(channel, selectionKey);
371        }
372    }
373
374    public interface ChannelSelectedCallback {
375        void onChannelSelected(SelectableChannel channel, SelectionKey selectionKey);
376    }
377
378    public void setReactorThreadCount(int reactorThreadCount) {
379        if (reactorThreadCount < 2) {
380            throw new IllegalArgumentException("Must have at least two reactor threads, but you requested " + reactorThreadCount);
381        }
382
383        synchronized (reactorThreads) {
384            int deltaThreads = reactorThreadCount - reactorThreads.size();
385            if (deltaThreads > 0) {
386                // Start new reactor thread. Note that we start the threads before we increase the permits of the
387                // actionsSemaphore.
388                for (int i = 0; i < deltaThreads; i++) {
389                    Reactor reactor = new Reactor();
390                    reactor.setDaemon(true);
391                    reactor.setName("Smack " + reactorName + " Thread #" + i);
392                    reactorThreads.add(reactor);
393                    reactor.start();
394                }
395
396                actionsSemaphore.release(deltaThreads);
397            } else {
398                // Stop existing reactor threads. First we change the sign of deltaThreads, then we decrease the permits
399                // of the actionsSemaphore *before* we signal the selected reactor threads that they should shut down.
400                deltaThreads -= deltaThreads;
401
402                for (int i = deltaThreads - 1; i > 0; i--) {
403                    // Note that this could potentially block forever, starving on the unfair semaphore.
404                    actionsSemaphore.acquireUninterruptibly();
405                }
406
407                for (int i = deltaThreads - 1; i > 0; i--) {
408                    Reactor reactor = reactorThreads.remove(i);
409                    reactor.requestShutdown();
410                }
411
412                selector.wakeup();
413            }
414        }
415    }
416
417    public static final class SelectionKeyAttachment {
418        private final ChannelSelectedCallback channelSelectedCallback;
419        private final AtomicBoolean reactorThreadRacing = new AtomicBoolean();
420
421        private SelectionKeyAttachment(ChannelSelectedCallback channelSelectedCallback) {
422            this.channelSelectedCallback = channelSelectedCallback;
423        }
424
425        private void setRacing() {
426            // We use lazySet here since it is sufficient if the value does not become visible immediately.
427            reactorThreadRacing.lazySet(true);
428        }
429
430        public void resetReactorThreadRacing() {
431            reactorThreadRacing.set(false);
432        }
433
434        public boolean isReactorThreadRacing() {
435            return reactorThreadRacing.get();
436        }
437
438    }
439}