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