<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/impl/QueuePool.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;


import static com.github.bizyun.LogUtils.debugLog;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.github.phantomthief.util.ThrowableSupplier;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterators;
import com.google.common.util.concurrent.Uninterruptibles;

/**
 * @author zhangbiyun
 */
class QueuePool<E> implements Iterable<BlockingQueueHolder<E>>, MigrationLifecycle {

    private static final Logger logger = LoggerFactory.getLogger(QueuePool.class);
    private static final AtomicInteger POOL_ID = new AtomicInteger();

    private final int poolId;
    private final AtomicInteger state = new AtomicInteger(NORMAL);
    private final List<BlockingQueueHolder<E>> queueList;
    private final BlockingQueue<BlockingQueueHolder<E>> qq;

    private final AtomicInteger migratedQueueCount = new AtomicInteger();
    private volatile BlockingQueue<BlockingQueueHolder<E>> needMigratedQueueList;

    public QueuePool(Supplier<BlockingQueue<E>> queueSupplier, int queueCount) {
        this.poolId = POOL_ID.getAndIncrement();
        List<BlockingQueueHolder<E>> qList = new ArrayList<>(queueCount);
        for (int i = 0; i < queueCount; i++) {
            qList.add(new BlockingQueueHolder<>(queueSupplier.get(), poolId, i));
        }
        this.queueList = ImmutableList.copyOf(qList);
        this.qq = new LinkedBlockingQueue<>();
    }

    public BlockingQueueHolder<E> selectQueue(long key) {
        return queueList.get((int) (key % queueList.size()));
    }

    public BlockingQueueHolder<E> bindQueueBlock() throws InterruptedException {
        return doBindQueue(() -> {
            BlockingQueueHolder<E> queue = qq.take();
            if (isMigrated()) {
                // producer may put element when migrated, so migrated queue maybe not empty
                assert queue.isMigrated();
                qq.offer(queue);
                debugLog(logger, "[bindQueue-take] pool-{} is migrated, {}", poolId, queue);
                throw new InterruptedException();
            }
            return queue;
        });
    }

    @Nullable
    public BlockingQueueHolder<E> bindQueue() {
        return doBindQueue(() -> {
            BlockingQueueHolder<E> queue = qq.poll();
            if (isMigrated()) {
                if (queue != null) {
                    assert queue.isMigrated();
                    qq.offer(queue);
                }
                debugLog(logger, "[bindQueue-poll] pool-{} is migrated, {}", poolId, queue);
                return null;
            }
            return queue;
        });
    }

    @Nullable
    public BlockingQueueHolder<E> bindQueue(long timeout, TimeUnit unit) throws InterruptedException {
        return doBindQueue(() -> {
            BlockingQueueHolder<E> queue = qq.poll(timeout, unit);
            if (isMigrated()) {
                if (queue != null) {
                    assert queue.isMigrated();
                    qq.offer(queue);
                }
                debugLog(logger, "[bindQueue-pollTimeout] pool-{} is migrated, {}", poolId,
                        queue);
                return null;
            }
            return queue;
        });
    }

    private <T extends Throwable> BlockingQueueHolder<E> doBindQueue(
            ThrowableSupplier<BlockingQueueHolder<E>, T> queueSupplier) throws T {
        outer:
        while (true) {
            BlockingQueueHolder<E> queue = queueSupplier.get();
            if (queue == null) {
                return null;
            }
            while (true) {
                if (!queue.bind()) {
                    continue outer;
                }
                if (isNotEmpty(queue)) {
                    return queue;
                } else {
                    queue.unbind();
                }
                if (!queue.getQueue().isEmpty()) {
                    continue;
                }
                continue outer;
            }
        }
    }

    private boolean isNotEmpty(BlockingQueueHolder<E> queue) {
        if (!queue.getQueue().isEmpty()) {
            return true;
        }
        queue.getIdle().set(true);
        if (!queue.getQueue().isEmpty()) {
            queue.getIdle().set(false);
            return true;
        }
        return false;
    }

    public E peek() {
        BlockingQueueHolder<E> queue = qq.peek();
        if (queue == null) {
            return null;
        }
        return queue.getQueue().peek();
    }

    public void unbindQueue(BlockingQueueHolder<E> queue) {
        putQueueUninterrupted(queue.unbind());
    }

    private void putQueueUninterrupted(BlockingQueueHolder<E> queue) {
        Uninterruptibles.putUninterruptibly(qq, queue);
    }

    public void markQueueNotIdle(BlockingQueueHolder<E> queue) {
        if (queue.getIdle().compareAndSet(true, false)) {
            putQueueUninterrupted(queue);
        }
    }

    @Nullable
    public BlockingQueueHolder<E> bindQueueForMigrating() {
        assert isMigrating() || isMigrated();
        if (needMigratedQueueList == null) {
            return null;
        }
        BlockingQueueHolder<E> queue = needMigratedQueueList.poll();
        if (queue == null) {
            return null;
        }
        assert !(queue.isMigrating() || queue.isMigrated());

        if (!queue.bind()) {
            debugLog(logger, "[bindQueueForMigrating] queue is bound, pool-{}, {}", poolId
                    , queue);
            needMigratedQueueList.offer(queue);
            return null;
        }
        return queue;
    }

    public void increaseMigratedQueueCount() {
        migratedQueueCount.incrementAndGet();
    }

    public boolean allQueueMigrated() {
        return migratedQueueCount.get() == queueList.size();
    }

    @Override
    public boolean isMigrating() {
        return state.get() == MIGRATING;
    }

    @Override
    public boolean isMigrated() {
        return state.get() == MIGRATED;
    }

    @Override
    public void startMigrating() {
        boolean result = state.compareAndSet(NORMAL, MIGRATING);
        if (result) {
            needMigratedQueueList = new LinkedBlockingQueue<>(queueList);
        }
    }

    @Override
    public void stopMigrating() {
        boolean result = state.compareAndSet(MIGRATING, MIGRATED);
        if (result) {
            qq.offer(queueList.get(0));
        }
    }

    public void clear() {
        for (BlockingQueueHolder<E> queue : queueList) {
            queue.getQueue().clear();
        }
    }

    public int remainingCapacity() {
        long count = 0;
        for (BlockingQueueHolder<E> queue : queueList) {
            count += queue.getQueue().remainingCapacity();
        }
        if (count > Integer.MAX_VALUE) {
            return Integer.MAX_VALUE;
        }
        return (int) count;
    }

    public int drainTo(Collection<? super E> c) {
        int count = 0;
        for (BlockingQueueHolder<E> queue : queueList) {
            count += queue.getQueue().drainTo(c);
        }
        return count;
    }

    public int drainTo(Collection<? super E> c, int maxElements) {
        int count = 0;
        int left = maxElements;
        for (BlockingQueueHolder<E> queue : queueList) {
            count += queue.getQueue().drainTo(c, left);
            left = maxElements - count;
            if (left <= 0) {
                break;
            }
        }
        return count;
    }

