com.alibaba.dts.client.executor.script.ShellJobProcessor Maven / Gradle / Ivy

Go to download
package com.alibaba.dts.client.executor.script;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.CountDownLatch;

import com.alibaba.dts.client.executor.job.processor.ExSimpleJobProcessor;
import com.alibaba.dts.client.executor.logcollector.AbstractLogCollector;
import com.alibaba.dts.client.executor.logcollector.StreamType;
import com.alibaba.dts.client.executor.simple.processor.SimpleJobContext;
import com.alibaba.dts.common.logger.SchedulerXLoggerFactory;
import com.alibaba.dts.common.logger.innerlog.Logger;
import com.alibaba.edas.schedulerx.ProcessResult;

/**
 *
 * @author xiaomeng.hxm
 */
public class ShellJobProcessor extends ExSimpleJobProcessor {
    private static final Logger LOGGER = SchedulerXLoggerFactory.getLogger(ShellJobProcessor.class);
    private Process shellProcess = null;
    private AbstractLogCollector logCollector = null;

    @Override
    public ProcessResult process(SimpleJobContext context) {
        boolean result = false;
        try {
            ProcessBuilder processBuilder = ShellUtils.createProcessBuilder(getContext(context));
            if (redirectStream()) {
                processBuilder.redirectErrorStream(true);
            }
            shellProcess = processBuilder.start();

            CountDownLatch countDownLatch = null;
            if (redirectStream()) {
                countDownLatch = new CountDownLatch(1);
                Thread stderrStreamProcessor = new ShellStreamProcessor(this, shellProcess.getInputStream(),
                        StreamType.STD_ERR, countDownLatch);
                stderrStreamProcessor.start();
            } else {
                countDownLatch = new CountDownLatch(2);
                Thread stdoutStreamProcessor = new ShellStreamProcessor(this, shellProcess.getInputStream(),
                        StreamType.STD_OUT, countDownLatch);
                stdoutStreamProcessor.start();

                Thread stderrStreamProcessor = new ShellStreamProcessor(this, shellProcess.getErrorStream(),
                        StreamType.STD_ERR, countDownLatch);
                stderrStreamProcessor.start();
            }

            countDownLatch.await();
            result = (shellProcess.waitFor() == 0);
        } catch (Exception e) {
            logCollector.collectStderr(e.getMessage());
        } finally {
            return new ProcessResult(result);
        }
    }

    protected String getContext(SimpleJobContext context) {
        return context.getJob().getJobArguments();
    }

    protected void processStdOutputStream(InputStream inputStream) {
        String line = null;
        try {
            BufferedReader br = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
            while ((line = br.readLine()) != null) {
                logCollector.collectStdout(line);
            }
        } catch (IOException e) {
            LOGGER.error("error ShellJobProcessor stdout stream", e);
        } finally {
            logCollector.collectStdout("", true);
        }
    }

    protected void processStdErrorStream(InputStream inputStream) {
        String line = null;
        try {
            BufferedReader br = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
            while ((line = br.readLine()) != null) {
                logCollector.collectStderr(line);
            }
        } catch (IOException e) {
            LOGGER.error("error ShellJobProcessor stderr stream", e);
        } finally {
            logCollector.collectStderr("", true);
        }
    }

    protected boolean redirectStream() {
        return false;
    }

    public boolean kill() {
        boolean result = true;
        try {
            ShellUtils.killProcess(ShellUtils.getPidOfProcess(shellProcess));
        } catch (Exception e) {
            result = false;
        }

        if (shellProcess != null) {
            shellProcess.destroy();
        }
        return result;
    }

    @Override
    public void preProcess(SimpleJobContext context) {
        logCollector = context.getLogCollector();
    }

    @Override
    public void postProcess(SimpleJobContext context) {
    }
}