org.jboss.netty.handler.execution.MemoryAwareThreadPoolExecutor - Netty 3.10.5.Final 源码

/*
 * Copyright 2012 The Netty Project
 *
 * The Netty Project licenses this file to you 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.jboss.netty.handler.execution;

import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.channel.Channel;
import org.jboss.netty.channel.ChannelEvent;
import org.jboss.netty.channel.ChannelFuture;
import org.jboss.netty.channel.ChannelHandlerContext;
import org.jboss.netty.channel.ChannelState;
import org.jboss.netty.channel.ChannelStateEvent;
import org.jboss.netty.channel.Channels;
import org.jboss.netty.channel.MessageEvent;
import org.jboss.netty.channel.WriteCompletionEvent;
import org.jboss.netty.logging.InternalLogger;
import org.jboss.netty.logging.InternalLoggerFactory;
import org.jboss.netty.util.DefaultObjectSizeEstimator;
import org.jboss.netty.util.ObjectSizeEstimator;
import org.jboss.netty.util.internal.ConcurrentIdentityHashMap;
import org.jboss.netty.util.internal.SharedResourceMisuseDetector;

import java.io.IOException;
import java.lang.reflect.Method;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * A {@link ThreadPoolExecutor} which blocks the task submission when there's
 * too many tasks in the queue.  Both per-{@link Channel} and per-{@link Executor}
 * limitation can be applied.
 * <p>
 * When a task (i.e. {@link Runnable}) is submitted,
 * {@link MemoryAwareThreadPoolExecutor} calls {@link ObjectSizeEstimator#estimateSize(Object)}
 * to get the estimated size of the task in bytes to calculate the amount of
 * memory occupied by the unprocessed tasks.
 * <p>
 * If the total size of the unprocessed tasks exceeds either per-{@link Channel}
 * or per-{@link Executor} threshold, any further {@link #execute(Runnable)}
 * call will block until the tasks in the queue are processed so that the total
 * size goes under the threshold.
 *
 * <h3>Using an alternative task size estimation strategy</h3>
 *
 * Although the default implementation does its best to guess the size of an
 * object of unknown type, it is always good idea to to use an alternative
 * {@link ObjectSizeEstimator} implementation instead of the
 * {@link DefaultObjectSizeEstimator} to avoid incorrect task size calculation,
 * especially when:
 * <ul>
 *   <li>you are using {@link MemoryAwareThreadPoolExecutor} independently from
 *       {@link ExecutionHandler},</li>
 *   <li>you are submitting a task whose type is not {@link ChannelEventRunnable}, or</li>
 *   <li>the message type of the {@link MessageEvent} in the {@link ChannelEventRunnable}
 *       is not {@link ChannelBuffer}.</li>
 * </ul>
 * Here is an example that demonstrates how to implement an {@link ObjectSizeEstimator}
 * which understands a user-defined object:
 * <pre>
 * public class MyRunnable implements {@link Runnable} {
 *
 *     <b>private final byte[] data;</b>
 *
 *     public MyRunnable(byte[] data) {
 *         this.data = data;
 *     }
 *
 *     public void run() {
 *         // Process 'data' ..
 *     }
 * }
 *
 * public class MyObjectSizeEstimator extends {@link DefaultObjectSizeEstimator} {
 *
 *     {@literal @Override}
 *     public int estimateSize(Object o) {
 *         if (<b>o instanceof MyRunnable</b>) {
 *             <b>return ((MyRunnable) o).data.length + 8;</b>
 *         }
 *         return super.estimateSize(o);
 *     }
 * }
 *
 * {@link ThreadPoolExecutor} pool = new {@link MemoryAwareThreadPoolExecutor}(
 *         16, 65536, 1048576, 30, {@link TimeUnit}.SECONDS,
 *         <b>new MyObjectSizeEstimator()</b>,
 *         {@link Executors}.defaultThreadFactory());
 *
 * <b>pool.execute(new MyRunnable(data));</b>
 * </pre>
 *
 * <h3>Event execution order</h3>
 *
 * Please note that this executor does not maintain the order of the
 * {@link ChannelEvent}s for the same {@link Channel}.  For example,
 * you can even receive a {@code "channelClosed"} event before a
 * {@code "messageReceived"} event, as depicted by the following diagram.
 *
 * For example, the events can be processed as depicted below:
 *
 * <pre>
 *           --------------------------------&gt; Timeline --------------------------------&gt;
 *
 * Thread X: --- Channel A (Event 1) --- Channel A (Event 2) ---------------------------&gt;
 *
 * Thread Y: --- Channel A (Event 3) --- Channel B (Event 2) --- Channel B (Event 3) ---&gt;
 *
 * Thread Z: --- Channel B (Event 1) --- Channel B (Event 4) --- Channel A (Event 4) ---&gt;
 * </pre>
 *
 * To maintain the event order, you must use {@link OrderedMemoryAwareThreadPoolExecutor}.
 *
 * @apiviz.has org.jboss.netty.util.ObjectSizeEstimator oneway - -
 * @apiviz.has org.jboss.netty.handler.execution.ChannelEventRunnable oneway - - executes
 */