    public Iterator<E> getElementIterator() {
        return Iterators.concat(queueList.stream()
                .map(BlockingQueueHolder::getQueue)
                .map(BlockingQueue::iterator).toArray(Iterator[]::new));
    }

    @Nonnull
    @Override
    public Iterator<BlockingQueueHolder<E>> iterator() {
        return queueList.iterator();
    }

    public int size() {
        return queueList.stream()
                .map(BlockingQueueHolder::getQueue)
                .mapToInt(BlockingQueue::size).sum();
    }


    public int getPoolId() {
        return poolId;
    }

    @Override
    public String toString() {
        return "QueuePool{" +
                "poolId=" + poolId +
                ", state=" + state +
                ", queueList.size=" + queueList.size() +
                ", qq.size=" + qq.size() +
                ", queueList=" + queueList +
                '}';
    }
}


<|start_filename|>src/test/java/com/github/bizyun/keythreadpool/impl/KeyThreadPoolExecutorTest.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

import static com.github.bizyun.keythreadpool.impl.KeyThreadPoolExecutor.newKeyThreadPool;
import static com.google.common.util.concurrent.MoreExecutors.shutdownAndAwaitTermination;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.junit.jupiter.api.Test;

import com.github.bizyun.keythreadpool.KeyRunnable;
import com.github.bizyun.keythreadpool.KeySupplier;
import com.github.bizyun.keythreadpool.impl.testtask.TestKeyCallable;
import com.github.bizyun.keythreadpool.impl.testtask.TestKeyRunnable;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.Uninterruptibles;

/**
 * @author zhangbiyun
 * @date 2020/6/21
 */
class KeyThreadPoolExecutorTest {

    @Test
    void test1() {
        AtomicInteger poolSize = new AtomicInteger(1);
        AtomicInteger queueCapacity = new AtomicInteger(10);
        AtomicInteger queueCount = new AtomicInteger(3);

        ExecutorService keyExecutor = newKeyThreadPool(() -> poolSize.get(),
                () -> new LinkedBlockingQueue<>(queueCapacity.get()), () -> queueCount.get());
        ExecutorService producerExecutor = Executors.newFixedThreadPool(10);
        ConcurrentHashMap<Long, LinkedBlockingQueue<Integer>> consumerRecords = new ConcurrentHashMap<>();
        int taskCount = 30000;
        AtomicInteger runTaskCount = new AtomicInteger();
        for (int i = 0; i < taskCount; i++) {
            final int key = i;
            producerExecutor.execute(() -> {
                keyExecutor.execute(new TestKeyRunnable(key, runTaskCount, key, consumerRecords));
            });
        }
        for (int i = 0; i < taskCount; i++) {
            final int key = i;
            producerExecutor.execute(() -> {
                keyExecutor.submit(new TestKeyRunnable(key, runTaskCount, key, consumerRecords));
            });
        }
        for (int i = 0; i < taskCount; i++) {
            final int key = i;
            producerExecutor.execute(() -> {
                keyExecutor.submit(new TestKeyCallable<>(key, runTaskCount, key, consumerRecords));
            });
        }
        for (int i = 0; i < taskCount; i++) {
            final int key = i;
            producerExecutor.execute(() -> {
                CompletableFuture.runAsync(new TestKeyRunnable(key, runTaskCount, key, consumerRecords), keyExecutor);
            });
        }
        ListeningExecutorService keyListeningExecutor = MoreExecutors.listeningDecorator(keyExecutor);
        AtomicInteger callbackCount = new AtomicInteger();
        for (int i = 0; i < taskCount; i++) {
            final int key = i;
            producerExecutor.execute(() -> {
                ListenableFuture<Object> future = keyListeningExecutor.submit(new TestKeyCallable<>(key, runTaskCount, key, consumerRecords));
                future.addListener(() -> {
                    callbackCount.incrementAndGet();
                }, MoreExecutors.directExecutor());
            });
        }

        // 扩容队列数
        queueCapacity.set(30);
        queueCount.set(10);
        for (int i = 0; i < taskCount; i++) {
            final int key = i;
            producerExecutor.execute(() -> {
                TestKeyRunnable runner = new TestKeyRunnable(key, runTaskCount, key, consumerRecords);
                keyExecutor.execute(() -> runner.run());
            });
        }
        Uninterruptibles.sleepUninterruptibly(200, TimeUnit.MILLISECONDS);

        // 扩容线程数，缩容队列数
        poolSize.set(8);
        queueCount.set(2);
        queueCapacity.set(1000);
        for (int i = 0; i < taskCount; i++) {
            final int key = i;
            producerExecutor.execute(() -> {
                TestKeyCallable<Void> callable = new TestKeyCallable<>(key, runTaskCount, key, consumerRecords);
                keyExecutor.submit(() -> callable.call());
            });
        }
        Uninterruptibles.sleepUninterruptibly(200, TimeUnit.MILLISECONDS);

        // 缩容线程数，扩容队列数
        poolSize.set(3);
        queueCount.set(20);
        queueCapacity.set(1000);
        for (int i = 0; i < taskCount; i++) {
            final int key = i;
            producerExecutor.execute(() -> {
                TestKeyRunnable runner = new TestKeyRunnable(key, runTaskCount, key, consumerRecords);
                keyExecutor.submit(() -> runner.run());
            });
        }

        shutdownAndAwaitTermination(producerExecutor, 1, TimeUnit.HOURS);
        assertEquals(3, ((KeyThreadPoolExecutor) keyExecutor).getCorePoolSize());
        shutdownAndAwaitTermination(keyExecutor, 1, TimeUnit.HOURS);
        assertEquals(taskCount * 8, runTaskCount.get());
        assertEquals(taskCount * 8, ((KeyThreadPoolExecutor) keyExecutor).getCompletedTaskCount());
        assertEquals(taskCount, callbackCount.get());
        assertEquals(taskCount, consumerRecords.size());
        for (int i = 0; i < taskCount; i++) {
            long key = i;
            LinkedBlockingQueue<Integer> records = consumerRecords.get(key);
            assertEquals(8, records.size());
            for (int j = 0; j < 8; j++) {
                assertEquals(i, records.poll());
            }
        }
    }

