/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License;
 * you may not use this file except in compliance with the Elastic License.
 */
package org.elasticsearch.xpack.ml.job.process.autodetect;

import org.apache.logging.log4j.Logger;
import org.elasticsearch.ElasticsearchException;
import org.elasticsearch.common.logging.Loggers;
import org.elasticsearch.xpack.ml.action.OpenJobAction.JobTask;
import org.elasticsearch.xpack.ml.utils.ExceptionsHelper;

import java.io.IOException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;

/**
 * The process context that encapsulates the job task, the process state and the autodetect communicator.
 */
final class ProcessContext {

    private static final Logger LOGGER = Loggers.getLogger(ProcessContext.class);

    private final ReentrantLock lock = new ReentrantLock();
    private final JobTask jobTask;
    private volatile AutodetectCommunicator autodetectCommunicator;
    private volatile ProcessState state;

    ProcessContext(JobTask jobTask) {
        this.jobTask = jobTask;
        this.state = new ProcessNotRunningState();
    }

    JobTask getJobTask() {
        return jobTask;
    }

    AutodetectCommunicator getAutodetectCommunicator() {
        return autodetectCommunicator;
    }

    private void setAutodetectCommunicator(AutodetectCommunicator autodetectCommunicator) {
        this.autodetectCommunicator = autodetectCommunicator;
    }

    ProcessStateName getState() {
        return state.getName();
    }

    private void setState(ProcessState state) {
        this.state = state;
    }

    void tryLock() {
        try {
            if (lock.tryLock(NativeAutodetectProcessFactory.PROCESS_STARTUP_TIMEOUT.getSeconds(), TimeUnit.SECONDS) == false) {
                LOGGER.error("Failed to acquire process lock for job [{}]", jobTask.getJobId());
                throw ExceptionsHelper.serverError("Failed to acquire process lock for job [" + jobTask.getJobId() + "]");
            }
        } catch (InterruptedException e) {
            throw new ElasticsearchException(e);
        }
    }

    void unlock() {
        lock.unlock();
    }

    void setRunning(AutodetectCommunicator autodetectCommunicator) {
        assert lock.isHeldByCurrentThread();
        state.setRunning(this, autodetectCommunicator);
    }

    boolean setDying() {
        assert lock.isHeldByCurrentThread();
        return state.setDying(this);
    }

    KillBuilder newKillBuilder() {
        return new ProcessContext.KillBuilder();
    }

    class KillBuilder {
        private boolean awaitCompletion;
        private boolean finish;
        private boolean silent;
        private String reason;

        KillBuilder setAwaitCompletion(boolean awaitCompletion) {
            this.awaitCompletion = awaitCompletion;
            return this;
        }

        KillBuilder setFinish(boolean finish) {
            this.finish = finish;
            return this;
        }

        KillBuilder setSilent(boolean silent) {
            this.silent = silent;
            return this;
        }

        KillBuilder setReason(String reason) {
            this.reason = reason;
            return this;
        }

        void kill() {
            if (autodetectCommunicator == null) {
                return;
            }
            String jobId = jobTask.getJobId();

            if (silent == false) {
                String extraInfo = (state.getName() == ProcessStateName.DYING) ? " while closing" : "";
                if (reason == null) {
                    LOGGER.info("Killing job [{}]{}", jobId, extraInfo);
                } else {
                    LOGGER.info("Killing job [{}]{}, because [{}]", jobId, extraInfo, reason);
                }
            }
            try {
                autodetectCommunicator.killProcess(awaitCompletion, finish);
            } catch (IOException e) {
                LOGGER.error("[{}] Failed to kill autodetect process for job", jobId);
            }
        }
    }

    enum ProcessStateName {
        NOT_RUNNING, RUNNING, DYING
    }

    private interface ProcessState {
        /**
         * @return was a state change made?
         * */
        boolean setRunning(ProcessContext processContext, AutodetectCommunicator autodetectCommunicator);
        /**
         * @return was a state change made?
         */
        boolean setDying(ProcessContext processContext);
        ProcessStateName getName();
    }

    private static class ProcessNotRunningState implements ProcessState {
        @Override
        public boolean setRunning(ProcessContext processContext, AutodetectCommunicator autodetectCommunicator) {
            processContext.setAutodetectCommunicator(autodetectCommunicator);
            processContext.setState(new ProcessRunningState());
            return true;
        }

        @Override
        public boolean setDying(ProcessContext processContext) {
            processContext.setState(new ProcessDyingState());
            return true;
        }

        @Override
        public ProcessStateName getName() {
            return ProcessStateName.NOT_RUNNING;
        }
    }

    private static class ProcessRunningState implements ProcessState {

        @Override
        public boolean setRunning(ProcessContext processContext, AutodetectCommunicator autodetectCommunicator) {
            LOGGER.debug("Process set to [running] while it was already in that state");
            return false;
        }

        @Override
        public boolean setDying(ProcessContext processContext) {
            processContext.setState(new ProcessDyingState());
            return true;
        }

        @Override
        public ProcessStateName getName() {
            return ProcessStateName.RUNNING;
        }
    }

    private static class ProcessDyingState implements ProcessState {

        @Override
        public boolean setRunning(ProcessContext processContext, AutodetectCommunicator autodetectCommunicator) {
            LOGGER.debug("Process set to [running] while it was in [dying]");
            return false;
        }

        @Override
        public boolean setDying(ProcessContext processContext) {
            LOGGER.debug("Process set to [dying] while it was already in that state");
            return false;
        }

        @Override
        public ProcessStateName getName() {
            return ProcessStateName.DYING;
        }
    }
}

@@ -433,16 +433,67 @@
 oggers;%0A
+import org.elasticsearch.xpack.ml.MachineLearning;%0A
 import o
@@ -1829,54 +1829,45 @@
 ock(
-NativeAutodetectProcessFactory.PROCESS_STARTUP
+MachineLearning.STATE_PERSIST_RESTORE
 _TIM

Use the correct timeout for the process context lock (elastic/x-pack-elasticsearch#2935)