public class MemoryAwareThreadPoolExecutor extends ThreadPoolExecutor {

    private static final InternalLogger logger =
        InternalLoggerFactory.getInstance(MemoryAwareThreadPoolExecutor.class);

    private static final SharedResourceMisuseDetector misuseDetector =
        new SharedResourceMisuseDetector(MemoryAwareThreadPoolExecutor.class);

    private volatile Settings settings;

    private final ConcurrentMap<Channel, AtomicLong> channelCounters =
        new ConcurrentIdentityHashMap<Channel, AtomicLong>();
    private final Limiter totalLimiter;

    private volatile boolean notifyOnShutdown;

    /**
     * Creates a new instance.
     *
     * @param corePoolSize          the maximum number of active threads
     * @param maxChannelMemorySize  the maximum total size of the queued events per channel.
     *                              Specify {@code 0} to disable.
     * @param maxTotalMemorySize    the maximum total size of the queued events for this pool
     *                              Specify {@code 0} to disable.
     */
    public MemoryAwareThreadPoolExecutor(
            int corePoolSize, long maxChannelMemorySize, long maxTotalMemorySize) {

        this(corePoolSize, maxChannelMemorySize, maxTotalMemorySize, 30, TimeUnit.SECONDS);
    }

    /**
     * Creates a new instance.
     *
     * @param corePoolSize          the maximum number of active threads
     * @param maxChannelMemorySize  the maximum total size of the queued events per channel.
     *                              Specify {@code 0} to disable.
     * @param maxTotalMemorySize    the maximum total size of the queued events for this pool
     *                              Specify {@code 0} to disable.
     * @param keepAliveTime         the amount of time for an inactive thread to shut itself down
     * @param unit                  the {@link TimeUnit} of {@code keepAliveTime}
     */
    public MemoryAwareThreadPoolExecutor(
            int corePoolSize, long maxChannelMemorySize, long maxTotalMemorySize,
            long keepAliveTime, TimeUnit unit) {

        this(
                corePoolSize, maxChannelMemorySize, maxTotalMemorySize, keepAliveTime, unit,
                Executors.defaultThreadFactory());
    }

    /**
     * Creates a new instance.
     *
     * @param corePoolSize          the maximum number of active threads
     * @param maxChannelMemorySize  the maximum total size of the queued events per channel.
     *                              Specify {@code 0} to disable.
     * @param maxTotalMemorySize    the maximum total size of the queued events for this pool
     *                              Specify {@code 0} to disable.
     * @param keepAliveTime         the amount of time for an inactive thread to shut itself down
     * @param unit                  the {@link TimeUnit} of {@code keepAliveTime}
     * @param threadFactory         the {@link ThreadFactory} of this pool
     */
    public MemoryAwareThreadPoolExecutor(
            int corePoolSize, long maxChannelMemorySize, long maxTotalMemorySize,
            long keepAliveTime, TimeUnit unit, ThreadFactory threadFactory) {

        this(
                corePoolSize, maxChannelMemorySize, maxTotalMemorySize, keepAliveTime, unit,
                new DefaultObjectSizeEstimator(), threadFactory);
    }