    @Test
    void test2() {
        AtomicInteger poolSize = new AtomicInteger(1);
        AtomicInteger queueCapacity = new AtomicInteger(100);
        AtomicInteger queueCount = new AtomicInteger(10);

        ExecutorService keyExecutor = newKeyThreadPool(() -> poolSize.get(),
                () -> new LinkedBlockingQueue<>(queueCapacity.get()), () -> queueCount.get());
        ConcurrentHashMap<Long, LinkedBlockingQueue<Integer>> consumerRecords = new ConcurrentHashMap<>();
        int taskCount = 30000;
        int keyMod = 17;
        Map<Integer, ExecutorService> producerExecutorMap = new HashMap<>();
        for (int i = 0; i < 17; i++) {
            producerExecutorMap.put(i, Executors.newFixedThreadPool(1));
        }

        AtomicInteger runTaskCount = new AtomicInteger();
        for (int i = 0; i < taskCount; i++) {
            final int key = i % keyMod;
            final int value = i;
            producerExecutorMap.get(key).execute(() -> {
                keyExecutor.execute(new TestKeyRunnable(key, runTaskCount, value, consumerRecords));
            });
        }
        for (int i = 0; i < taskCount; i++) {
            final int key = i % keyMod;
            final int value = i;
            producerExecutorMap.get(key).execute(() -> {
                keyExecutor.submit(new TestKeyRunnable(key, runTaskCount, value, consumerRecords));
            });
        }
        for (int i = 0; i < taskCount; i++) {
            final int key = i % keyMod;
            final int value = i;
            producerExecutorMap.get(key).execute(() -> {
                keyExecutor.submit(new TestKeyCallable<>(key, runTaskCount, value,
                        consumerRecords));
            });
        }
        for (int i = 0; i < taskCount; i++) {
            final int key = i % keyMod;
            final int value = i;
            producerExecutorMap.get(key).execute(() -> {
                CompletableFuture.runAsync(new TestKeyRunnable(key, runTaskCount, value,
                        consumerRecords), keyExecutor);
            });
        }
        ListeningExecutorService keyListeningExecutor = MoreExecutors.listeningDecorator(keyExecutor);
        AtomicInteger callbackCount = new AtomicInteger();
        for (int i = 0; i < taskCount; i++) {
            final int key = i % keyMod;
            final int value = i;
            producerExecutorMap.get(key).execute(() -> {
                ListenableFuture<Object> future =
                        keyListeningExecutor.submit(new TestKeyCallable<>(key, runTaskCount, value,
                                consumerRecords));
                future.addListener(() -> {
                    callbackCount.incrementAndGet();
                }, MoreExecutors.directExecutor());
            });
        }

        // 扩容队列数
        queueCount.set(100000);
        queueCapacity.set(300);
        for (int i = 0; i < taskCount; i++) {
            final int key = i % keyMod;
            final int value = i;
            producerExecutorMap.get(key).execute(() -> {
                TestKeyRunnable runner = new TestKeyRunnable(key, runTaskCount, value, consumerRecords);
                keyExecutor.execute(() -> runner.run());
            });
        }
        Uninterruptibles.sleepUninterruptibly(200, TimeUnit.MILLISECONDS);

        // 扩容线程数，缩容队列数
        poolSize.set(20);
        queueCount.set(2);
        queueCapacity.set(400000);
        for (int i = 0; i < taskCount; i++) {
            final int key = i % keyMod;
            final int value = i;
            producerExecutorMap.get(key).execute(() -> {
                TestKeyCallable<Void> callable = new TestKeyCallable<>(key, runTaskCount, value,
                        consumerRecords);
                keyExecutor.submit(() -> callable.call());
            });
        }
        Uninterruptibles.sleepUninterruptibly(200, TimeUnit.MILLISECONDS);

        // 缩容线程数，扩容队列数
        poolSize.set(3);
        queueCount.set(20);
        queueCapacity.set(400100);
        for (int i = 0; i < taskCount; i++) {
            final int key = i % keyMod;
            final int value = i;
            producerExecutorMap.get(key).execute(() -> {
                TestKeyRunnable runner = new TestKeyRunnable(key, runTaskCount, value, consumerRecords);
                keyExecutor.submit(() -> runner.run());
            });
        }
        producerExecutorMap.values().forEach(producerExecutor -> shutdownAndAwaitTermination(producerExecutor, 1, TimeUnit.HOURS));
        assertEquals(3, ((KeyThreadPoolExecutor) keyExecutor).getCorePoolSize());
        shutdownAndAwaitTermination(keyExecutor, 1, TimeUnit.HOURS);
        assertEquals(taskCount * 8, runTaskCount.get());
        assertEquals(taskCount * 8, ((KeyThreadPoolExecutor) keyExecutor).getCompletedTaskCount());
        assertEquals(taskCount, callbackCount.get());
        assertEquals(keyMod, consumerRecords.size());
        for (int j = 0; j < 8; j++) {
            for (int i = 0; i < taskCount; i++) {
                long key = i % keyMod;
                LinkedBlockingQueue<Integer> records = consumerRecords.get(key);
                assertEquals(i, records.poll());
            }
        }
    }

    @Test
    void test3() throws InterruptedException {
        AtomicInteger poolSize = new AtomicInteger(10);
        AtomicInteger queueCapacity = new AtomicInteger(1);
        AtomicInteger queueCount = new AtomicInteger(100);

        ExecutorService keyExecutor = newKeyThreadPool(() -> poolSize.get(),
                () -> new LinkedBlockingQueue<>(queueCapacity.get()), () -> queueCount.get());
        CountDownLatch latch = new CountDownLatch(10);
        for (int i = 0; i < 100; i++) {
            final int k = i;
            keyExecutor.execute(new KeyRunnable() {
                @Override
                public long getKey() {
                    return k;
                }

                @Override
                public void run() {
                    latch.countDown();
                    try {
                        Thread.sleep(1000000);
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                    }
                }
            });
        }
        latch.await();
        List<Runnable> list = keyExecutor.shutdownNow();
        assertEquals(90, list.size());
        keyExecutor.awaitTermination(1, TimeUnit.HOURS);
        long completedTaskCount = ((KeyThreadPoolExecutor) keyExecutor).getCompletedTaskCount();
        assertEquals(10, completedTaskCount);
    }

    @Test
    void test4() throws InterruptedException {
        ExecutorService keyExecutor = newKeyThreadPool(() -> 2,
                () -> new LinkedBlockingQueue<>(1), () -> 2,
               new KeyThreadPoolExecutor.DiscardOldestPolicy());

        CountDownLatch latch = new CountDownLatch(2);
        BlockingQueue<Runnable> queue = ((KeyThreadPoolExecutor) keyExecutor).getQueue();
        for (int i = 0; i < 101; i++) {
            final int k = i;
            keyExecutor.execute(new KeyRunnable() {
                @Override
                public long getKey() {
                    return k;
                }

                @Override
                public void run() {
                    latch.countDown();
                    try {
                        Thread.sleep(100000);
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                    }
                }
            });
            if (i == 1) {
                latch.await();
            }
            if (i <= 1) {
                continue;
            }
            assertNull(queue.poll());
            assertNull(queue.poll(2, TimeUnit.MILLISECONDS));
            if (i == 2) {
                assertEquals(1, queue.size());
                Runnable next = queue.iterator().next();
                assertEquals(2, ((KeySupplier) next).getKey());
            }
            if (i >= 3) {
                if (i % 2 == 1) {
                    checkQueue(queue, i-1, i);
                } else {
                    checkQueue(queue, i, i - 1);
                }
            }
        }
        List<Runnable> list = new ArrayList<>();
        queue.drainTo(list);
        assertEquals(2, list.size());
        keyExecutor.shutdownNow();
        keyExecutor.awaitTermination(1, TimeUnit.HOURS);
    }

