diff --git a/gobblin-utility/src/main/java/gobblin/util/HeapDumpForTaskUtils.java b/gobblin-utility/src/main/java/gobblin/util/HeapDumpForTaskUtils.java new file mode 100644 index 00000000000..190740b8cd7 --- /dev/null +++ b/gobblin-utility/src/main/java/gobblin/util/HeapDumpForTaskUtils.java @@ -0,0 +1,91 @@ +/* (c) 2014 LinkedIn Corp. All rights reserved. + * + * 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. + */ + +package gobblin.util; + +import java.io.BufferedWriter; +import java.io.IOException; +import java.io.OutputStreamWriter; +import java.nio.charset.Charset; + +import org.apache.hadoop.fs.FileSystem; +import org.apache.hadoop.fs.Path; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import com.google.common.io.Closer; + +import gobblin.configuration.ConfigurationKeys; + + +/** + * A utility class for generating script to move the heap dump .prof files to HDFS for hadoop tasks, when Java heap out of memory error is thrown. + */ +public class HeapDumpForTaskUtils { + + private static final Logger LOG = LoggerFactory.getLogger(HeapDumpForTaskUtils.class); + private static final String DUMP_FOLDER = "dumps"; + + /** + * Generate the dumpScript, which is used when OOM error is thrown during task execution. + * The current content dumpScript puts the .prof files to the DUMP_FOLDER within the same directory of the dumpScript. + * + * User needs to add the following options to the task java.opts: + * + * -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=./heapFileName.hprof -XX:OnOutOfMemoryError=./dumpScriptFileName + * + * @param dumpScript The path to the dumpScript, which needs to be added to the Distributed cache. + * To use it, simply put the path of dumpScript to the gobblin config: job.hdfs.files. + * @param fs File system + * @param heapFileName the name of the .prof file. + * @param chmod chmod for the dump script. For hdfs file, e.g, "hadoop fs -chmod 755" + * @throws IOException + */ + public static void generateDumpScript(Path dumpScript, FileSystem fs, String heapFileName, String chmod) + throws IOException { + if (fs.exists(dumpScript)) { + LOG.info("Heap dump script already exists: " + dumpScript); + return; + } + + Closer closer = Closer.create(); + try { + Path dumpDir = new Path(dumpScript.getParent(), DUMP_FOLDER); + if (!fs.exists(dumpDir)) { + fs.mkdirs(dumpDir); + } + BufferedWriter scriptWriter = + closer.register(new BufferedWriter(new OutputStreamWriter(fs.create(dumpScript), Charset + .forName(ConfigurationKeys.DEFAULT_CHARSET_ENCODING)))); + + scriptWriter.write("#!/bin/sh\n"); + scriptWriter.write("if [ -n \"$HADOOP_PREFIX\" ]; then\n"); + scriptWriter.write(" ${HADOOP_PREFIX}/bin/hadoop dfs -put " + heapFileName + " " + dumpDir + + "/${PWD//\\//_}.hprof\n"); + scriptWriter.write("else\n"); + scriptWriter.write(" ${HADOOP_HOME}/bin/hadoop dfs -put " + heapFileName + " " + dumpDir + + "/${PWD//\\//_}.hprof\n"); + scriptWriter.write("fi\n"); + + } catch (IOException ioe) { + LOG.error("Heap dump script is not generated successfully."); + if (fs.exists(dumpScript)) { + fs.delete(dumpScript, true); + } + throw ioe; + } catch (Throwable t) { + closer.rethrow(t); + } finally { + closer.close(); + } + Runtime.getRuntime().exec(chmod + " " + dumpScript); + } +} diff --git a/gobblin-utility/src/test/java/gobblin/util/HeapDumpForTaskUtilsTest.java b/gobblin-utility/src/test/java/gobblin/util/HeapDumpForTaskUtilsTest.java new file mode 100644 index 00000000000..e44d2ad7179 --- /dev/null +++ b/gobblin-utility/src/test/java/gobblin/util/HeapDumpForTaskUtilsTest.java @@ -0,0 +1,73 @@ +/* (c) 2014 LinkedIn Corp. All rights reserved. + * + * 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. + */ + +package gobblin.util; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStreamReader; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.FileSystem; +import org.apache.hadoop.fs.Path; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import com.google.common.io.Closer; + + +@Test(groups = { "gobblin.util" }) +public class HeapDumpForTaskUtilsTest { + + private FileSystem fs; + + private static final String TEST_DIR = "dumpScript"; + private static final String SCRIPT_NAME = "dump.sh"; + + @BeforeClass + public void setUp() throws IOException { + this.fs = FileSystem.getLocal(new Configuration()); + this.fs.mkdirs(new Path(TEST_DIR)); + } + + @Test + public void testGenerateDumpScript() throws IOException { + Path dumpScript = new Path(TEST_DIR, SCRIPT_NAME); + HeapDumpForTaskUtils.generateDumpScript(dumpScript, this.fs, "test.hprof", "chmod 777 "); + Assert.assertEquals(true, this.fs.exists(dumpScript)); + Assert.assertEquals(true, this.fs.exists(new Path(dumpScript.getParent(), "dumps"))); + Closer closer = Closer.create(); + try { + BufferedReader scriptReader = + closer.register(new BufferedReader(new InputStreamReader(this.fs.open(dumpScript)))); + Assert.assertEquals("#!/bin/sh", scriptReader.readLine()); + Assert.assertEquals("if [ -n \"$HADOOP_PREFIX\" ]; then", scriptReader.readLine()); + Assert.assertEquals(" ${HADOOP_PREFIX}/bin/hadoop dfs -put test.hprof dumpScript/dumps/${PWD//\\//_}.hprof", + scriptReader.readLine()); + Assert.assertEquals("else", scriptReader.readLine()); + Assert.assertEquals(" ${HADOOP_HOME}/bin/hadoop dfs -put test.hprof dumpScript/dumps/${PWD//\\//_}.hprof", + scriptReader.readLine()); + Assert.assertEquals("fi", scriptReader.readLine()); + } catch (Throwable t) { + closer.rethrow(t); + } finally { + closer.close(); + } + } + + @AfterClass + public void tearDown() throws IOException { + fs.delete(new Path(TEST_DIR), true); + fs.close(); + } +}