    /**
     * Creates a new instance.
     *
     * @param corePoolSize          the maximum number of active threads
     * @param maxChannelMemorySize  the maximum total size of the queued events per channel.
     *                              Specify {@code 0} to disable.
     * @param maxTotalMemorySize    the maximum total size of the queued events for this pool
     *                              Specify {@code 0} to disable.
     * @param keepAliveTime         the amount of time for an inactive thread to shut itself down
     * @param unit                  the {@link TimeUnit} of {@code keepAliveTime}
     * @param threadFactory         the {@link ThreadFactory} of this pool
     * @param objectSizeEstimator   the {@link ObjectSizeEstimator} of this pool
     */
    public MemoryAwareThreadPoolExecutor(
            int corePoolSize, long maxChannelMemorySize, long maxTotalMemorySize,
            long keepAliveTime, TimeUnit unit, ObjectSizeEstimator objectSizeEstimator,
            ThreadFactory threadFactory) {

        super(corePoolSize, corePoolSize, keepAliveTime, unit,
              new LinkedBlockingQueue<Runnable>(), threadFactory, new NewThreadRunsPolicy());

        if (objectSizeEstimator == null) {
            throw new NullPointerException("objectSizeEstimator");
        }
        if (maxChannelMemorySize < 0) {
            throw new IllegalArgumentException(
                    "maxChannelMemorySize: " + maxChannelMemorySize);
        }
        if (maxTotalMemorySize < 0) {
            throw new IllegalArgumentException(
                    "maxTotalMemorySize: " + maxTotalMemorySize);
        }

        // Call allowCoreThreadTimeOut(true) using reflection
        // because it is not supported in Java 5.
        try {
            Method m = getClass().getMethod("allowCoreThreadTimeOut", new Class[] { boolean.class });
            m.invoke(this, Boolean.TRUE);
        } catch (Throwable t) {
            // Java 5
            logger.debug(
                    "ThreadPoolExecutor.allowCoreThreadTimeOut() is not " +
                    "supported in this platform.");
        }

        settings = new Settings(
                objectSizeEstimator, maxChannelMemorySize);

        if (maxTotalMemorySize == 0) {
            totalLimiter = null;
        } else {
            totalLimiter = new Limiter(maxTotalMemorySize);
        }

        // Misuse check
        misuseDetector.increase();
    }

    @Override
    protected void terminated() {
        super.terminated();
        misuseDetector.decrease();
    }

    /**
     * This will call {@link #shutdownNow(boolean)} with the value of {@link #getNotifyChannelFuturesOnShutdown()}.
     */
    @Override
    public List<Runnable> shutdownNow() {
        return shutdownNow(notifyOnShutdown);
    }

    /**
     * See {@link ThreadPoolExecutor#shutdownNow()} for how it handles the shutdown.
     * If {@code true} is given to this method it also notifies all {@link ChannelFuture}'s
     * of the not executed {@link ChannelEventRunnable}'s.
     *
     * <p>
     * Be aware that if you call this with {@code false} you will need to handle the
     * notification of the {@link ChannelFuture}'s by your self. So only use this if you
     * really have a use-case for it.
     * </p>
     *
     */
    public List<Runnable> shutdownNow(boolean notify) {
        if (!notify) {
            return super.shutdownNow();
        }
        Throwable cause = null;
        Set<Channel> channels = null;

        List<Runnable> tasks = super.shutdownNow();

        // loop over all tasks and cancel the ChannelFuture of the ChannelEventRunable's
        for (Runnable task: tasks) {
            if (task instanceof ChannelEventRunnable) {
                if (cause == null) {
                    cause = new IOException("Unable to process queued event");
                }
                ChannelEvent event = ((ChannelEventRunnable) task).getEvent();
                event.getFuture().setFailure(cause);

                if (channels == null) {
                    channels = new HashSet<Channel>();
                }

                // store the Channel of the event for later notification of the exceptionCaught event
                channels.add(event.getChannel());
            }
        }

        // loop over all channels and fire an exceptionCaught event
        if (channels != null) {
            for (Channel channel: channels) {
                Channels.fireExceptionCaughtLater(channel, cause);
            }
        }
        return tasks;
    }

    /**
     * Returns the {@link ObjectSizeEstimator} of this pool.
     */
    public ObjectSizeEstimator getObjectSizeEstimator() {
        return settings.objectSizeEstimator;
    }

    /**
     * Sets the {@link ObjectSizeEstimator} of this pool.
     */
    public void setObjectSizeEstimator(ObjectSizeEstimator objectSizeEstimator) {
        if (objectSizeEstimator == null) {
            throw new NullPointerException("objectSizeEstimator");
        }

        settings = new Settings(
                objectSizeEstimator,
                settings.maxChannelMemorySize);
    }