    @Test
    void test5() throws InterruptedException {
        ExecutorService keyExecutor = newKeyThreadPool(() -> 1,
                () -> new LinkedBlockingQueue<>(), () -> 3);

        CountDownLatch latch = new CountDownLatch(1);
        BlockingQueue<Runnable> queue = ((KeyThreadPoolExecutor) keyExecutor).getQueue();
        for (int i = 0; i < 100; i++) {
            final int k = i;
            keyExecutor.execute(new KeyRunnable() {
                @Override
                public long getKey() {
                    return k;
                }

                @Override
                public void run() {
                    if (k == 0) {
                        latch.countDown();
                        try {
                            Thread.sleep(100000);
                        } catch (InterruptedException e) {
                            Thread.currentThread().interrupt();
                        }
                    }
                }
            });
            if (i == 0) {
                latch.await();
            }
        }
        Runnable task1 = queue.poll(2, TimeUnit.MILLISECONDS);
        assertNotNull(task1);
        Runnable task2 = queue.poll();
        assertNotNull(task2);
        Runnable task3 = queue.poll();
        assertNull(task3);
        task3 = queue.poll(2, TimeUnit.MILLISECONDS);
        assertNull(task3);
        task1.run();
        task3 = queue.poll();
        assertNotNull(task3);
        Runnable task4 = queue.poll();
        assertNull(task4);
        task2.run();
        Runnable task5 = queue.take();
        assertNotNull(task5);
        task4 = queue.poll();
        assertNull(task4);
        task5.run();
        task4 = queue.poll();
        assertNotNull(task4);
        assertNull(queue.peek());
        task4.run();
        assertNotNull(queue.peek());
        task4 = queue.poll();
        assertNotNull(task4);

        assertTrue(queue.size() > 0);
        queue.clear();
        assertTrue(queue.size() == 0);

        keyExecutor.shutdownNow();
        keyExecutor.awaitTermination(1, TimeUnit.HOURS);
    }

    @Test
    void test6() {
        assertThrows(IllegalArgumentException.class, () -> newKeyThreadPool(() -> 1,
                () -> new LinkedBlockingQueue<>(), () -> 3,
                new ThreadPoolExecutor.DiscardOldestPolicy()));

        assertThrows(IllegalArgumentException.class, () -> newKeyThreadPool(() -> 1,
                () -> new LinkedBlockingQueue<>(), () -> 3,
                new ThreadPoolExecutor.CallerRunsPolicy()));

        assertThrows(IllegalArgumentException.class, () -> newKeyThreadPool(() -> 0,
                () -> new LinkedBlockingQueue<>(),
                () -> 1));

        assertThrows(IllegalArgumentException.class, () -> newKeyThreadPool(() -> 1,
                () -> new LinkedBlockingQueue<>(),
                () -> 0));

        KeyThreadPoolExecutor keyExecutor = newKeyThreadPool(() -> 1,
                () -> new LinkedBlockingQueue<>(),
                () -> 1);

        assertThrows(IllegalArgumentException.class, () ->
                keyExecutor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy()));
    }

    @Test
    void test7() throws InterruptedException {
        AtomicInteger queueCount = new AtomicInteger(5);
        AtomicInteger queueCapacity = new AtomicInteger(10);
        KeyThreadPoolExecutor keyExecutor = newKeyThreadPool(() -> 1,
                () -> new LinkedBlockingQueue<>(queueCapacity.get()),
                queueCount::get);
        CountDownLatch latch = new CountDownLatch(1);
        CountDownLatch latch1 = new CountDownLatch(1);
        AtomicInteger runTaskCount = new AtomicInteger();
        for (int i = 0; i < 50; i++) {
            final int k = i;
            KeySupplier keySupplier = () -> k;
            keyExecutor.execute(() -> {
                try {
                    latch1.countDown();
                    latch.await();
                    runTaskCount.getAndIncrement();
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    KeySupplier ks = keySupplier;
                }
            });
        }
        latch1.await();
        queueCount.set(1);
        queueCapacity.set(20);
        latch.countDown();
        MoreExecutors.shutdownAndAwaitTermination(keyExecutor, 1, TimeUnit.HOURS);
        assertEquals(21, runTaskCount.get());
        assertEquals(21, keyExecutor.getCompletedTaskCount());
    }

    @Test
    void test8() {
        AtomicInteger queueCount = new AtomicInteger(200000);
        AtomicInteger queueCapacity = new AtomicInteger(1);
        KeyThreadPoolExecutor keyExecutor = newKeyThreadPool(() -> 100,
                () -> new LinkedBlockingQueue<>(queueCapacity.get()),
                queueCount::get);
        ConcurrentHashMap<Long, LinkedBlockingQueue<Integer>> consumerRecords = new ConcurrentHashMap<>();
        long key1 = 37;
        AtomicInteger runTaskCount = new AtomicInteger();
        keyExecutor.execute(new TestKeyRunnable(key1, runTaskCount, 12, consumerRecords));
        while (runTaskCount.get() != 1);
        queueCount.set(10);
        long key2 = key1;
        keyExecutor.execute(new TestKeyRunnable(key2, runTaskCount, 8, consumerRecords));
        MoreExecutors.shutdownAndAwaitTermination(keyExecutor, 1, TimeUnit.MINUTES);

        assertEquals(1, consumerRecords.size());
        assertEquals(keyExecutor.getCompletedTaskCount(), runTaskCount.get());
        LinkedBlockingQueue<Integer> records = consumerRecords.get(key1);
        assertEquals(12, records.poll());
        records = consumerRecords.get(key2);
        assertEquals(8, records.poll());
        assertNull(records.poll());
    }

    @Test
    void test9() {
        AtomicInteger queueCount = new AtomicInteger(2);
        AtomicInteger queueCapacity = new AtomicInteger(1);
        KeyThreadPoolExecutor keyExecutor = newKeyThreadPool(() -> 100,
                () -> new LinkedBlockingQueue<>(queueCapacity.get()),
                queueCount::get);
        ConcurrentHashMap<Long, LinkedBlockingQueue<Integer>> consumerRecords = new ConcurrentHashMap<>();
        long key1 = 37;
        AtomicInteger runTaskCount = new AtomicInteger();
        TestKeyRunnable command = new TestKeyRunnable(key1, runTaskCount, 12, consumerRecords);
        CountDownLatch latch = new CountDownLatch(1);
        keyExecutor.execute(() -> {
            command.run();
            Uninterruptibles.awaitUninterruptibly(latch);
        });
        while (runTaskCount.get() != 1);
        queueCount.set(10);
        long key2 = 38;
        keyExecutor.execute(new TestKeyRunnable(key2, runTaskCount, 8, consumerRecords));
        Uninterruptibles.sleepUninterruptibly(1, TimeUnit.SECONDS);
        latch.countDown();
        Uninterruptibles.sleepUninterruptibly(500, TimeUnit.MILLISECONDS);
        MoreExecutors.shutdownAndAwaitTermination(keyExecutor, 1, TimeUnit.MINUTES);

        assertEquals(2, consumerRecords.size());
        assertEquals(keyExecutor.getCompletedTaskCount(), runTaskCount.get());
        LinkedBlockingQueue<Integer> records = consumerRecords.get(key1);
        assertEquals(12, records.poll());
        records = consumerRecords.get(key2);
        assertEquals(8, records.poll());
        assertNull(records.poll());

    }

    private void checkQueue(BlockingQueue<Runnable> queue, int i2, int i3) {
        assertEquals(2, queue.size());
        Iterator<Runnable> iterator = queue.iterator();
        Runnable next = iterator.next();
        assertEquals(i2, ((KeySupplier) next).getKey());
        next = iterator.next();
        assertEquals(i3, ((KeySupplier) next).getKey());
    }
}

