/*
* 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.pig.scripting;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.Reader;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.pig.ExecType;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.MRConfiguration;
import org.apache.pig.impl.PigImplConstants;
import org.apache.pig.impl.util.UDFContext;
import com.google.common.base.Charsets;
/**
* This class helps a scripting UDF capture user output by managing when to capture output and where
* the output is written to.
*
* For illustrate, we will only capture output for the last run (with the final set of data) and
* we need to keep track of the file containing that output for returning w/ the illustrate results.
* For runs, all standard output is written to the user logs.
*/
public class ScriptingOutputCapturer {
private static Log log = LogFactory.getLog(ScriptingOutputCapturer.class);
private static Map<String, String> outputFileNames = new HashMap<String, String>();
private static String runId = UUID.randomUUID().toString(); //Unique ID for this run to ensure udf output files aren't corrupted from previous runs.
//Illustrate will set the static flag telling udf to start capturing its output. It's up to each
//instance to react to it and set its own flag.
private static boolean captureOutput = false;
private boolean instancedCapturingOutput = false;
private ExecType execType;
public ScriptingOutputCapturer(ExecType execType) {
this.execType = execType;
}
public String getStandardOutputRootWriteLocation() throws IOException {
Configuration conf = UDFContext.getUDFContext().getJobConf();
String jobId = conf.get(MRConfiguration.JOB_ID);
String taskId = conf.get(MRConfiguration.TASK_ID);
String hadoopLogDir = System.getProperty("yarn.app.container.log.dir");
if (hadoopLogDir == null) {
hadoopLogDir = conf.get("yarn.app.container.log.dir");
}
if (hadoopLogDir == null) {
hadoopLogDir = System.getProperty("hadoop.log.dir");
}
if (hadoopLogDir == null) {
hadoopLogDir = conf.get("hadoop.log.dir");
}
String tmpDir = conf.get("hadoop.tmp.dir");
boolean fallbackToTmp = (hadoopLogDir == null);
if (!fallbackToTmp) {
try {
if (!(new File(hadoopLogDir).canWrite())) {
fallbackToTmp = true;
}
}
catch (SecurityException e) {
fallbackToTmp = true;
}
finally {
if (fallbackToTmp)
log.warn(String.format("Insufficient permission to write into %s. Change path to: %s", hadoopLogDir, tmpDir));
}
}
if (fallbackToTmp) {
hadoopLogDir = tmpDir;
}
log.debug("JobId: " + jobId);
log.debug("TaskId: " + taskId);
log.debug("hadoopLogDir: " + hadoopLogDir);
if (execType.isLocal() || conf.getBoolean(PigImplConstants.CONVERTED_TO_FETCH, false)) {
String logDir = System.getProperty("pig.udf.scripting.log.dir");
if (logDir == null)
logDir = ".";
return logDir + "/" + (taskId == null ? "" : (taskId + "_"));
} else {
String taskLogDir = getTaskLogDir(jobId, taskId, hadoopLogDir);
return taskLogDir + "/";
}
}
public String getTaskLogDir(String jobId, String taskId, String hadoopLogDir) throws IOException {
String taskLogDir;
String defaultUserLogDir = hadoopLogDir + File.separator + "userlogs";
if ( new File(defaultUserLogDir + File.separator + jobId).exists() ) {
taskLogDir = defaultUserLogDir + File.separator + jobId + File.separator + taskId;
} else if ( new File(defaultUserLogDir + File.separator + taskId).exists() ) {
taskLogDir = defaultUserLogDir + File.separator + taskId;
} else if ( new File(defaultUserLogDir).exists() ){
taskLogDir = defaultUserLogDir;
} else {
taskLogDir = hadoopLogDir + File.separator + "udfOutput";
File dir = new File(taskLogDir);
dir.mkdirs();
if (!dir.exists()) {
throw new IOException("Could not create directory: " + taskLogDir);
}
}
return taskLogDir;
}
public static void startCapturingOutput() {
ScriptingOutputCapturer.captureOutput = true;
}
public static Map<String, String> getUdfOutput() throws IOException {
Map<String, String> udfFuncNameToOutput = new HashMap<String,String>();
for (Map.Entry<String, String> funcToOutputFileName : outputFileNames.entrySet()) {
StringBuffer udfOutput = new StringBuffer();
FileInputStream fis = new FileInputStream(funcToOutputFileName.getValue());
Reader fr = new InputStreamReader(fis, Charsets.UTF_8);
BufferedReader br = new BufferedReader(fr);
try {
String line = br.readLine();
while (line != null) {
udfOutput.append("\t" + line + "\n");
line = br.readLine();
}
} finally {
br.close();
}
udfFuncNameToOutput.put(funcToOutputFileName.getKey(), udfOutput.toString());
}
return udfFuncNameToOutput;
}
public void registerOutputLocation(String functionName, String fileName) {
outputFileNames.put(functionName, fileName);
}
public static String getRunId() {
return runId;
}
public static boolean isClassCapturingOutput() {
return ScriptingOutputCapturer.captureOutput;
}
public boolean isInstanceCapturingOutput() {
return this.instancedCapturingOutput;
}
public void setInstanceCapturingOutput(boolean instanceCapturingOutput) {
this.instancedCapturingOutput = instanceCapturingOutput;
}
}