com.baifendian.swordfish.execserver.job.ProcessJob.java Source code

Java tutorial

Introduction

Here is the source code for com.baifendian.swordfish.execserver.job.ProcessJob.java

Source

/*
 * Copyright (C) 2017 Baifendian Corporation
 *
 * 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 com.baifendian.swordfish.execserver.job;

import com.baifendian.swordfish.execserver.exception.ExecException;
import com.baifendian.swordfish.execserver.exception.ExecTimeoutException;
import com.baifendian.swordfish.execserver.utils.Constants;
import com.baifendian.swordfish.execserver.utils.ProcessUtil;
import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.reflect.Field;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.function.BooleanSupplier;
import java.util.function.Consumer;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;

/**
 *  shell ?? job
 */
public class ProcessJob {

    /**
     *  process 
     */
    private ProcessBuilder processBuilder;

    /**
     * 
     */
    private Process process;

    /**
     * ?
     */
    private Consumer<List<String>> logHandler;

    /**
     * ??
     **/
    private BooleanSupplier isCompleted;

    /**
     * ?
     */
    private boolean isLongJob;

    /**
     * 
     */
    private String workDir;

    /**
     * ? id,  workDir ? command ??
     */
    private String jobAppId;

    /**
     * ???
     */
    private String proxyUser;

    /**
     * 
     */
    private String envFile;

    /**
     * ?
     */
    private Date startTime;

    /**
     * 
     */
    private int timeout;

    /**
     * 
     */
    private Logger logger;

    /**
     * 
     */
    private final List<String> logs;

    public ProcessJob(Consumer<List<String>> logHandler, BooleanSupplier isCompleted, boolean isLongJob,
            String workDir, String jobAppId, String proxyUser, String envFile, Date startTime, int timeout,
            Logger logger) {
        this.logHandler = logHandler;
        this.isCompleted = isCompleted;
        this.isLongJob = isLongJob;
        this.workDir = workDir;
        this.jobAppId = jobAppId;
        this.proxyUser = proxyUser;
        this.envFile = envFile;
        this.startTime = startTime;
        this.timeout = timeout;
        this.logger = logger;
        this.logs = Collections.synchronizedList(new ArrayList<>());
    }

    /**
     * s <p>
     *
     * @return 
     */
    private long calcNodeTimeout() {
        long usedTime = (System.currentTimeMillis() - startTime.getTime()) / 1000;
        long remainTime = timeout - usedTime;

        if (remainTime <= 0 && !isLongJob) {
            throw new ExecTimeoutException("workflow or streaming job execution time out");
        }

        return remainTime;
    }

    /**
     * ?, ? shell 
     *
     * @param command ?
     * @return ??, 0 ?, 
     */
    public int runCommand(String command) {
        // , 
        long remainTime = calcNodeTimeout();
        int exitCode;

        try {
            // ?
            processBuilder = new ProcessBuilder();

            // ? job ?
            if (StringUtils.isEmpty(command)) {
                exitCode = 0;
                return exitCode;
            }

            // ?
            String commandFile = String.format("%s/%s.command", workDir, jobAppId);

            logger.info("proxy user:{}, work dir:{}", proxyUser, workDir);

            // ?, ??
            if (!Files.exists(Paths.get(commandFile))) {
                logger.info("generate command file:{}", commandFile);

                StringBuilder stringBuilder = new StringBuilder();
                stringBuilder.append("#!/bin/sh\n");
                stringBuilder.append("BASEDIR=$(cd `dirname $0`; pwd)\n");
                stringBuilder.append("cd $BASEDIR\n");

                if (envFile != null) {
                    stringBuilder.append("source " + envFile + "\n");
                }

                stringBuilder.append("\n\n");
                stringBuilder.append(command);

                // ?
                FileUtils.writeStringToFile(new File(commandFile), stringBuilder.toString(),
                        Charset.forName("UTF-8"));
            }

            // ?

            processBuilder.command("sudo", "-u", proxyUser, "sh", commandFile);

            // 
            processBuilder.directory(new File(workDir));

            //  error ? merge ?
            processBuilder.redirectErrorStream(true);
            process = processBuilder.start();

            // ??
            printCommand(processBuilder);

            // ??
            readProcessOutput();

            int pid = getProcessId(process);

            logger.info("Process start, process id is: {}", pid);

            // 
            if (isLongJob) {
                // ?, , ??
                // ?,  10 , ?
                while (!isCompleted.getAsBoolean() && process.isAlive()) {
                    Thread.sleep(3000);
                }

                logger.info("streaming job has exit, work dir:{}, pid:{}", workDir, pid);

                // ?,  storm, ????
                exitCode = (isCompleted.getAsBoolean()) ? 0 : -1;
            } else {// ?
                boolean status = process.waitFor(remainTime, TimeUnit.SECONDS);

                if (status) {
                    exitCode = process.exitValue();
                    logger.info("job has exit, work dir:{}, pid:{}", workDir, pid);
                } else {
                    cancel();
                    exitCode = -1;
                    logger.info("job has timeout, work dir:{}, pid:{}", workDir, pid);
                }
            }
        } catch (InterruptedException e) {
            logger.error("interrupt exception, maybe task has been cancel or killed.", e);
            exitCode = -1;
            throw new ExecException("Process has been interrupted. Exit code is " + exitCode);
        } catch (Exception e) {
            logger.error(e.getMessage(), e);
            exitCode = -1;
            throw new ExecException("Process error. Exit code is " + exitCode);
        }

        return exitCode;
    }