<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/impl/MigrationLifecycle.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

/**
 * @author zhangbiyun
 */
interface MigrationLifecycle {

    int NORMAL = 0;
    int MIGRATING = 1;
    int MIGRATED = 2;

    void startMigrating();

    void stopMigrating();

    boolean isMigrating();

    boolean isMigrated();
}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/impl/KeyFutureTask.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

import java.util.concurrent.Callable;
import java.util.concurrent.FutureTask;

import javax.annotation.Nonnull;

import com.github.bizyun.keythreadpool.KeySupplier;


/**
 * @author zhangbiyun
 */
class KeyFutureTask<V> extends FutureTask<V> implements KeySupplier {

    private final KeySupplier keySupplier;

    public KeyFutureTask(@Nonnull Callable<V> callable) {
        super(callable);
        if (callable instanceof KeySupplier) {
            this.keySupplier = (KeySupplier) callable;
        } else {
            this.keySupplier = scanKeySupplier(callable);
        }
    }

    public KeyFutureTask(@Nonnull Runnable runnable, V result) {
        super(runnable, result);
        if (runnable instanceof KeySupplier) {
            this.keySupplier = (KeySupplier) runnable;
        } else {
            this.keySupplier = scanKeySupplier(runnable);
        }
    }

    private KeySupplier scanKeySupplier(@Nonnull Object object) {
        KeySupplier keySupplier = new KeySupplierScanner().scan(object);
        if (keySupplier == null) {
            throw new IllegalStateException("not found keySupplier");
        }
        return keySupplier;
    }

    @Override
    public long getKey() {
        return keySupplier.getKey();
    }
}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/impl/BlockingQueueHolder.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

import static com.github.bizyun.LogUtils.debugLog;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author zhangbiyun
 */
class BlockingQueueHolder<E> implements MigrationLifecycle {

    private static final Logger logger = LoggerFactory.getLogger(BlockingQueueHolder.class);

    private final int poolId;
    private final int queueId;
    private final BlockingQueueProxy<E> queue;
    private final AtomicBoolean idle = new AtomicBoolean(true);
    private final AtomicReference<Thread> bindThread = new AtomicReference<>();
    private final AtomicInteger state = new AtomicInteger(NORMAL);

    public BlockingQueueHolder(BlockingQueue<E> queue, int poolId, int queueId) {
        this.poolId = poolId;
        this.queueId = queueId;
        this.queue = new BlockingQueueProxy<>(queue);
    }

    public BlockingQueue<E> getQueue() {
        return queue;
    }

    boolean bind() {
        boolean result = bindThread.compareAndSet(null, Thread.currentThread());
        debugLog(logger, "[bind-queue], {}, {}", result, this);
        return result;
    }

    BlockingQueueHolder<E> unbind() {
        boolean result = bindThread.compareAndSet(Thread.currentThread(), null);
        assert result;
        debugLog(logger, "[unbind-queue], {}", this);
        return this;
    }

    public AtomicBoolean getIdle() {
        return idle;
    }

    @Override
    public String toString() {
        Thread thread = bindThread.get();
        return "BlockingQueueHolder{" +
                "poolId=" + poolId +
                ", queueId=" + queueId +
                ", queue.size=" + queue.size() +
                ", state=" + state +
                ", idle=" + idle.get() +
                ", bindThread=" + (thread == null ? "null" : thread.getName()) +
                '}';
    }

    @Override
    public void startMigrating() {
        state.compareAndSet(NORMAL, MIGRATING);
    }

    @Override
    public void stopMigrating() {
        state.compareAndSet(MIGRATING, MIGRATED);
    }

    @Override
    public boolean isMigrating() {
        return state.get() == MIGRATING;
    }

    @Override
    public boolean isMigrated() {
        return state.get() == MIGRATED;
    }
}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/impl/KeyThreadPoolExecutor.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.Executors;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.RunnableFuture;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.IntSupplier;
import java.util.function.Supplier;

import javax.annotation.Nonnull;

import com.github.bizyun.keythreadpool.KeyRunnable;
import com.github.bizyun.keythreadpool.KeySupplier;
import com.google.common.base.Preconditions;

/**
 * @author zhangbiyun
 */
public class KeyThreadPoolExecutor extends ThreadPoolExecutor {

    private final IntSupplier poolSizeSupplier;
    private volatile boolean hasStartCoreThreads = false;

    public static KeyThreadPoolExecutor newKeyThreadPool(@Nonnull IntSupplier poolSizeSupplier,
            @Nonnull Supplier<BlockingQueue<Runnable>> queueSupplier,
            @Nonnull IntSupplier queueCountSupplier) {
        return new KeyThreadPoolExecutor(poolSizeSupplier.getAsInt(),
                poolSizeSupplier,
                Executors.defaultThreadFactory(),
                new CallerBlockingPolicy(),
                queueSupplier,
                queueCountSupplier);
    }

    public static KeyThreadPoolExecutor newKeyThreadPool(@Nonnull IntSupplier poolSizeSupplier,
            @Nonnull ThreadFactory threadFactory,
            @Nonnull Supplier<BlockingQueue<Runnable>> queueSupplier,
            @Nonnull IntSupplier queueCountSupplier) {
        return new KeyThreadPoolExecutor(poolSizeSupplier.getAsInt(),
                poolSizeSupplier,
                threadFactory,
                new CallerBlockingPolicy(),
                queueSupplier,
                queueCountSupplier);
    }

    public static KeyThreadPoolExecutor newKeyThreadPool(@Nonnull IntSupplier poolSizeSupplier,
            @Nonnull Supplier<BlockingQueue<Runnable>> queueSupplier,
            @Nonnull IntSupplier queueCountSupplier,
            @Nonnull RejectedExecutionHandler handler) {
        return new KeyThreadPoolExecutor(poolSizeSupplier.getAsInt(),
                poolSizeSupplier,
                Executors.defaultThreadFactory(),
                handler,
                queueSupplier,
                queueCountSupplier);
    }

    public static KeyThreadPoolExecutor newKeyThreadPool(@Nonnull IntSupplier poolSizeSupplier,
            @Nonnull ThreadFactory threadFactory,
            @Nonnull Supplier<BlockingQueue<Runnable>> queueSupplier,
            @Nonnull IntSupplier queueCountSupplier,
            @Nonnull RejectedExecutionHandler handler) {
        return new KeyThreadPoolExecutor(poolSizeSupplier.getAsInt(),
                poolSizeSupplier,
                threadFactory,
                handler,
                queueSupplier,
                queueCountSupplier);
    }