    /**
     * Returns the maximum total size of the queued events per channel.
     */
    public long getMaxChannelMemorySize() {
        return settings.maxChannelMemorySize;
    }

    /**
     * Sets the maximum total size of the queued events per channel.
     * Specify {@code 0} to disable.
     */
    public void setMaxChannelMemorySize(long maxChannelMemorySize) {
        if (maxChannelMemorySize < 0) {
            throw new IllegalArgumentException(
                    "maxChannelMemorySize: " + maxChannelMemorySize);
        }

        if (getTaskCount() > 0) {
            throw new IllegalStateException(
                    "can't be changed after a task is executed");
        }

        settings = new Settings(
                settings.objectSizeEstimator,
                maxChannelMemorySize);
    }

    /**
     * Returns the maximum total size of the queued events for this pool.
     */
    public long getMaxTotalMemorySize() {
        if (totalLimiter == null) {
            return 0;
        }
        return totalLimiter.limit;
    }

    /**
     * If set to {@code false} no queued {@link ChannelEventRunnable}'s {@link ChannelFuture}
     * will get notified once {@link #shutdownNow()} is called.  If set to {@code true} every
     * queued {@link ChannelEventRunnable} will get marked as failed via {@link ChannelFuture#setFailure(Throwable)}.
     *
     * <p>
     * Please only set this to {@code false} if you want to handle the notification by yourself
     * and know what you are doing. Default is {@code true}.
     * </p>
     */
    public void setNotifyChannelFuturesOnShutdown(boolean notifyOnShutdown) {
        this.notifyOnShutdown = notifyOnShutdown;
    }

    /**
     * Returns if the {@link ChannelFuture}'s of the {@link ChannelEventRunnable}'s should be
     * notified about the shutdown of this {@link MemoryAwareThreadPoolExecutor}.
     */
    public boolean getNotifyChannelFuturesOnShutdown() {
        return notifyOnShutdown;
    }

    @Override
    public void execute(Runnable command) {
        if (command instanceof ChannelDownstreamEventRunnable) {
            throw new RejectedExecutionException("command must be enclosed with an upstream event.");
        }
        if (!(command instanceof ChannelEventRunnable)) {
            command = new MemoryAwareRunnable(command);
        }

        increaseCounter(command);
        doExecute(command);
    }

    /**
     * Put the actual execution logic here.  The default implementation simply
     * calls {@link #doUnorderedExecute(Runnable)}.
     */
    protected void doExecute(Runnable task) {
        doUnorderedExecute(task);
    }

    /**
     * Executes the specified task without maintaining the event order.
     */
    protected final void doUnorderedExecute(Runnable task) {
        super.execute(task);
    }

    @Override
    public boolean remove(Runnable task) {
        boolean removed = super.remove(task);
        if (removed) {
            decreaseCounter(task);
        }
        return removed;
    }

    @Override
    protected void beforeExecute(Thread t, Runnable r) {
        super.beforeExecute(t, r);
        decreaseCounter(r);
    }

    protected void increaseCounter(Runnable task) {
        if (!shouldCount(task)) {
            return;
        }

        Settings settings = this.settings;
        long maxChannelMemorySize = settings.maxChannelMemorySize;

        int increment = settings.objectSizeEstimator.estimateSize(task);

        if (task instanceof ChannelEventRunnable) {
            ChannelEventRunnable eventTask = (ChannelEventRunnable) task;
            eventTask.estimatedSize = increment;
            Channel channel = eventTask.getEvent().getChannel();
            long channelCounter = getChannelCounter(channel).addAndGet(increment);
            //System.out.println("IC: " + channelCounter + ", " + increment);
            if (maxChannelMemorySize != 0 && channelCounter >= maxChannelMemorySize && channel.isOpen()) {
                if (channel.isReadable()) {
                    //System.out.println("UNREADABLE");
                    ChannelHandlerContext ctx = eventTask.getContext();
                    if (ctx.getHandler() instanceof ExecutionHandler) {
                        // readSuspended = true;
                        ctx.setAttachment(Boolean.TRUE);
                    }
                    channel.setReadable(false);
                }
            }
        } else {
            ((MemoryAwareRunnable) task).estimatedSize = increment;
        }

        if (totalLimiter != null) {
            totalLimiter.increase(increment);
        }
    }

