/*
Copyright 2017 Red Hat, Inc.

Licensed 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.as.test.manualmode.suspend;

import static org.jboss.as.controller.descriptions.ModelDescriptionConstants.NAME;
import static org.jboss.as.controller.descriptions.ModelDescriptionConstants.OP;
import static org.jboss.as.controller.descriptions.ModelDescriptionConstants.READ_ATTRIBUTE_OPERATION;
import static org.jboss.as.controller.descriptions.ModelDescriptionConstants.SUSPEND_STATE;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.HttpURLConnection;
import java.net.SocketPermission;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import javax.inject.Inject;

import org.jboss.as.test.integration.common.HttpRequest;
import org.jboss.as.test.shared.PermissionUtils;
import org.jboss.as.test.shared.TestSuiteEnvironment;
import org.jboss.as.test.shared.TimeoutUtil;
import org.jboss.dmr.ModelNode;
import org.jboss.logging.Logger;
import org.jboss.msc.service.ServiceActivator;
import org.jboss.shrinkwrap.api.ShrinkWrap;
import org.jboss.shrinkwrap.api.asset.StringAsset;
import org.jboss.shrinkwrap.api.spec.JavaArchive;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.AssumptionViolatedException;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.wildfly.core.testrunner.ServerControl;
import org.wildfly.core.testrunner.ServerController;
import org.wildfly.core.testrunner.UnsuccessfulOperationException;
import org.wildfly.core.testrunner.WildflyTestRunner;
import org.wildfly.test.suspendresumeendpoint.SuspendResumeHandler;
import org.wildfly.test.suspendresumeendpoint.TestSuspendServiceActivator;
import org.wildfly.test.suspendresumeendpoint.TestUndertowService;
import org.xnio.IoUtils;

/**
 * WFCORE-3073. Tests that a softly killed process shuts down gracefully.
 *
 * @author Brian Stansberry
 */
@RunWith(WildflyTestRunner.class)
@ServerControl(manual = true)
public class SuspendOnSoftKillTestCase {

    private static final String WEB_SUSPEND_JAR = "web-suspend.jar";

    @SuppressWarnings("unused")
    @Inject
    private static ServerController serverController;

    private String jbossArgs;
    private ProcessUtil processUtil;

    @After
    public void stopContainer() {
        try {
            // The os kills we use in the test don't result in serverController
            // being in an internal state that reflects its server has stopped.
            // So, regardless of whether we think it was killed, stop the container
            // If we did do the kill, this puts the serverController in the correct state.
            // And, if the test failed before we got to the kill, this will do
            // the necessary stop.
            serverController.stop();
        } finally {
            if (jbossArgs == null) {
                System.clearProperty("jboss.args");
            } else {
                System.setProperty("jboss.args", jbossArgs);
            }
        }
    }

    @AfterClass
    public static void cleanDeployment() {
        // When the test methods are done, start up the server again so we can clean out the deployment
        try {
            serverController.start();
            serverController.undeploy(WEB_SUSPEND_JAR);
        } catch (Exception e) {
            Logger.getLogger(SuspendOnSoftKillTestCase.class.getName()).error("Failed undeploying in stopContainer", e);
        } finally {
            // Stop the container
            serverController.stop();
        }
    }

    @Test
    public void testSuspendOnSoftKillPositiveTimeout() throws Exception {
        startContainer(15);
        suspendOnSoftKillTest();
    }

    @Test
    public void testSuspendOnSoftKillNegativeTimeout() throws Exception {
        startContainer(-1);
        suspendOnSoftKillTest();
    }

    private void startContainer(int suspendTimeout) throws Exception {

        suspendTimeout = suspendTimeout < 1 ? suspendTimeout : TimeoutUtil.adjust(suspendTimeout);

        jbossArgs = System.getProperty("jboss.args");
        String newArgs = jbossArgs == null ? "" : jbossArgs;
        newArgs += " -D[Standalone] -Dorg.wildfly.sigterm.suspend.timeout=" + suspendTimeout;
        System.setProperty("jboss.args", newArgs);

        if (File.pathSeparatorChar == ':'){
            processUtil = new UnixProcessUtil();
        } else {
            processUtil = new WindowsProcessUtil();
        }

        // Start the server
        serverController.start();

        // If there's already the deployment there from a previous test, remove it
        try {
            serverController.undeploy(WEB_SUSPEND_JAR);
        } catch (Exception ignored) {
            // assume it wasn't deployed. if it was we'll fail below
        }

        JavaArchive war = ShrinkWrap.create(JavaArchive.class, WEB_SUSPEND_JAR);
        war.addPackage(SuspendResumeHandler.class.getPackage());
        war.addAsServiceProvider(ServiceActivator.class, TestSuspendServiceActivator.class);
        war.addAsResource(new StringAsset("Dependencies: org.jboss.dmr, org.jboss.as.controller, io.undertow.core, org.jboss.as.server,org.wildfly.extension.request-controller, org.jboss.as.network\n"), "META-INF/MANIFEST.MF");
        war.addAsManifestResource(PermissionUtils.createPermissionsXmlAsset(
                new RuntimePermission("createXnioWorker"),
                new SocketPermission(TestSuiteEnvironment.getServerAddress() + ":8080", "listen,resolve"),
                new SocketPermission("*", "accept,resolve")
        ), "permissions.xml");
        serverController.deploy(war, WEB_SUSPEND_JAR);
    }