    private KeyThreadPoolExecutor(int poolSize, IntSupplier poolSizeSupplier,
            ThreadFactory threadFactory, RejectedExecutionHandler handler,
            Supplier<BlockingQueue<Runnable>> queueSupplier, IntSupplier queueCountSupplier) {
        super(poolSize, poolSize, 0, TimeUnit.MILLISECONDS,
                new KeyBlockingQueue(queueSupplier, queueCountSupplier), threadFactory, handler);
        this.poolSizeSupplier = poolSizeSupplier;
        checkNotSupportPolicy(handler);
    }

    private void checkNotSupportPolicy(RejectedExecutionHandler handler) {
        Preconditions.checkArgument(!(handler instanceof CallerRunsPolicy));
        Preconditions.checkArgument(!(handler instanceof ThreadPoolExecutor.DiscardOldestPolicy));
    }

    @Override
    public void execute(Runnable command) {
        Preconditions.checkNotNull(command);

        if (!(command instanceof KeySupplier)) {
            command = new KeyRunnableWrapper(command);
        }

        if (!hasStartCoreThreads) {
            hasStartCoreThreads = true;
            prestartAllCoreThreads();
        }
        if (!isShutdown() && getQueue().offer(command)) {
            tryChangePoolSize();
            if (isShutdown() && getQueue().remove(command)) {
                getRejectedExecutionHandler().rejectedExecution(command, this);
            }
        } else {
            getRejectedExecutionHandler().rejectedExecution(command, this);
        }
    }

    private void tryChangePoolSize() {
        int poolSize = poolSizeSupplier.getAsInt();
        if (poolSize <= 0) {
            return;
        }
        if (getCorePoolSize() != poolSize) {
            if (poolSize > getCorePoolSize()) {
                setMaximumPoolSize(poolSize);
                setCorePoolSize(poolSize);
                prestartAllCoreThreads();
            } else {
                setCorePoolSize(poolSize);
                setMaximumPoolSize(poolSize);
            }
        }
    }

    @Override
    protected <T> RunnableFuture<T> newTaskFor(Runnable runnable, T value) {
        return new KeyFutureTask<>(runnable, value);
    }

    @Override
    protected <T> RunnableFuture<T> newTaskFor(Callable<T> callable) {
        return new KeyFutureTask<>(callable);
    }

    @Override
    public void setRejectedExecutionHandler(RejectedExecutionHandler handler) {
        super.setRejectedExecutionHandler(handler);
        checkNotSupportPolicy(handler);
    }

    private static class KeyRunnableWrapper implements KeyRunnable {

        private final Runnable original;
        private final KeySupplier keySupplier;

        public KeyRunnableWrapper(Runnable original) {
            this.original = original;
            this.keySupplier = new KeySupplierScanner().scan(original);
            if (this.keySupplier == null) {
                throw new IllegalStateException("not found keySupplier");
            }
        }

        @Override
        public void run() {
            original.run();
        }

        @Override
        public long getKey() {
            return keySupplier.getKey();
        }
    }

    public static class CallerBlockingPolicy implements RejectedExecutionHandler {
        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            BlockingQueue<Runnable> queue = executor.getQueue();
            boolean interrupted = false;
            try {
                while (true) {
                    try {
                        if (executor.isShutdown()) {
                            throw new RejectedExecutionException("executor is shutdown");
                        }
                        queue.put(r);
                        break;
                    } catch (InterruptedException e) {
                        interrupted = true;
                    }
                }
            } finally {
                if (interrupted) {
                    Thread.currentThread().interrupt();
                }
            }
        }
    }

    public static class DiscardOldestPolicy implements RejectedExecutionHandler {
        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            if (!executor.isShutdown()) {
                BlockingQueue<Runnable> queue = executor.getQueue();
                if (queue instanceof KeyQueue && r instanceof KeySupplier) {
                    ((KeyQueue<Runnable>) queue).poll((KeySupplier) r);
                    executor.execute(r);
                } else {
                    throw new IllegalStateException("is not KeyThreadPoolExecutor");
                }
            }
        }
    }

}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/impl/KeyBlockingQueue.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

import static com.github.bizyun.LogUtils.debugLog;
import static com.google.common.base.Suppliers.memoize;

import java.util.AbstractQueue;
import java.util.Collection;
import java.util.Iterator;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.IntSupplier;
import java.util.function.Supplier;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.github.bizyun.keythreadpool.KeySupplier;
import com.github.phantomthief.util.ThrowablePredicate;
import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.Uninterruptibles;

/**
 * @author zhangbiyun
 */
