/**
 * 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.activemq.transport.nio;

import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.broker.BrokerService;
import org.apache.activemq.broker.TransportConnector;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.jms.Connection;
import javax.jms.Session;
import java.lang.management.ManagementFactory;
import java.lang.management.ThreadInfo;
import java.lang.management.ThreadMXBean;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import static junit.framework.TestCase.assertTrue;
import static org.junit.Assert.assertEquals;

public class NIOSSThreadLeakTest {

    private static final Logger LOG = LoggerFactory.getLogger(NIOSSThreadLeakTest.class);

    BrokerService broker;
    TransportConnector connector;

    public static final String KEYSTORE_TYPE = "jks";
    public static final String PASSWORD = "password";
    public static final String SERVER_KEYSTORE = "src/test/resources/server.keystore";
    public static final String TRUST_KEYSTORE = "src/test/resources/client.keystore";


    @Before
    public void setUp() throws Exception {
        System.setProperty("javax.net.ssl.trustStore", TRUST_KEYSTORE);
        System.setProperty("javax.net.ssl.trustStorePassword", PASSWORD);
        System.setProperty("javax.net.ssl.trustStoreType", KEYSTORE_TYPE);
        System.setProperty("javax.net.ssl.keyStore", SERVER_KEYSTORE);
        System.setProperty("javax.net.ssl.keyStoreType", KEYSTORE_TYPE);
        System.setProperty("javax.net.ssl.keyStorePassword", PASSWORD);
        System.setProperty("javax.net.ssl.keyStorePassword", PASSWORD);

        // ensure a new thread per 2 connections
        System.setProperty("org.apache.activemq.transport.nio.SelectorManager.maxChannelsPerWorker", "2");

        broker = new BrokerService();
        broker.setPersistent(false);
        broker.setUseJmx(false);
        connector = broker.addConnector("nio+ssl://localhost:0?transport.needClientAuth=true");
        broker.start();
        broker.waitUntilStarted();

    }

    @After
    public void tearDown() throws Exception {
        if (broker != null) {
            broker.stop();
            broker.waitUntilStopped();
        }
    }


    @Test(timeout=360000)
    public void testThreadUsage() throws Exception {
        final ActiveMQConnectionFactory factory = new ActiveMQConnectionFactory("nio+ssl://localhost:" + connector.getConnectUri().getPort());
        int threadNumber = 10;
        ExecutorService executorService = Executors.newFixedThreadPool(threadNumber);
        final CountDownLatch[] latch = new CountDownLatch[]{new CountDownLatch(threadNumber)};
        final AtomicInteger errors = new AtomicInteger(0);

        Runnable newConnection = new Runnable() {
            @Override
            public void run() {
                Connection conn = null;
                try {
                    conn = factory.createConnection();
                    conn.start();
                    Session sess = conn.createSession(false, Session.AUTO_ACKNOWLEDGE);
                    TimeUnit.MILLISECONDS.sleep(600);
                } catch (Exception e) {
                    e.printStackTrace();
                    errors.incrementAndGet();
                } finally {
                    try {
                        conn.close();
                    } catch (Exception e) {
                    }
                    latch[0].countDown();
                }
            }
        };
        for (int i = 0; i < threadNumber; i++) {
            executorService.submit(newConnection);
        }

        latch[0].await(5, TimeUnit.MINUTES);
        LOG.info("First batch...");
        assertTrue(noMoreThanXSelectorThreads(threadNumber, countSelectorThreads()));

        // repeat a few more batches to verify no leaks
        for (int i=0; i<4; i++) {
            LOG.info("batch..." + (i+2));
            latch[0] = new CountDownLatch(threadNumber);

            for (int j = 0; j < threadNumber; j++) {
                executorService.submit(newConnection);
            }

            latch[0].await(5, TimeUnit.MINUTES);
        }
        assertTrue(noMoreThanXSelectorThreads(threadNumber, countSelectorThreads()));

        LOG.info("errors " + errors.get());
        assertEquals(0, errors.get());
        executorService.shutdownNow();
    }

    private boolean noMoreThanXSelectorThreads(int expected, int i) {
        boolean result = i <= expected;
        if (!result) {
            LOG.error("too many threads: " + i);
        }
        return result;
    }

    private int countSelectorThreads() {
        int count = 0;
        ThreadMXBean thbean = ManagementFactory.getThreadMXBean();
        for (long id : thbean.getAllThreadIds()) {
            ThreadInfo threadInfo = thbean.getThreadInfo(id);
            if (threadInfo != null) {
                if (threadInfo.getThreadName().contains("Selector Worker:")) {
                    LOG.info(thbean.getThreadInfo(id).getThreadName());
                    count++;
                }
            }
        }
        return count;
    }
}
