src/java/org/apache/cassandra/concurrent/ThreadPoolExecutorJMXAdapter.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF 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.apache.cassandra.concurrent;

import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.TimeUnit;

import com.google.common.annotations.VisibleForTesting;

import org.apache.cassandra.metrics.ThreadPoolMetrics;
import org.apache.cassandra.utils.MBeanWrapper;

/**
 * A {@link ThreadPoolExecutorBase} adapter to expose it via JMX.
 * The executor is not itself modified to maximise code re-use.
 * Only its rejected execution handler is updated, and a shutdown listener is registered.
 */
@VisibleForTesting
public class ThreadPoolExecutorJMXAdapter implements Runnable, ResizableThreadPoolMXBean
{
    /**
     * A builder wrapper that delegates all methods except {@link Builder#build()}
     * @param <E>
     */
    static class Builder<E extends ThreadPoolExecutorBase> implements ExecutorBuilder<E>
    {
        final ExecutorBuilder<E> wrapped;
        final String jmxPath;
        Builder(ExecutorBuilder<E> wrapped, String jmxPath)
        {
            this.wrapped = wrapped;
            this.jmxPath = jmxPath;
        }

        @Override
        public ExecutorBuilder<E> withKeepAlive(long keepAlive, TimeUnit keepAliveUnits)
        {
            wrapped.withKeepAlive(keepAlive, keepAliveUnits);
            return this;
        }

        @Override
        public ExecutorBuilder<E> withKeepAlive()
        {
            wrapped.withKeepAlive();
            return this;
        }

        @Override
        public ExecutorBuilder<E> withThreadPriority(int threadPriority)
        {
            wrapped.withThreadPriority(threadPriority);
            return this;
        }

        @Override
        public ExecutorBuilder<E> withQueueLimit(int queueLimit)
        {
            wrapped.withQueueLimit(queueLimit);
            return this;
        }

        @Override
        public ExecutorBuilder<E> withThreadGroup(ThreadGroup threadGroup)
        {
            wrapped.withThreadGroup(threadGroup);
            return this;
        }

        @Override
        public ExecutorBuilder<E> withDefaultThreadGroup()
        {
            wrapped.withDefaultThreadGroup();
            return this;
        }

        @Override
        public ExecutorBuilder<E> withRejectedExecutionHandler(RejectedExecutionHandler rejectedExecutionHandler)
        {
            wrapped.withRejectedExecutionHandler(rejectedExecutionHandler);
            return this;
        }

        @Override
        public ExecutorBuilder<E> withUncaughtExceptionHandler(Thread.UncaughtExceptionHandler uncaughtExceptionHandler)
        {
            wrapped.withUncaughtExceptionHandler(uncaughtExceptionHandler);
            return this;
        }

        /**
         * Invoke {@link ExecutorBuilder#build()} on {@link #wrapped}, and register the resultant
         * {@link ThreadPoolExecutorBase} with a new {@link ThreadPoolExecutorJMXAdapter}.
         *
         * The executor constructed by {@link #wrapped} is returned.
         */
        @Override
        public E build()
        {
            E result = wrapped.build();
            register(jmxPath, result);
            return result;
        }
    }

    public static void register(String jmxPath, ThreadPoolExecutorBase executor)
    {
        new ThreadPoolExecutorJMXAdapter(jmxPath, executor);
    }

    final String mbeanName;
    final ThreadPoolExecutorBase executor;
    final ThreadPoolMetrics metrics;
    boolean released;

    private ThreadPoolExecutorJMXAdapter(String jmxPath, ThreadPoolExecutorBase executor)
    {
        this.executor = executor;
        this.mbeanName = "org.apache.cassandra." + jmxPath + ":type=" + executor.getThreadFactory().id;
        this.metrics = new ThreadPoolMetrics(executor, jmxPath, executor.getThreadFactory().id).register();
        executor.setRejectedExecutionHandler(rejectedExecutionHandler(metrics, executor.getRejectedExecutionHandler()));
        MBeanWrapper.instance.registerMBean(this, mbeanName);
        executor.onShutdown(this);
    }

    @Override
    public synchronized void run()
    {
        if (released)
            return;

        MBeanWrapper.instance.unregisterMBean(mbeanName);
        metrics.release();
        released = true;
    }

    public ThreadPoolMetrics metrics()
    {
        return metrics;
    }

    @Override
    public int getActiveTaskCount()
    {
        return executor.getActiveTaskCount();
    }

    @Override
    public int getPendingTaskCount()
    {
        return executor.getPendingTaskCount();
    }

    @Override
    public int getCoreThreads()
    {
        return executor.getCoreThreads();
    }

    @Override
    public void setCoreThreads(int number)
    {
        executor.setCoreThreads(number);
    }

    @Override
    public int getMaximumThreads()
    {
        return executor.getMaximumThreads();
    }

    @Override
    public void setMaximumThreads(int number)
    {
        executor.setMaximumThreads(number);
    }

    @Override
    public void setCorePoolSize(int corePoolSize)
    {
        executor.setCorePoolSize(corePoolSize);
    }

    @Override
    public int getCorePoolSize()
    {
        return executor.getCorePoolSize();
    }

    @Override
    public void setMaximumPoolSize(int maximumPoolSize)
    {
        executor.setMaximumPoolSize(maximumPoolSize);
    }

    @Override
    public int getMaximumPoolSize()
    {
        return executor.getMaximumPoolSize();
    }

    @Override
    public long getCompletedTaskCount()
    {
        return executor.getCompletedTaskCount();
    }

    @Override
    public int getMaxTasksQueued()
    {
        return executor.getMaxTasksQueued();
    }

    static RejectedExecutionHandler rejectedExecutionHandler(ThreadPoolMetrics metrics, RejectedExecutionHandler wrap)
    {
        return (task, executor) ->
        {
            metrics.totalBlocked.inc();
            metrics.currentBlocked.inc();
            try
            {
                wrap.rejectedExecution(task, executor);
            }
            finally
            {
                metrics.currentBlocked.dec();
            }
        };
    }
}