class KeyBlockingQueue extends AbstractQueue<Runnable> implements BlockingQueue<Runnable>,
        KeyQueue<Runnable> {

    private static final Logger logger = LoggerFactory.getLogger(KeyBlockingQueue.class);

    private final Supplier<BlockingQueue<Runnable>> queueSupplier;
    private final IntSupplier queueCountSupplier;
    private volatile int curQueueCount;
    private volatile Supplier<QueuePool<Runnable>> queuePoolSupplier;

    private final AtomicBoolean expandOrShrink = new AtomicBoolean();
    private volatile CountDownLatch expandOrShrinkProducerLock;
    private volatile Supplier<QueuePool<Runnable>> backupQueuePoolSupplier;

    public KeyBlockingQueue(Supplier<BlockingQueue<Runnable>> queueSupplier,
            IntSupplier queueCountSupplier) {
        this.curQueueCount = queueCountSupplier.getAsInt();
        Preconditions.checkArgument(this.curQueueCount > 0, "queue count <= 0");
        this.queueCountSupplier = queueCountSupplier;
        this.queueSupplier = queueSupplier;
        this.queuePoolSupplier = memoize(() -> new QueuePool<>(queueSupplier, this.curQueueCount));
    }

    @Override
    public Iterator<Runnable> iterator() {
        return queuePool().getElementIterator();
    }

    @Override
    public int size() {
        return queuePool().size();
    }

    @Override
    public void put(@Nonnull Runnable runnable) throws InterruptedException {
        addAndMark(runnable, queue -> {
            queue.getQueue().put(runnable);
            return true;
        });
    }

    @Override
    public boolean offer(Runnable runnable, long timeout, @Nonnull TimeUnit unit) throws InterruptedException {
        return addAndMark(runnable, queue -> queue.getQueue().offer(runnable, timeout, unit));
    }

    @Override
    public boolean offer(Runnable runnable) {
        return addAndMark(runnable, queue -> queue.getQueue().offer(runnable));
    }


    private <T extends Throwable> boolean addAndMark(Runnable runnable,
            ThrowablePredicate<BlockingQueueHolder<Runnable>, T> predicate) throws T {
        while (true) {
            isExpandOrShrink();
            QueuePool<Runnable> pool = queuePool();
            BlockingQueueHolder<Runnable> queue = pool.selectQueue(getKey(runnable));
            if (!predicate.test(queue)) {
                return false;
            }
            if (queue.isMigrating() || queue.isMigrated()) {
                if (removeFromQueue(runnable, queue)) {
                    continue;
                }
            }
            pool.markQueueNotIdle(queue);
            return true;
        }
    }

    private boolean removeFromQueue(Runnable runnable, BlockingQueueHolder<Runnable> queue) {
        if (queue.getQueue().remove(runnable)) {
            debugLog(logger, "[producer] put to a migrating or migrated queue, remove success, {}",
                    queue);
            return true;
        }
        debugLog(logger, "[producer] put to a migrating or migrated queue, remove failed，maybe "
                        + "migrated to new queue, {}", queue);
        return false;
    }

    private void isExpandOrShrink() {
        if (expandOrShrink.get()) {
            CountDownLatch latch = expandOrShrinkProducerLock;
            if (latch != null) {
                debugLog(logger, "[expandOrShrink] producer await");
                Uninterruptibles.awaitUninterruptibly(latch);
                debugLog(logger, "[expandOrShrink] producer wakeup");
            }
        }
    }

    private QueuePool<Runnable> queuePool() {
        return queuePoolSupplier.get();
    }

    @Nonnull
    @Override
    public Runnable take() throws InterruptedException {
        tryExpandOrShrink();
        QueuePool<Runnable> pool = queuePool();
        BlockingQueueHolder<Runnable> queue = pool.bindQueueBlock();
        assert !queue.isMigrating() && !queue.isMigrated();
        Runnable runnable;
        try {
            runnable = queue.getQueue().take();
        } catch (InterruptedException e) {
            pool.unbindQueue(queue);
            throw e;
        }
        return () -> {
            try {
                runnable.run();
            } finally {
                pool.unbindQueue(queue);
            }
        };
    }

    @Nullable
    @Override
    public Runnable poll(long timeout, @Nonnull TimeUnit unit) throws InterruptedException {
        tryExpandOrShrink();
        QueuePool<Runnable> pool = queuePool();
        BlockingQueueHolder<Runnable> queue = pool.bindQueue(timeout, unit);
        if (queue == null) {
            return null;
        }
        assert !queue.isMigrating() && !queue.isMigrated();
        Runnable runnable;
        try {
            runnable = queue.getQueue().poll(timeout, unit);
            assert runnable != null;
        } catch (InterruptedException e) {
            pool.unbindQueue(queue);
            throw e;
        }
        return () -> {
            try {
                runnable.run();
            } finally {
                pool.unbindQueue(queue);
            }
        };
    }

    @Nullable
    @Override
    public Runnable poll() {
        tryExpandOrShrink();
        QueuePool<Runnable> pool = queuePool();
        BlockingQueueHolder<Runnable> queue = pool.bindQueue();
        if (queue == null) {
            return null;
        }
        assert !queue.isMigrating() && !queue.isMigrated();
        Runnable runnable = queue.getQueue().poll();
        assert runnable != null;
        return () -> {
            try {
                runnable.run();
            } finally {
                pool.unbindQueue(queue);
            }
        };
    }

    private void tryExpandOrShrink() {
        QueuePool<Runnable> pool = queuePool();
        int queueCount = this.queueCountSupplier.getAsInt();
        if (queueCount > 0 && queueCount != this.curQueueCount && expandOrShrink.compareAndSet(false, true)) {
            debugLog(logger, "[expandOrShrink] start, oldQueueCount:{}, newQueueCount:{}, " +
                    "pool:{}", this.curQueueCount, queueCount, pool);
            this.curQueueCount = queueCount;
            this.backupQueuePoolSupplier = memoize(() -> new QueuePool<>(queueSupplier, this.curQueueCount));
            this.expandOrShrinkProducerLock = new CountDownLatch(1);
            pool.startMigrating();
        }

        if (pool.isMigrating()) {
            doMigrate(pool);
        }
    }

    private void doMigrate(QueuePool<Runnable> pool) {
        while (true) {
            BlockingQueueHolder<Runnable> queue = pool.bindQueueForMigrating();
            if (queue != null) {
                migrateOneQueue(queue);
                pool.increaseMigratedQueueCount();
            } else if (pool.allQueueMigrated()) {
                synchronized (expandOrShrink) {
                    if (!pool.isMigrated()) {
                        this.queuePoolSupplier = this.backupQueuePoolSupplier;
                        pool.stopMigrating();
                        this.expandOrShrinkProducerLock.countDown();
                        this.backupQueuePoolSupplier = null;
                        this.expandOrShrinkProducerLock = null;
                        this.expandOrShrink.set(false);
                        debugLog(logger, "[expandOrShrink] complete, oldPool:{}, " +
                                "newPool:{}", pool, queuePool());
                    } else {
                        debugLog(logger, "[expandOrShrink] complete, oldPoolId:{}, " +
                                "newPoolId:{}", pool.getPoolId(), queuePool().getPoolId());
                    }
                    return;
                }
            } else {
                Uninterruptibles.sleepUninterruptibly(5, TimeUnit.MILLISECONDS);
            }
        }
    }

    private void migrateOneQueue(BlockingQueueHolder<Runnable> queue) {
        QueuePool<Runnable> backupPool = this.backupQueuePoolSupplier.get();
        queue.startMigrating();
        debugLog(logger, "[migrateOneQueue] {}", queue);
        for (Runnable r; (r = queue.getQueue().poll()) != null; ) {
            BlockingQueueHolder<Runnable> backupQueue = backupPool.selectQueue(getKey(r));
            if (!backupQueue.getQueue().offer(r)) {
                debugLog(logger, "[migrateOneQueue] queue is full, {}", queue);
                continue;
            }
            backupPool.markQueueNotIdle(backupQueue);
//            debugLog(logger, "[migrateOneQueue] backupQueue:{}", backupQueue);
        }
        debugLog(logger, "[migrateOneQueue] complete {}", queue);
        queue.stopMigrating();
    }

    @Override
    public Runnable peek() {
        return queuePool().peek();
    }

    @Override
    public void clear() {
        queuePool().clear();
    }

    @Override
    public int remainingCapacity() {
        return queuePool().remainingCapacity();
    }

    @Override
    public int drainTo(@Nonnull Collection<? super Runnable> c) {
        return queuePool().drainTo(c);
    }

    @Override
    public int drainTo(@Nonnull Collection<? super Runnable> c, int maxElements) {
        return queuePool().drainTo(c, maxElements);
    }

    @Override
    public boolean remove(Object o) {
        if (o instanceof KeySupplier) {
            BlockingQueueHolder<Runnable> queue = queuePool().selectQueue(((KeySupplier) o).getKey());
            return queue.getQueue().remove(o);
        }
        return false;
    }

    private long getKey(Runnable runnable) {
        return ((KeySupplier) runnable).getKey();
    }

    @Override
    public Runnable poll(@Nonnull KeySupplier keySupplier) {
        BlockingQueueHolder<Runnable> queue = queuePool().selectQueue(keySupplier.getKey());
        return queue.getQueue().poll();
    }
}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/impl/KeyQueue.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

import javax.annotation.Nonnull;

import com.github.bizyun.keythreadpool.KeySupplier;

/**
 * @author zhangbiyun
 */
interface KeyQueue<E> {
    E poll(@Nonnull KeySupplier keySupplier);
}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/impl/KeySupplierScanner.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

