package com.trip.tools.webdriver;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * 任务中心，每个queueName使用一个队列.
 * 
 */
public class MultiQueueCenter<T> {

    /**
     * 每个queueName一个队列.
     */
    private final ConcurrentHashMap<String, LinkedBlockingQueue<T>> queueMap = new ConcurrentHashMap<String, LinkedBlockingQueue<T>>();

    /**
     * 根据queueName取得对应的队列.
     * 
     * @param queueName
     * @return
     */
    public LinkedBlockingQueue<T> getQueue(String queueName) {

        // 存在直接返回
        if (queueMap.containsKey(queueName)) {
            return queueMap.get(queueName);
        }

        // 避免并发put同名的队列
        // putIfAbsent:不存在则放入并返回null，存在则返回对应value。线程安全
        LinkedBlockingQueue<T> queue = new LinkedBlockingQueue<T>();
        LinkedBlockingQueue<T> rtnQueue = queueMap
                .putIfAbsent(queueName, queue);
        if (rtnQueue != null) {
            return rtnQueue;
        }
        return queue;
    }

    /**
     * 放入一条消息（队列满的话，等待）.
     * 
     * @param msg
     *            消息内容
     */
    public void put(String queueName, T t) {
        if (t == null) {
            return;
        }

        try {
            getQueue(queueName).put(t);
        } catch (InterruptedException e) {
        }
    }

    /**
     * 取走一条消息（队列无消息，返回null）.
     * 
     * @return 消息内容
     */
    public T poll(String queueName) {

        if (queueMap.containsKey(queueName)) {
            return queueMap.get(queueName).poll();
        }
        return null;
    }

    /**
     * 清空指定队列.
     */
    public void clearQueue(String queueName) {
        if (queueMap.containsKey(queueName)) {
            queueMap.get(queueName).clear();
        }
    }

    /**
     * 返回队列名称.
     */
    public Set<String> getQueueNameSet() {
        return queueMap.keySet();
    }

    /**
     * 每个子队列的大小.
     */
    public Map<String, Integer> getSizeMap() {
        Map<String, Integer> sizeMap = new HashMap<String, Integer>();
        for (String key : queueMap.keySet()) {
            sizeMap.put(key, queueMap.get(key).size());
        }
        return sizeMap;
    }
}

package thread.queue2;

import java.util.Random;

public class ExecutorTest {

    public static MultiQueueCenter<String> center = new MultiQueueCenter<String>();

    public static void main(String[] args) {

        for (int i = 0; i < 3; i++) {
            new putThread222(i).start();
        }
        new monitorThread().start();

        while (true) {
            System.out.println("-----get-begin------");
            for (int i = 0; i < 3; i++) {
                String msg = ExecutorTest.center.poll("" + i);
                System.out.println("" + i + ":" + msg);
            }

            try {
                Thread.sleep(500);
            } catch (InterruptedException e) {
            }
        }
    }

}

class monitorThread extends Thread {

    public void run() {
        while (true) {
            System.out.println("queue:" + ExecutorTest.center.getSizeMap());
            try {
                Thread.sleep(100);
            } catch (InterruptedException e) {
            }
        }
    }
}

class putThread222 extends Thread {

    private int n = 0;

    putThread222(int m) {
        n = m;
    }

    public void run() {

        while (true) {
            // 随机延迟
            try {
                int num = 50 + new Random().nextInt(1000);
                Thread.sleep(num);
                ExecutorTest.center.put("" + n, "" + num);

            } catch (InterruptedException e) {
            }
        }

        // System.out.println("put thread" + this.n + " end.");
    }
}

package com.trip.tools.webdriver;

import java.util.Set;

import javax.servlet.ServletContextEvent;
import javax.servlet.ServletContextListener;
import javax.servlet.annotation.WebListener;

import org.openqa.selenium.WebDriver;
import org.openqa.selenium.chrome.ChromeDriver;
import org.openqa.selenium.chrome.ChromeDriverService;
import org.openqa.selenium.chrome.ChromeOptions;

/**
 * WebDriver的池，池里只放空闲的浏览器，拿出使用，用完放回.
 */
@WebListener
public class DriverPool implements ServletContextListener {

    private static final MultiQueueCenter<WebDriver> driverQueue = new MultiQueueCenter<WebDriver>();

    /**
     * 取得一个空闲的WebDriver.
     * 
     * @param queueName
     *            队列名称
     * @return
     */
    public static WebDriver getOneFreeDriver(String queueName) {

        // 从队列里取
        WebDriver driver = driverQueue.poll(queueName);

        // 没取到，创建一个
        if (driver == null) {

            // 懒得写配置，写死路径
            String driverPath = "D:\\chromedriver.exe";

            // 判断系统
            String os = System.getProperties().getProperty("os.name");
            if (os != null && os.toLowerCase().indexOf("linux") > -1) {
                driverPath = "/opt/webDriver/chromedriver";
            }

            // 设置Driver路径
            System.setProperty("webdriver.chrome.driver", driverPath);

            // 创建一个Driver对象
            ChromeDriverService service = ChromeDriverService
                    .createDefaultService();
            ChromeOptions options = new ChromeOptions();
            options.addArguments("--headless", "--disable-gpu", "--no-sandbox");
            driver = new ChromeDriver(service, options);
        }
        return driver;
    }

    /**
     * 归还WebDriver对象.
     * 
     * @param queueName
     *            队列名称
     * @param driver
     *            WebDriver对象
     */
    public static void putBackDriver(String queueName, WebDriver driver) {

        // 放回队列
        driverQueue.put(queueName, driver);

        // 队列初始为空，随着归还会越来越大，当使用数量与队列里数量一致时，达到平衡，不再增加
        // 想空闲时减小池的大小，需要另外写一个线程，根据空闲时间来关闭WebDriver

    }

    @Override
    public void contextInitialized(ServletContextEvent sce) {
        System.out.println("listener init.");
    }

    @Override
    public void contextDestroyed(ServletContextEvent sce) {

        // 循环每个队列
        Set<String> queueSet = driverQueue.getQueueNameSet();
        for (String queueName : queueSet) {
            System.out.println("try close queue:" + queueName);

            // 取空队列，都关闭
            WebDriver driver = driverQueue.poll(queueName);
            while (driver != null) {
                driver.quit();
                driver = driverQueue.poll(queueName);
            }
        }
    }

}