    protected void decreaseCounter(Runnable task) {
        if (!shouldCount(task)) {
            return;
        }

        Settings settings = this.settings;
        long maxChannelMemorySize = settings.maxChannelMemorySize;

        int increment;
        if (task instanceof ChannelEventRunnable) {
            increment = ((ChannelEventRunnable) task).estimatedSize;
        } else {
            increment = ((MemoryAwareRunnable) task).estimatedSize;
        }

        if (totalLimiter != null) {
            totalLimiter.decrease(increment);
        }

        if (task instanceof ChannelEventRunnable) {
            ChannelEventRunnable eventTask = (ChannelEventRunnable) task;
            Channel channel = eventTask.getEvent().getChannel();
            long channelCounter = getChannelCounter(channel).addAndGet(-increment);
            //System.out.println("DC: " + channelCounter + ", " + increment);
            if (maxChannelMemorySize != 0 && channelCounter < maxChannelMemorySize && channel.isOpen()) {
                if (!channel.isReadable()) {
                    //System.out.println("READABLE");
                    ChannelHandlerContext ctx = eventTask.getContext();
                    if (ctx.getHandler() instanceof ExecutionHandler) {
                        // check if the attachment was set as this means that we suspend the channel
                        // from reads. This only works when this pool is used with ExecutionHandler
                        // but I guess thats good enough for us.
                        //
                        // See #215
                        if (ctx.getAttachment() != null) {
                            // readSuspended = false;
                            ctx.setAttachment(null);
                            channel.setReadable(true);
                        }
                    } else {
                        channel.setReadable(true);
                    }
                }
            }
        }
    }

    private AtomicLong getChannelCounter(Channel channel) {
        AtomicLong counter = channelCounters.get(channel);
        if (counter == null) {
            counter = new AtomicLong();
            AtomicLong oldCounter = channelCounters.putIfAbsent(channel, counter);
            if (oldCounter != null) {
                counter = oldCounter;
            }
        }

        // Remove the entry when the channel closes.
        if (!channel.isOpen()) {
            channelCounters.remove(channel);
        }
        return counter;
    }

    /**
     * Returns {@code true} if and only if the specified {@code task} should
     * be counted to limit the global and per-channel memory consumption.
     * To override this method, you must call {@code super.shouldCount()} to
     * make sure important tasks are not counted.
     */
    protected boolean shouldCount(Runnable task) {
        if (task instanceof ChannelUpstreamEventRunnable) {
            ChannelUpstreamEventRunnable r = (ChannelUpstreamEventRunnable) task;
            ChannelEvent e = r.getEvent();
            if (e instanceof WriteCompletionEvent) {
                return false;
            } else if (e instanceof ChannelStateEvent) {
                if (((ChannelStateEvent) e).getState() == ChannelState.INTEREST_OPS) {
                    return false;
                }
            }
        }
        return true;
    }

    private static final class Settings {
        final ObjectSizeEstimator objectSizeEstimator;
        final long maxChannelMemorySize;

        Settings(ObjectSizeEstimator objectSizeEstimator,
                 long maxChannelMemorySize) {
            this.objectSizeEstimator = objectSizeEstimator;
            this.maxChannelMemorySize = maxChannelMemorySize;
        }
    }

    private static final class NewThreadRunsPolicy implements RejectedExecutionHandler {
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            try {
                final Thread t = new Thread(r, "Temporary task executor");
                t.start();
            } catch (Throwable e) {
                throw new RejectedExecutionException(
                        "Failed to start a new thread", e);
            }
        }
    }

    private static final class MemoryAwareRunnable implements Runnable {
        final Runnable task;
        int estimatedSize;

        MemoryAwareRunnable(Runnable task) {
            this.task = task;
        }

        public void run() {
            task.run();
        }
    }

    private static class Limiter {

        final long limit;
        private long counter;
        private int waiters;

        Limiter(long limit) {
            this.limit = limit;
        }

        synchronized void increase(long amount) {
            while (counter >= limit) {
                waiters ++;
                try {
                    wait();
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                } finally {
                    waiters --;
                }
            }
            counter += amount;
        }

        synchronized void decrease(long amount) {
            counter -= amount;
            if (counter < limit && waiters > 0) {
                notifyAll();
            }
        }
    }
}