    private void suspendOnSoftKillTest() throws Exception {

        final String address = "http://" + TestSuiteEnvironment.getServerAddress() + ":8080/web-suspend";
        ExecutorService executorService = Executors.newSingleThreadExecutor();

        try {
            // Send a request that will block
            Future<Object> result = executorService.submit(new Callable<Object>() {
                @Override
                public Object call() throws Exception {
                    return HttpRequest.get(address, 60, TimeUnit.SECONDS);
                }
            });

            Thread.sleep(TimeoutUtil.adjust(1000)); //nasty, but we need to make sure the HTTP request has started

            softKillServer();

            ModelNode op = new ModelNode();
            op.get(OP).set(READ_ATTRIBUTE_OPERATION);
            op.get(NAME).set(SUSPEND_STATE);

            String suspendState;
            long timeout = System.currentTimeMillis() + TimeoutUtil.adjust(10000);
            do {
                suspendState = serverController.getClient().executeForResult(op).asString();
                if (!"SUSPENDING".equals(suspendState)) {
                    Thread.sleep(50);
                } else {
                    break;
                }
            } while (System.currentTimeMillis() < timeout);
            if ("RUNNING".equals(suspendState) && processUtil instanceof WindowsProcessUtil) {
                // taskkill with no /f doesn't seem to soft kill Windows vms. I can't find
                // any other way to test this than trying taskkill without the hard-kill /f arg.
                // I tried wmic and terminate but that results in a hard kill.
                // So just ignore this test. The point of this test is to test how the shutdown hook
                // we install handles things. If the hook is never invoked by the VM we
                // can't do anything about that anyway.
                // I let the test run all the way to here before ignoring so that if
                // taskkill happens to actually work in some Windows environment, then
                // we get the coverage
                throw new AssumptionViolatedException("taskkill did not produce a 'soft kill' of the server on Windows");
            }
            Assert.assertEquals("SUSPENDING", suspendState);

            final HttpURLConnection conn = (HttpURLConnection) new URL(address).openConnection();
            try {
                conn.setDoInput(true);
                int responseCode = conn.getResponseCode();
                Assert.assertEquals(503, responseCode);
            } finally {
                conn.disconnect();
            }

            // Send a request that will trigger the first request to complete
            HttpRequest.get(address + "?" + TestUndertowService.SKIP_GRACEFUL + "=true", TimeoutUtil.adjust(30), TimeUnit.SECONDS);
            // Confirm 1st request completed
            Assert.assertEquals(SuspendResumeHandler.TEXT, result.get());

            // Check server behaves like a suspended, suspending or stopped server; which it will be is unknown.
            // If we are still in suspending state, wait a bit longer until get suspended or the server is down.
            try {
                timeout = System.currentTimeMillis() + TimeoutUtil.adjust(10000);
                do {
                    suspendState = serverController.getClient().executeForResult(op).asString();
                    if ("SUSPENDING".equals(suspendState)) {
                        Thread.sleep(50);
                    } else {
                        break;
                    }
                } while (System.currentTimeMillis() < timeout);
                Assert.assertEquals("SUSPENDED", suspendState);
            } catch (UnsuccessfulOperationException | RuntimeException ok) {
                // ignore; it's fine if the server's down
            }
        } finally {
            executorService.shutdown();
        }

        waitForServerShutdown();
    }

    private void softKillServer() {
        List<RunningProcess> runningProcesses = processUtil.getRunningProcesses();
        Assert.assertEquals("Incorrect number of processes -- " + runningProcesses, 1, runningProcesses.size());
        processUtil.softKillProcess(runningProcesses.get(0));
    }