import java.lang.reflect.Field;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.Callable;

import org.apache.commons.lang3.reflect.FieldUtils;

import com.github.bizyun.keythreadpool.KeySupplier;

/**
 * @author zhangbiyun
 */
class KeySupplierScanner {

    private final Set<Object> scannedObjects = new HashSet<>();

    public KeySupplier scan(Object obj) {
        if (obj instanceof KeySupplier) {
            return (KeySupplier) obj;
        }

        for (Field field : FieldUtils.getAllFields(obj.getClass())) {
            KeySupplier keySupplier = doScan(obj, field, KeySupplier.class);
            if (keySupplier != null) {
                return keySupplier;
            }
            keySupplier = doScan(obj, field, Runnable.class);
            if (keySupplier != null) {
                return keySupplier;
            }
            keySupplier = doScan(obj, field, Callable.class);
            if (keySupplier != null) {
                return keySupplier;
            }
        }
        return null;
    }

    private <T> KeySupplier doScan(Object obj, Field field, Class<T> matchClass) {
        if (matchClass.isAssignableFrom(field.getType())) {
            boolean isAccessible = field.isAccessible();
            if (!isAccessible) {
                field.setAccessible(true);
            }
            T o;
            try {
                o = matchClass.cast(field.get(obj));
            } catch (IllegalAccessException e) {
                return null;
            } finally {
                if (!isAccessible) {
                    field.setAccessible(false);
                }
            }

            if (o instanceof KeySupplier) {
                return (KeySupplier) o;
            }
            if (!scannedObjects.contains(o)) {
                scannedObjects.add(o);
                return scan(o);
            }
        }
        return null;
    }
}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/impl/BlockingQueueProxy.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

import java.util.concurrent.BlockingQueue;

import com.google.common.util.concurrent.ForwardingBlockingQueue;

/**
 * @author zhangbiyun
 */
class BlockingQueueProxy<E> extends ForwardingBlockingQueue<E> {

    private final BlockingQueue<E> queue;

    public BlockingQueueProxy(BlockingQueue<E> queue) {
        this.queue = queue;
    }

    @Override
    protected BlockingQueue<E> delegate() {
        return queue;
    }
}


<|start_filename|>src/main/java/com/github/bizyun/LogUtils.java<|end_filename|>
package com.github.bizyun;

import org.slf4j.Logger;

/**
 * @author zhangbiyun
 */
public final class LogUtils {

    public static void debugLog(Logger logger, String log, Object... args) {
        if (logger.isDebugEnabled()) {
            logger.debug(log, args);
        }
    }
}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/KeyCallable.java<|end_filename|>
package com.github.bizyun.keythreadpool;

import java.util.concurrent.Callable;

/**
 * @author zhangbiyun
 */
public interface KeyCallable<V> extends Callable<V>, KeySupplier {
}


<|start_filename|>src/test/java/com/github/bizyun/keythreadpool/impl/KeySupplierScannerTest.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl;

import static org.junit.jupiter.api.Assertions.*;

import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

import org.junit.jupiter.api.Test;

import com.github.bizyun.keythreadpool.KeySupplier;
import com.github.bizyun.keythreadpool.impl.testtask.TestKeyRunnable;

/**
 * @author zhangbiyun
 * @date 2020/6/21
 */
class KeySupplierScannerTest {

    @Test
    void test1() {
        TestKeyRunnable testKeyRunner = new TestKeyRunnable(1, new AtomicInteger(), 1, new ConcurrentHashMap<>());
        Runnable r1 = () -> testKeyRunner.run();
        KeySupplier keySupplier = new KeySupplierScanner().scan(r1);
        assertNotNull(keySupplier);
        assertEquals(1, keySupplier.getKey());

        Runnable r2 = () -> r1.run();

        keySupplier = new KeySupplierScanner().scan(r2);
        assertNotNull(keySupplier);
        assertEquals(1, keySupplier.getKey());

        Callable<Void> c1 = () -> {
            r2.run();
            return null;
        };

        keySupplier = new KeySupplierScanner().scan(c1);
        assertNotNull(keySupplier);
        assertEquals(1, keySupplier.getKey());

        Runnable r3 = () -> {
            try {
                c1.call();
            } catch (Exception e) {
            }
        };
        keySupplier = new KeySupplierScanner().scan(r3);
        assertNotNull(keySupplier);
        assertEquals(1, keySupplier.getKey());

        Callable<Void> c2 = () -> c1.call();

        keySupplier = new KeySupplierScanner().scan(c2);
        assertNotNull(keySupplier);
        assertEquals(1, keySupplier.getKey());

    }

    @Test
    void test2() {
        TestRunnable t1 = new TestRunnable();
        TestRunnable t2 = new TestRunnable();
        t1.setR(t2);
        t2.setR(t1);

        KeySupplier keySupplier = new KeySupplierScanner().scan(t2);
        assertNull(keySupplier);
    }

    @Test
    void test3() {
        Runnable r1 = () -> new TestKeyRunnable(1, new AtomicInteger(), 1, new ConcurrentHashMap<>()).run();
        KeySupplier keySupplier = new KeySupplierScanner().scan(r1);
        assertNull(keySupplier);
    }

    private static class TestRunnable implements Runnable {
        private Runnable r;

        @Override
        public void run() {

        }

        public void setR(Runnable r) {
            this.r = r;
        }
    }

}

<|start_filename|>src/test/java/com/github/bizyun/keythreadpool/impl/testtask/TestKeyCallable.java<|end_filename|>
package com.github.bizyun.keythreadpool.impl.testtask;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicInteger;

import com.github.bizyun.keythreadpool.KeyCallable;

/**
 * @author zhangbiyun
 * @date 2020/6/21
 */
public class TestKeyCallable<Void> implements KeyCallable<Void> {

    private final long key;
    private final AtomicInteger count;
    private final ConcurrentHashMap<Long, LinkedBlockingQueue<Integer>> consumerRecords;
    private final int value;

    public TestKeyCallable(long key, AtomicInteger count,
            int value, ConcurrentHashMap<Long, LinkedBlockingQueue<Integer>> consumerRecords) {
        this.key = key;
        this.count = count;
        this.value = value;
        this.consumerRecords = consumerRecords;
    }

    @Override
    public long getKey() {
        return key;
    }

    @Override
    public Void call() throws Exception {
        count.incrementAndGet();
        LinkedBlockingQueue<Integer> records = consumerRecords.get(key);
        if (records == null) {
            records = consumerRecords.computeIfAbsent(key, k -> new LinkedBlockingQueue<>());
        }
        records.offer(value);
        return null;
    }
}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/KeyRunnable.java<|end_filename|>
package com.github.bizyun.keythreadpool;

/**
 * @author zhangbiyun
 */
public interface KeyRunnable extends Runnable, KeySupplier {
}


<|start_filename|>src/main/java/com/github/bizyun/keythreadpool/KeySupplier.java<|end_filename|>
package com.github.bizyun.keythreadpool;

/**
 * @author zhangbiyun
 */
public interface KeySupplier {
    long getKey();
}