    /**
     * ? shell 
     */
    public void cancel() throws Exception {
        if (process == null) {
            return;
        }

        // ?
        clear();

        int processId = getProcessId(process);

        logger.info("cancel process: {}", processId);

        // kill, ?
        boolean killed = softKill(processId);

        if (!killed) {
            // 
            hardKill(processId);

            // destory
            process.destroy();

            //  null
            process = null;
        }
    }

    /**
     * @param processId  id
     */
    private boolean softKill(int processId) throws InterruptedException {
        // ? 0 ?, 
        if (processId != 0 && process.isAlive()) {
            try {
                // ? sudo -u user command ?, ? user ? kill 
                String cmd = String.format("sudo kill %d", processId);

                logger.info("softkill job:{}, process id:{}, cmd:{}", jobAppId, processId, cmd);

                Runtime.getRuntime().exec(cmd);
            } catch (IOException e) {
                logger.info("kill attempt failed.", e);
            }
        }

        return process.isAlive();
    }

    /**
     *  kill
     *
     * @param processId  id
     */
    private void hardKill(int processId) {
        // ? 0 ?, 
        if (processId != 0 && process.isAlive()) {
            try {
                String cmd = String.format("sudo kill -9 %d", processId);

                logger.info("hardKill job:{}, process id:{}, cmd:{}", jobAppId, processId, cmd);

                Runtime.getRuntime().exec(cmd);
            } catch (IOException e) {
                logger.error("Kill attempt failed.", e);
            }
        }
    }

    /**
     * ?
     *
     * @param processBuilder 
     */
    private void printCommand(ProcessBuilder processBuilder) {
        String cmdStr;

        try {
            cmdStr = ProcessUtil.genCmdStr(processBuilder.command());
            logger.info("job run command:\n{}", cmdStr);
        } catch (IOException e) {
            logger.error(e.getMessage(), e);
        }
    }

    /**
     *  id
     *
     * @param process ?
     * @return  id
     */
    private int getProcessId(Process process) {
        int processId = 0;

        try {
            Field f = process.getClass().getDeclaredField("pid");
            f.setAccessible(true);

            processId = f.getInt(process);
        } catch (Throwable e) {
            logger.error(e.getMessage(), e);
        }

        return processId;
    }

    /**
     * ?
     */
    private void clear() {
        if (!logs.isEmpty()) {
            // ?
            logHandler.accept(logs);

            logs.clear();
        }
    }

    /**
     * ?, ?
     */
    private void readProcessOutput() {
        String threadLoggerInfoName = String.format("LoggerInfo-%s", jobAppId);

        Thread loggerInfoThread = new Thread(() -> {
            BufferedReader reader = null;

            try {
                reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
                String line;

                long preFlushTime = System.currentTimeMillis();

                while ((line = reader.readLine()) != null) {
                    logs.add(line);

                    long now = System.currentTimeMillis();

                    // ??
                    if (logs.size() >= Constants.defaultLogBufferSize
                            || now - preFlushTime > Constants.defaultLogFlushInterval) {
                        preFlushTime = now;

                        // ?
                        logHandler.accept(logs);

                        logs.clear();
                    }
                }
            } catch (Exception e) {
                // Do Nothing
            } finally {
                // , 
                if (!logs.isEmpty()) {
                    // ?
                    logHandler.accept(logs);

                    logs.clear();
                }

                if (reader != null) {
                    try {
                        reader.close();
                    } catch (IOException e) {
                        logger.error(e.getMessage(), e);
                    }
                }
            }
        }, threadLoggerInfoName);

        try {
            loggerInfoThread.setDaemon(true);
            loggerInfoThread.start();
        } catch (Exception e) {
            logger.error(e.getMessage(), e);
        }
    }
}