    private void waitForServerShutdown() throws Exception {
        // Server should be suspended already with no in-flight requests, so it should stop quickly.
        // But we're willing to be patient and wait to be sure it stops even on a slow system.
        // The point is just to make sure it eventually stops, not to check how fast
        final long time = System.currentTimeMillis() + TimeoutUtil.adjust(20000);
        List<RunningProcess> runningProcesses;
        do {
            runningProcesses = processUtil.getRunningProcesses();
            if (runningProcesses.size() == 0){
                return;
            }
            Thread.sleep(200);
        } while(System.currentTimeMillis() < time);
        Assert.fail("Did not have all running processes " + runningProcesses);
    }

    private abstract static class ProcessUtil {

        List<String> initialProcessIds;

        ProcessUtil(){
            initialProcessIds = getInitialProcessIds();
        }

        List<String> getInitialProcessIds(){
            List<String> processes = listProcesses();
            List<String> ids = new ArrayList<String>();
            for (String proc : processes){
                ids.add(parseProcessId(proc));
            }
            return ids;
        }

        String parseProcessId(String proc){
            proc = proc.trim();
            int i = proc.indexOf(' ');
            return proc.substring(0, i);
        }

        List<RunningProcess> getRunningProcesses(){
            List<RunningProcess> running = new ArrayList<RunningProcess>();
            List<String> processes = listProcesses();
            for (String proc : processes){
                String id = parseProcessId(proc);
                if (!initialProcessIds.contains(id)){
                    if (proc.contains("-D[Standalone]")){
                        running.add(new RunningProcess(id, "Standalone"));
                    }
                }
            }
            return running;
        }

        List<String> listProcesses() {
            final Process p;
            try {
                p = Runtime.getRuntime().exec(getJpsCommand());
            } catch (IOException e) {
                throw new RuntimeException(e);
            }

            List<String> processes = new ArrayList<String>();
            BufferedReader input = new BufferedReader(new InputStreamReader(p.getInputStream()));
            try {
                String line;
                while ((line = input.readLine()) != null) {
                    if (line.contains("jboss-modules.jar")) {
                        processes.add(line);
                    }

                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            } finally {
                IoUtils.safeClose(input);
            }
            return processes;
        }

        void softKillProcess(RunningProcess process) {
            try {
                Runtime.getRuntime().exec(getSoftKillCommand(process));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }

        abstract String[] getJpsCommand();

        abstract String getSoftKillCommand(RunningProcess process);
    }

    private static class UnixProcessUtil extends ProcessUtil {
        @Override
        String[] getJpsCommand() {
            final File jreHome = new File(System.getProperty("java.home"));
            Assert.assertTrue("JRE home not found. File: " + jreHome.getAbsoluteFile(), jreHome.exists());
            if (TestSuiteEnvironment.isIbmJvm()) {
                return new String[] { "sh", "-c", "ps -ef | awk '{$1=\"\"; print $0}'" };
            } else {
                File jpsExe = new File(jreHome, "bin/jps");
                if (!jpsExe.exists()) {
                    jpsExe = new File(jreHome, "../bin/jps");
                }
                Assert.assertTrue("JPS executable not found. File: " + jpsExe, jpsExe.exists());
                return new String[] { jpsExe.getAbsolutePath(), "-l", "-v" };
            }
        }

        @Override
        String getSoftKillCommand(RunningProcess process) {
            return "kill -15 " + process.getProcessId();
        }
    }

    private static class WindowsProcessUtil extends ProcessUtil {

        @Override
        String[] getJpsCommand() {
            final File jreHome = new File(System.getProperty("java.home"));
            Assert.assertTrue("JRE home not found. File: " + jreHome.getAbsoluteFile(), jreHome.exists());
            File jpsExe = new File(jreHome, "bin/jps.exe");
            if (!jpsExe.exists()) {
                jpsExe = new File(jreHome, "../bin/jps.exe");
            }
            Assert.assertTrue("JPS executable not found. File: " + jpsExe, jpsExe.exists());
            return new String[] { jpsExe.getAbsolutePath(), "-l", "-v" };
        }

        @Override
        String getSoftKillCommand(RunningProcess process) {
            return "taskkill /pid " + process.getProcessId();
        }
    }

    private static class RunningProcess {
        final String processId;
        final String process;

        private RunningProcess(String processId, String process) {
            this.processId = processId;
            this.process = process;
        }

        String getProcessId() {
            return processId;
        }

        @Override
        public int hashCode() {
            return processId.hashCode();
        }

        @Override
        public boolean equals(Object obj) {
            return obj instanceof RunningProcess && ((RunningProcess) obj).processId.equals(processId);
        }

        @Override
        public String toString() {
            return "Process{id=" + processId + ", process=" + process + "}";
        }
    }
}
