001/**
002 *
003 * Copyright 2018-2023 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);
280                selectedKeySet.clear();
281            }
282
283            int selectedKeysCount = selectedKeys.size();
284            int currentReactorThreadCount = reactorThreads.size();
285            int myKeyCount;
286            if (selectedKeysCount > currentReactorThreadCount) {
287                myKeyCount = selectedKeysCount / currentReactorThreadCount;
288            } else {
289                myKeyCount = selectedKeysCount;
290            }
291
292            final Level reactorSelectStatsLogLevel = Level.FINE;
293            if (LOGGER.isLoggable(reactorSelectStatsLogLevel)) {
294                LOGGER.log(reactorSelectStatsLogLevel,
295                                "New selected key count: " + newSelectedKeysCount
296                                + ". Total selected key count " + selectedKeysCount
297                                + ". My key count: " + myKeyCount
298                                + ". Current reactor thread count: " + currentReactorThreadCount);
299            }
300
301            Collection<SelectionKey> mySelectedKeys = new ArrayList<>(myKeyCount);
302            Iterator<SelectionKey> it = selectedKeys.iterator();
303            for (int i = 0; i < myKeyCount; i++) {
304                SelectionKey selectionKey = it.next();
305                mySelectedKeys.add(selectionKey);
306            }
307            while (it.hasNext()) {
308                // Drain to pendingSelectionKeys.
309                SelectionKey selectionKey = it.next();
310                pendingSelectionKeys.add(selectionKey);
311            }
312
313            if (selectedKeysCount - myKeyCount > 0) {
314                // There where pending selection keys: Wakeup another reactor thread to handle them.
315                selector.wakeup();
316            }
317
318            handleSelectedKeys(mySelectedKeys);
319        }
320
321        private void handlePendingSelectionKeys() {
322            final int pendingSelectionKeysSize = pendingSelectionKeys.size();
323            if (pendingSelectionKeysSize == 0) {
324                return;
325            }
326
327            int currentReactorThreadCount = reactorThreads.size();
328            int myKeyCount = pendingSelectionKeysSize / currentReactorThreadCount;
329            // The division could result in myKeyCount being zero, even though there are pending selection keys.
330            // Therefore, ensure that this thread tries to get at least one pending selection key by invoking poll().
331            // Otherwise, it could happen that we end up in a busy loop, where myKeyCount is zero and this thread invokes
332            // selector.wakeup() below because pendingSelectionsKeys is not empty, but the woken up reactor thread wil
333            // end up with myKeyCount being zero again, restarting the busy-loop cycle.
334            if (myKeyCount == 0) myKeyCount = 1;
335            Collection<SelectionKey> selectedKeys = new ArrayList<>(myKeyCount);
336            for (int i = 0; i < myKeyCount; i++) {
337                SelectionKey selectionKey = pendingSelectionKeys.poll();
338                if (selectionKey == null) {
339                    // We lost a race and can abort here since the pendingSelectionKeys queue is empty.
340                    break;
341                }
342                selectedKeys.add(selectionKey);
343            }
344
345            if (!pendingSelectionKeys.isEmpty()) {
346                // There are more pending selection keys, wakeup a thread blocked in select() to handle them.
347                selector.wakeup();
348            }
349
350            handleSelectedKeys(selectedKeys);
351        }
352
353        private void setInterestOpsCancelledKeySafe(SelectionKey selectionKey, int interestOps) {
354            try {
355                selectionKey.interestOps(interestOps);
356            }
357            catch (CancelledKeyException e) {
358                final Level keyCancelledLogLevel = Level.FINER;
359                if (LOGGER.isLoggable(keyCancelledLogLevel)) {
360                    LOGGER.log(keyCancelledLogLevel, "Key '" + selectionKey + "' has been cancelled", e);
361                }
362            }
363        }
364
365        void requestShutdown() {
366            shutdownRequestTimestamp = System.currentTimeMillis();
367        }
368    }
369
370    private static void handleSelectedKeys(Collection<SelectionKey> selectedKeys) {
371        for (SelectionKey selectionKey : selectedKeys) {
372            SelectableChannel channel = selectionKey.channel();
373            SelectionKeyAttachment selectionKeyAttachment = (SelectionKeyAttachment) selectionKey.attachment();
374            ChannelSelectedCallback channelSelectedCallback = selectionKeyAttachment.channelSelectedCallback;
375            channelSelectedCallback.onChannelSelected(channel, selectionKey);
376        }
377    }
378
379    public interface ChannelSelectedCallback {
380        void onChannelSelected(SelectableChannel channel, SelectionKey selectionKey);
381    }
382
383    public void setReactorThreadCount(int reactorThreadCount) {
384        if (reactorThreadCount < 2) {
385            throw new IllegalArgumentException("Must have at least two reactor threads, but you requested " + reactorThreadCount);
386        }
387
388        synchronized (reactorThreads) {
389            int deltaThreads = reactorThreadCount - reactorThreads.size();
390            if (deltaThreads > 0) {
391                // Start new reactor thread. Note that we start the threads before we increase the permits of the
392                // actionsSemaphore.
393                for (int i = 0; i < deltaThreads; i++) {
394                    Reactor reactor = new Reactor();
395                    reactor.setDaemon(true);
396                    reactor.setName("Smack " + reactorName + " Thread #" + i);
397                    reactorThreads.add(reactor);
398                    reactor.start();
399                }
400
401                actionsSemaphore.release(deltaThreads);
402            } else {
403                // Stop existing reactor threads. First we change the sign of deltaThreads, then we decrease the permits
404                // of the actionsSemaphore *before* we signal the selected reactor threads that they should shut down.
405                deltaThreads -= deltaThreads;
406
407                for (int i = deltaThreads - 1; i > 0; i--) {
408                    // Note that this could potentially block forever, starving on the unfair semaphore.
409                    actionsSemaphore.acquireUninterruptibly();
410                }
411
412                for (int i = deltaThreads - 1; i > 0; i--) {
413                    Reactor reactor = reactorThreads.remove(i);
414                    reactor.requestShutdown();
415                }
416
417                selector.wakeup();
418            }
419        }
420    }
421
422    public static final class SelectionKeyAttachment {
423        private final ChannelSelectedCallback channelSelectedCallback;
424        private final AtomicBoolean reactorThreadRacing = new AtomicBoolean();
425
426        private SelectionKeyAttachment(ChannelSelectedCallback channelSelectedCallback) {
427            this.channelSelectedCallback = channelSelectedCallback;
428        }
429
430        private void setRacing() {
431            // We use lazySet here since it is sufficient if the value does not become visible immediately.
432            reactorThreadRacing.lazySet(true);
433        }
434
435        public void resetReactorThreadRacing() {
436            reactorThreadRacing.set(false);
437        }
438
439        public boolean isReactorThreadRacing() {
440            return reactorThreadRacing.get();
441        }
442
443    }
444}