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