org.apache.hadoop.mapred.TestTaskEnvironment.java Source code

Introduction

Here is the source code for org.apache.hadoop.mapred.TestTaskEnvironment.java

Source

/**
 * 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.hadoop.mapred;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.util.Vector;
import java.util.HashMap;
import java.util.Map;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.filecache.TrackerDistributedCacheManager;
import org.apache.hadoop.fs.FileUtil;
import org.apache.hadoop.mapred.JvmManager.JvmManagerForType;
import org.apache.hadoop.mapred.JvmManager.JvmManagerForType.JvmRunner;
import org.apache.hadoop.mapred.JvmManager.JvmEnv;
import org.apache.hadoop.mapred.TaskTracker.RunningJob;
import org.apache.hadoop.mapred.TaskTracker.TaskInProgress;
import org.apache.hadoop.mapreduce.TaskType;
import org.apache.hadoop.mapreduce.server.tasktracker.userlogs.UserLogManager;
import org.junit.After;
import static org.junit.Assert.*;
import org.junit.Before;
import org.junit.Test;

public class TestTaskEnvironment {
    private static File TEST_DIR = new File(System.getProperty("test.build.data", "/tmp"),
            TestJvmManager.class.getSimpleName());
    private static int MAP_SLOTS = 1;
    private static int REDUCE_SLOTS = 1;
    private TaskTracker tt;
    private JvmManager jvmManager;
    private JobConf ttConf;

    @Before
    public void setUp() {
        TEST_DIR.mkdirs();
    }

    @After
    public void tearDown() throws IOException {
        FileUtil.fullyDelete(TEST_DIR);
    }

    public TestTaskEnvironment() throws Exception {
        tt = new TaskTracker();
        ttConf = new JobConf();
        ttConf.setLong("mapred.tasktracker.tasks.sleeptime-before-sigkill", 2000);
        tt.setConf(ttConf);
        tt.setMaxMapSlots(MAP_SLOTS);
        tt.setMaxReduceSlots(REDUCE_SLOTS);
        tt.setTaskController(new DefaultTaskController());
        jvmManager = new JvmManager(tt);
        tt.setJvmManagerInstance(jvmManager);
        tt.setUserLogManager(new UserLogManager(ttConf));
    }

    // write a shell script to execute the command.
    private File writeScript(String fileName, String cmd, File pidFile) throws IOException {
        File script = new File(TEST_DIR, fileName);
        FileOutputStream out = new FileOutputStream(script);
        // write pid into a file
        out.write(("echo $$ >" + pidFile.toString() + ";").getBytes());
        // ignore SIGTERM
        out.write(("trap '' 15\n").getBytes());
        // write the actual command it self.
        out.write(cmd.getBytes());
        out.close();
        script.setExecutable(true);
        return script;
    }

    @Test
    public void testTaskEnv() throws Throwable {
        ttConf.set("mapreduce.admin.user.shell", "/bin/testshell");
        ttConf.set("mapreduce.admin.user.env", "key1=value1,key2=value2");
        ttConf.set("mapred.child.env", "ROOT=$HOME");
        final Map<String, String> env = new HashMap<String, String>();
        String user = "test";
        JobConf taskConf = new JobConf(ttConf);
        TaskAttemptID attemptID = new TaskAttemptID("test", 0, true, 0, 0);
        Task task = new MapTask(null, attemptID, 0, null, MAP_SLOTS);
        task.setConf(taskConf);
        TaskInProgress tip = tt.new TaskInProgress(task, taskConf);
        RunningJob rjob = new RunningJob(attemptID.getJobID());
        TaskController taskController = new DefaultTaskController();
        taskController.setConf(ttConf);
        rjob.distCacheMgr = new TrackerDistributedCacheManager(ttConf, taskController)
                .newTaskDistributedCacheManager(attemptID.getJobID(), taskConf);

        final TaskRunner taskRunner = task.createRunner(tt, tip, rjob);
        String errorInfo = "Child error";
        String mapredChildEnv = taskRunner.getChildEnv(taskConf);
        taskRunner.updateUserLoginEnv(errorInfo, user, taskConf, env);
        taskRunner.setEnvFromInputString(errorInfo, env, mapredChildEnv);

        final Vector<String> vargs = new Vector<String>(1);
        File pidFile = new File(TEST_DIR, "pid");
        vargs.add(writeScript("ENV", "/bin/env ", pidFile).getAbsolutePath());
        final File workDir = new File(TEST_DIR, "work");
        workDir.mkdir();
        final File stdout = new File(TEST_DIR, "stdout");
        final File stderr = new File(TEST_DIR, "stderr");
        Map<String, String> jvmenvmap = env;
        String javaOpts = taskRunner.getChildJavaOpts(ttConf, JobConf.MAPRED_MAP_TASK_JAVA_OPTS);

        assertTrue(jvmenvmap.containsKey("SHELL"));
        assertTrue(jvmenvmap.containsValue("/bin/testshell"));
        assertTrue(jvmenvmap.containsKey("key2"));
        assertTrue(jvmenvmap.containsValue("value2"));
        assertTrue(javaOpts, javaOpts.contains("Xmx"));
        assertTrue(javaOpts, javaOpts.contains("IPv4"));

        String root = jvmenvmap.get("ROOT");
        assertTrue(root, root.contentEquals(System.getenv("HOME")));
    }
}