001    /**
002     * Licensed to the Apache Software Foundation (ASF) under one
003     * or more contributor license agreements.  See the NOTICE file
004     * distributed with this work for additional information
005     * regarding copyright ownership.  The ASF licenses this file
006     * to you under the Apache License, Version 2.0 (the
007     * "License"); you may not use this file except in compliance
008     * with the License.  You may obtain a copy of the License at
009     *
010     *     http://www.apache.org/licenses/LICENSE-2.0
011     *
012     * Unless required by applicable law or agreed to in writing, software
013     * distributed under the License is distributed on an "AS IS" BASIS,
014     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015     * See the License for the specific language governing permissions and
016     * limitations under the License.
017     */
018    package org.apache.hadoop.util;
019    
020    import java.io.BufferedReader;
021    import java.io.File;
022    import java.io.IOException;
023    import java.io.InputStreamReader;
024    import java.util.Map;
025    import java.util.Timer;
026    import java.util.TimerTask;
027    import java.util.concurrent.atomic.AtomicBoolean;
028    
029    import org.apache.commons.logging.Log;
030    import org.apache.commons.logging.LogFactory;
031    import org.apache.hadoop.classification.InterfaceAudience;
032    import org.apache.hadoop.classification.InterfaceStability;
033    import org.apache.hadoop.conf.Configuration;
034    
035    /** 
036     * A base class for running a Unix command.
037     * 
038     * <code>Shell</code> can be used to run unix commands like <code>du</code> or
039     * <code>df</code>. It also offers facilities to gate commands by 
040     * time-intervals.
041     */
042    @InterfaceAudience.LimitedPrivate({"HDFS", "MapReduce"})
043    @InterfaceStability.Unstable
044    abstract public class Shell {
045      
046      public static final Log LOG = LogFactory.getLog(Shell.class);
047      
048      /** a Unix command to get the current user's name */
049      public final static String USER_NAME_COMMAND = "whoami";
050      /** a Unix command to get the current user's groups list */
051      public static String[] getGroupsCommand() {
052        return new String[]{"bash", "-c", "groups"};
053      }
054      /** a Unix command to get a given user's groups list */
055      public static String[] getGroupsForUserCommand(final String user) {
056        //'groups username' command return is non-consistent across different unixes
057        return new String [] {"bash", "-c", "id -Gn " + user};
058      }
059      /** a Unix command to get a given netgroup's user list */
060      public static String[] getUsersForNetgroupCommand(final String netgroup) {
061        //'groups username' command return is non-consistent across different unixes
062        return new String [] {"bash", "-c", "getent netgroup " + netgroup};
063      }
064      /** a Unix command to set permission */
065      public static final String SET_PERMISSION_COMMAND = "chmod";
066      /** a Unix command to set owner */
067      public static final String SET_OWNER_COMMAND = "chown";
068      public static final String SET_GROUP_COMMAND = "chgrp";
069      /** a Unix command to create a link */
070      public static final String LINK_COMMAND = "ln";
071      /** a Unix command to get a link target */
072      public static final String READ_LINK_COMMAND = "readlink";
073      /** Return a Unix command to get permission information. */
074      public static String[] getGET_PERMISSION_COMMAND() {
075        //force /bin/ls, except on windows.
076        return new String[] {(WINDOWS ? "ls" : "/bin/ls"), "-ld"};
077      }
078    
079      /**Time after which the executing script would be timedout*/
080      protected long timeOutInterval = 0L;
081      /** If or not script timed out*/
082      private AtomicBoolean timedOut;
083    
084      /** Set to true on Windows platforms */
085      public static final boolean WINDOWS /* borrowed from Path.WINDOWS */
086                    = System.getProperty("os.name").startsWith("Windows");
087      
088      private long    interval;   // refresh interval in msec
089      private long    lastTime;   // last time the command was performed
090      private Map<String, String> environment; // env for the command execution
091      private File dir;
092      private Process process; // sub process used to execute the command
093      private int exitCode;
094    
095      /**If or not script finished executing*/
096      private volatile AtomicBoolean completed;
097      
098      public Shell() {
099        this(0L);
100      }
101      
102      /**
103       * @param interval the minimum duration to wait before re-executing the 
104       *        command.
105       */
106      public Shell( long interval ) {
107        this.interval = interval;
108        this.lastTime = (interval<0) ? 0 : -interval;
109      }
110      
111      /** set the environment for the command 
112       * @param env Mapping of environment variables
113       */
114      protected void setEnvironment(Map<String, String> env) {
115        this.environment = env;
116      }
117    
118      /** set the working directory 
119       * @param dir The directory where the command would be executed
120       */
121      protected void setWorkingDirectory(File dir) {
122        this.dir = dir;
123      }
124    
125      /** check to see if a command needs to be executed and execute if needed */
126      protected void run() throws IOException {
127        if (lastTime + interval > System.currentTimeMillis())
128          return;
129        exitCode = 0; // reset for next run
130        runCommand();
131      }
132    
133      /** Run a command */
134      private void runCommand() throws IOException { 
135        ProcessBuilder builder = new ProcessBuilder(getExecString());
136        Timer timeOutTimer = null;
137        ShellTimeoutTimerTask timeoutTimerTask = null;
138        timedOut = new AtomicBoolean(false);
139        completed = new AtomicBoolean(false);
140        
141        if (environment != null) {
142          builder.environment().putAll(this.environment);
143        }
144        if (dir != null) {
145          builder.directory(this.dir);
146        }
147        
148        process = builder.start();
149        if (timeOutInterval > 0) {
150          timeOutTimer = new Timer("Shell command timeout");
151          timeoutTimerTask = new ShellTimeoutTimerTask(
152              this);
153          //One time scheduling.
154          timeOutTimer.schedule(timeoutTimerTask, timeOutInterval);
155        }
156        final BufferedReader errReader = 
157                new BufferedReader(new InputStreamReader(process
158                                                         .getErrorStream()));
159        BufferedReader inReader = 
160                new BufferedReader(new InputStreamReader(process
161                                                         .getInputStream()));
162        final StringBuffer errMsg = new StringBuffer();
163        
164        // read error and input streams as this would free up the buffers
165        // free the error stream buffer
166        Thread errThread = new Thread() {
167          @Override
168          public void run() {
169            try {
170              String line = errReader.readLine();
171              while((line != null) && !isInterrupted()) {
172                errMsg.append(line);
173                errMsg.append(System.getProperty("line.separator"));
174                line = errReader.readLine();
175              }
176            } catch(IOException ioe) {
177              LOG.warn("Error reading the error stream", ioe);
178            }
179          }
180        };
181        try {
182          errThread.start();
183        } catch (IllegalStateException ise) { }
184        try {
185          parseExecResult(inReader); // parse the output
186          // clear the input stream buffer
187          String line = inReader.readLine();
188          while(line != null) { 
189            line = inReader.readLine();
190          }
191          // wait for the process to finish and check the exit code
192          exitCode  = process.waitFor();
193          try {
194            // make sure that the error thread exits
195            errThread.join();
196          } catch (InterruptedException ie) {
197            LOG.warn("Interrupted while reading the error stream", ie);
198          }
199          completed.set(true);
200          //the timeout thread handling
201          //taken care in finally block
202          if (exitCode != 0) {
203            throw new ExitCodeException(exitCode, errMsg.toString());
204          }
205        } catch (InterruptedException ie) {
206          throw new IOException(ie.toString());
207        } finally {
208          if (timeOutTimer != null) {
209            timeOutTimer.cancel();
210          }
211          // close the input stream
212          try {
213            inReader.close();
214          } catch (IOException ioe) {
215            LOG.warn("Error while closing the input stream", ioe);
216          }
217          if (!completed.get()) {
218            errThread.interrupt();
219          }
220          try {
221            errReader.close();
222          } catch (IOException ioe) {
223            LOG.warn("Error while closing the error stream", ioe);
224          }
225          process.destroy();
226          lastTime = System.currentTimeMillis();
227        }
228      }
229    
230      /** return an array containing the command name & its parameters */ 
231      protected abstract String[] getExecString();
232      
233      /** Parse the execution result */
234      protected abstract void parseExecResult(BufferedReader lines)
235      throws IOException;
236    
237      /** get the current sub-process executing the given command 
238       * @return process executing the command
239       */
240      public Process getProcess() {
241        return process;
242      }
243    
244      /** get the exit code 
245       * @return the exit code of the process
246       */
247      public int getExitCode() {
248        return exitCode;
249      }
250    
251      /**
252       * This is an IOException with exit code added.
253       */
254      public static class ExitCodeException extends IOException {
255        int exitCode;
256        
257        public ExitCodeException(int exitCode, String message) {
258          super(message);
259          this.exitCode = exitCode;
260        }
261        
262        public int getExitCode() {
263          return exitCode;
264        }
265      }
266      
267      /**
268       * A simple shell command executor.
269       * 
270       * <code>ShellCommandExecutor</code>should be used in cases where the output 
271       * of the command needs no explicit parsing and where the command, working 
272       * directory and the environment remains unchanged. The output of the command 
273       * is stored as-is and is expected to be small.
274       */
275      public static class ShellCommandExecutor extends Shell {
276        
277        private String[] command;
278        private StringBuffer output;
279        
280        
281        public ShellCommandExecutor(String[] execString) {
282          this(execString, null);
283        }
284        
285        public ShellCommandExecutor(String[] execString, File dir) {
286          this(execString, dir, null);
287        }
288       
289        public ShellCommandExecutor(String[] execString, File dir, 
290                                     Map<String, String> env) {
291          this(execString, dir, env , 0L);
292        }
293    
294        /**
295         * Create a new instance of the ShellCommandExecutor to execute a command.
296         * 
297         * @param execString The command to execute with arguments
298         * @param dir If not-null, specifies the directory which should be set
299         *            as the current working directory for the command.
300         *            If null, the current working directory is not modified.
301         * @param env If not-null, environment of the command will include the
302         *            key-value pairs specified in the map. If null, the current
303         *            environment is not modified.
304         * @param timeout Specifies the time in milliseconds, after which the
305         *                command will be killed and the status marked as timedout.
306         *                If 0, the command will not be timed out. 
307         */
308        public ShellCommandExecutor(String[] execString, File dir, 
309            Map<String, String> env, long timeout) {
310          command = execString.clone();
311          if (dir != null) {
312            setWorkingDirectory(dir);
313          }
314          if (env != null) {
315            setEnvironment(env);
316          }
317          timeOutInterval = timeout;
318        }
319            
320    
321        /** Execute the shell command. */
322        public void execute() throws IOException {
323          this.run();    
324        }
325    
326        public String[] getExecString() {
327          return command;
328        }
329    
330        protected void parseExecResult(BufferedReader lines) throws IOException {
331          output = new StringBuffer();
332          char[] buf = new char[512];
333          int nRead;
334          while ( (nRead = lines.read(buf, 0, buf.length)) > 0 ) {
335            output.append(buf, 0, nRead);
336          }
337        }
338        
339        /** Get the output of the shell command.*/
340        public String getOutput() {
341          return (output == null) ? "" : output.toString();
342        }
343    
344        /**
345         * Returns the commands of this instance.
346         * Arguments with spaces in are presented with quotes round; other
347         * arguments are presented raw
348         *
349         * @return a string representation of the object.
350         */
351        public String toString() {
352          StringBuilder builder = new StringBuilder();
353          String[] args = getExecString();
354          for (String s : args) {
355            if (s.indexOf(' ') >= 0) {
356              builder.append('"').append(s).append('"');
357            } else {
358              builder.append(s);
359            }
360            builder.append(' ');
361          }
362          return builder.toString();
363        }
364      }
365      
366      /**
367       * To check if the passed script to shell command executor timed out or
368       * not.
369       * 
370       * @return if the script timed out.
371       */
372      public boolean isTimedOut() {
373        return timedOut.get();
374      }
375      
376      /**
377       * Set if the command has timed out.
378       * 
379       */
380      private void setTimedOut() {
381        this.timedOut.set(true);
382      }
383      
384      /** 
385       * Static method to execute a shell command. 
386       * Covers most of the simple cases without requiring the user to implement  
387       * the <code>Shell</code> interface.
388       * @param cmd shell command to execute.
389       * @return the output of the executed command.
390       */
391      public static String execCommand(String ... cmd) throws IOException {
392        return execCommand(null, cmd, 0L);
393      }
394      
395      /** 
396       * Static method to execute a shell command. 
397       * Covers most of the simple cases without requiring the user to implement  
398       * the <code>Shell</code> interface.
399       * @param env the map of environment key=value
400       * @param cmd shell command to execute.
401       * @param timeout time in milliseconds after which script should be marked timeout
402       * @return the output of the executed command.o
403       */
404      
405      public static String execCommand(Map<String, String> env, String[] cmd,
406          long timeout) throws IOException {
407        ShellCommandExecutor exec = new ShellCommandExecutor(cmd, null, env, 
408                                                              timeout);
409        exec.execute();
410        return exec.getOutput();
411      }
412    
413      /** 
414       * Static method to execute a shell command. 
415       * Covers most of the simple cases without requiring the user to implement  
416       * the <code>Shell</code> interface.
417       * @param env the map of environment key=value
418       * @param cmd shell command to execute.
419       * @return the output of the executed command.
420       */
421      public static String execCommand(Map<String,String> env, String ... cmd) 
422      throws IOException {
423        return execCommand(env, cmd, 0L);
424      }
425      
426      /**
427       * Timer which is used to timeout scripts spawned off by shell.
428       */
429      private static class ShellTimeoutTimerTask extends TimerTask {
430    
431        private Shell shell;
432    
433        public ShellTimeoutTimerTask(Shell shell) {
434          this.shell = shell;
435        }
436    
437        @Override
438        public void run() {
439          Process p = shell.getProcess();
440          try {
441            p.exitValue();
442          } catch (Exception e) {
443            //Process has not terminated.
444            //So check if it has completed 
445            //if not just destroy it.
446            if (p != null && !shell.completed.get()) {
447              shell.setTimedOut();
448              p.destroy();
449            }
450          }
451        }
452      }
453    }