pump.in

#!/bin/sh
#
# Copyright 2007 Google Inc. All Rights Reserved.
#
# Authors: klarlund@google.com, fergus@google.com
#
# 'pump': a script for using distcc-pump with build commands.

# This file is processed by configure, which substitutes in the right
# value for @PYTHON@ and @PYTHON_RELATIVE_LIB@.

PYTHON=@PYTHON@

# This is where include_server.py lives.  It's set to the empty string
# before pump is installed into its final location.  At install-time,
# the 'make install' command will rewrite this line of pump to point
# to the installed include_server.py.
# NOTE: DO NOT CHANGE THE LINE BELOW WITHOUT CHANGING THE SED IN
#       Makefile.in:install-include-server.
include_server=''


CheckUsage() {
  if [ "$1" = "" -o "$1" = '-h' -o "$1" = '--help' ]; then
    cat <<'EOF'
Usage:
    pump COMMAND [ARG...]
or
    pump --startup
    pump --shutdown

Description:
  Pump, also known as distcc-pump, accelerates remote compilation with
  distcc by also distributing preprocessing to the servers.

  The simplest usage is the form "pump COMMAND [ARG...]".
  This will start an include server for distcc-pump; set some environment
  variables; change PATH to use the distcc-pump "distcc" client; execute
  COMMAND with the specified ARG(s); and then shutdown the include server.
  The COMMAND is typically a parallel build command, such as
  "make -j80", that will do many concurrent invocations of distcc.

  An alternative way of invoking pump is to explicitly invoke "pump --startup"
  to start the include server and "pump --shutdown" to stop the include server.
  The "pump --startup" command will start up the include server, and will print
  out some environment variable settings.  These environment variables are used
  to communicate between the distcc-pump "distcc" client and the include
  server, and to communicate between "pump --startup" and "pump --shutdown".
  The caller of "pump --startup" is responsible for setting those environment
  variables before invoking "distcc" or "pump --shutdown".  For example:

      eval `pump --startup`
      make -j80
      pump --shutdown

  Note that distcc-pump assumes that sources files will not be modified during
  the lifetime of the include server, so modifying source files during a build
  may cause inconsistent results.

Environment variables (all optional):

  DISTCC_LOCATION        Either the distcc 'bin' install directory (when the
                         pump script in the installation executes) or the build
                         directory (when the pump script is run during 
                         development).

                         This location is normally inferred from the dirname
                         of argv[0].  It is resolved with respect to links.

                         If this location is the bin directory inside an
                         installation, then Python executables and distcc
                         itself will be retrieved from the parent directory of
                         the location.

                         If this location is the build directory, then Python 
                         sources are retrived from the source directory.

  DISTCC_POTENTIAL_HOSTS The distcc servers that will be queried by lsdistcc
                         in order to produce a value for DISTCC_HOSTS.
                         This value may be unset or null. In such
                         cases, use DISTCC_HOSTS.

  DISTCC_HOSTS           This variable is passed through to distcc but only if 
                         DISTCC_POTENTIAL_HOSTS is not set.

  LSDISTCC_ARGS          Extra arguments to lsdistcc.

  INCLUDE_SERVER_ARGS    Extra arguments to the include server.

  PYTHONOPTIMIZE         If set to "", then Python optimization is disabled.

Example: 
  pump make
EOF
    exit 0
  fi
}

# Get the directory of a command. The argument is $0 as received from
# argv[0].  If $0 is of the form /path/symlink and symlink points to
# another directory than /path, then `basename $0` does not return
# the directory where the executable lives. We fix this here.
GetScriptDir() {
  if [ -h "$0" ]; then
    dir=`dirname "$0"`
    link=`readlink "$0"`
    linkdir=`dirname "$link"`
    case $linkdir in
         /*) echo "$linkdir" ;;
         *)  echo "$dir/$linkdir" ;;
    esac
  else
    dirname "$0"
  fi
}

# Variables inherited from the environment of the caller.
if [ -z "$DISTCC_LOCATION" ]; then
   # Set the default for location of 'pump' script'.
   DISTCC_LOCATION=`GetScriptDir "$0"`
fi


# Variables used by this process.
program_name="$0"
socket_dir=''           # Temporary directory created by this process.
tmp_pid_file=''         # Temporary file created by this process.
available_hosts_file='' # Temporary file for newline separated list of hosts.
socket=''               # Temporary socket file, inside $socket_dir.
include_server_stdout=''  # Temporary file, inside $socket_dir.
include_server_stderr=''  # Temporary file, inside $socket_dir.
include_server_pid=''
include_server_relative=''  # Set by Initialize(). This path is relative to
                            # the 'bin' directory of a distcc-pump installation.
distcc_location='' # Set by Initialize().  The directory of the distcc binary.
DISTCC='distcc'    # the name of the distcc binary
verbose=1          # Print progress messages to stdout? (1 means yes.)
redirect_io=0      # Redirect include server's I/O streams? (1 means yes.)

# Make temp file using distinguishing prefix $1. Use optional argument
# $2="-d" to make a directory. The name of the created temp file or
# directory is written to stdout.
MakeTmpFile() {
  if mktemp $2 /tmp/$1.XXXXXX; then
    :   # mktemp prints the output we want; no need to do more
  else
    echo "$program_name: Could not make temp \"$1\"" 1>&2
    exit 1
  fi
}

Initialize() {
  # We know the value of $PYTHON when this function is executed. The value is
  # set when the 'pump' script is made from 'pump.in'. Use this value in the
  # formation of the relative path that takes us from the 'bin' directory of an
  # installation to the 'include_server' directory.
  include_server_relative="../$PYTHON_RELATIVE_LIB/include_server"

  # We use a little heuristic to determine whether this pump script is part of
  # an installation. Specifically, we check whether we're a bin directory, and
  # more importantly, we look for the include_server.py script. If we're running
  # out of the distcc_pump source directory, both conditions below will fail.
  if [ "`basename $DISTCC_LOCATION`" = 'bin' ] \
     && [ -f "$DISTCC_LOCATION/$include_server_relative/include_server.py" ]
  then
    # Running from installed directory.
    is_in_installation=true
  else
    # Running from configured directory, usually the source directory.
    is_in_installation=""
  fi

  distcc_location="$DISTCC_LOCATION"

  # Check that we're getting exactly the distcc client we want: the
  # one that is part of this build.
  if [ ! -x "$distcc_location/$DISTCC" ]; then
    echo "$0: error: can't find distcc (looked in $distcc_location)" 1>&2
    exit 1
  fi
}

PrintIncludeServerStatusMessage() {
  include_server_status=$1
  if [ "$verbose" = 1 ]; then
    if [ "$include_server_status" = 0 ]; then
      echo '__________Started distcc-pump include server'
    else
      echo '__________Could not start distcc-pump include server' 1>&2
    fi
  fi
}

Announce() {
  echo "__________Using distcc-pump from $DISTCC_LOCATION"
}

# Starts up the include server.  Sets $socket, $socket_dir, and
# $include_server_pid.  If successful (with exit status 0), sets exported
# variable $INCLUDE_SERVER_PORT to the socket file ($socket), to tell the distcc
# clients where to find the include server.
StartIncludeServer() {
  # If include_server already exists, that means we're a
  # installed pump (in /usr/local/bin somewhere or something), and
  # include_server points to the installed include_server.py.
  if [ -n "$include_server" ]; then
    pythonpath=`dirname "$include_server"`
  else
    # We assume this script is run from the build directory.  We pick up .py
    # files from the include_server directory in the source tree, and we pick up
    # the .so file from the include_server/build/libXXX/include_server
    # directory. 
    #
    # Calculate the location of the source directory.
    distcc_srcdir=`cd $DISTCC_LOCATION; cd @srcdir@; pwd`
    # The source tree location must be passed.
    include_server="$distcc_srcdir/include_server/include_server.py"
    # Now locate the single directory containing the .so file from the build
    # directory. Possibly there may be more than one such file; first identify
    # them all.
    so_dir=`"$distcc_srcdir/find_c_extension.sh" "$DISTCC_LOCATION"`
    pythonpath="$so_dir"
  fi

  # Create a temporary directory $socket_dir.
  socket_dir=`MakeTmpFile "distcc-pump" -d`

  # The socket file on which the include server accepts connections.
  socket="$socket_dir/socket"

  # Files for the include server's stdout/stderr.
  # When a build tool invokes 'pump --startup', stdout/stderr may be
  # pipes, in which case the invocation may hang unless the include
  # server process closes them.  So to avoid this, we need to redirect
  # the include server's output to temporary files.
  # We print these files during shutdown; better late than never!
  include_server_stdout="$socket_dir/stdout"
  include_server_stderr="$socket_dir/stderr"

  # File for the include server process id.
  tmp_pid_file=`MakeTmpFile "distcc-pump-pid"`

  # Start include server in optimized mode (no assertions) and with
  # debug level 1 for tracing warnings.
  # The include server will fork off
  # a background process to handle the requests;
  # the main process will exit only when the background
  # process is ready to accept connections.

  (
    # Optionally redirect the I/O streams for the include server.
    case $redirect_io in
      1) exec < /dev/null            \
           > $include_server_stdout  \
           2> $include_server_stderr
         ;;
      *)
          rm -f $include_server_stdout $include_server_stderr
          ;;
    esac

    if [ -z "$PYTHONOPTIMIZE" ]; then PYTHONOPTIMIZE=1; fi
    PYTHONOPTIMIZE=$PYTHONOPTIMIZE \
    PYTHONPATH="$pythonpath::$PYTHONPATH" \
    $PYTHON                          \
        "$include_server"            \
        --port "$socket"             \
        --pid_file "$tmp_pid_file"   \
        --email                      \
        --email_bound=5              \
        -d1                          \
        $INCLUDE_SERVER_ARGS
  )

  # solaris sh's built-in test does not support -S, so we need to use
  # the binary.
  if [ -x /bin/test ]; then
    TEST=/bin/test
  elif [ -x /usr/bin/test ]; then
    TEST=/usr/bin/test
  else
    TEST=test
  fi
  if $TEST ! -S "$socket"; then
    echo "__________Expected a socket at '$socket'" 1>&2
    PrintIncludeServerStatusMessage 1
    return 1
  fi

  include_server_pid=`cat $tmp_pid_file`
  rm "$tmp_pid_file"

  # Make sure $include_server_pid looks like a pid
  if echo "$include_server_pid" | grep '^[0-9][0-9]*$' >/dev/null; then
    # We got a number.  Tell the distcc clients where to find it.
    INCLUDE_SERVER_PORT="$socket"
    export INCLUDE_SERVER_PORT
  else
    # We got something that's not a number.  This indicates the
    # socket is not working.
    include_server_pid=''
    PrintIncludeServerStatusMessage 1
    return 1
  fi
}

ShutDown() {
  # Always -- at exit -- shut down include_server and remove $socket_dir
  if [ -n "$include_server_pid" ] && \
    ps -p "$include_server_pid" > /dev/null; then
    echo '__________Shutting down distcc-pump include server'
    kill $include_server_pid
  fi

  if [ -f "$include_server_stdout" ]; then
    cat "$include_server_stdout"
  fi
  if [ -f "$include_server_stderr" ]; then
    cat "$include_server_stderr" >&2
  fi

  if [ -n "$socket_dir" ];  then
    rm -rf "$socket_dir"
  fi
  if [ -n "$tmp_pid_file" ]; then
    rm -f "$tmp_pid_file"
  fi
  if [ -n "$available_hosts_file" ]; then
    rm -f "$available_hosts_file"
  fi
}

# Invokes lsdistcc to find the available servers. This list is
# \n-separated and written to the filepath provided as $1.
AvailableHosts() {
  available_hosts="$1"
  lsdistcc="$DISTCC_LOCATION/lsdistcc"
  if [ ! -x "$lsdistcc" ]; then
     echo "$0: error: can't find lsdistcc (looked in $DISTCC_LOCATION)" 1>&2
     exit 1
  fi
  # Call lsdistcc and let it wait no more than 150ms unless overridden in
  # LSDISTCC_ARGS.
  "$lsdistcc" -c150 $LSDISTCC_ARGS $DISTCC_POTENTIAL_HOSTS > "$available_hosts"
}

# Exports DISTCC_HOSTS as a function of the single argument and the global
# variable $available_hosts_file.
ExportDISTCC_HOSTS() {
  include_server_status="$1"
  opts=''
  if [ "$include_server_status" = 0 ]; then
    # The include server is up.
    # Make 'cpp' mode (with lzo) the default.
    opts=",cpp,lzo"
  else
    # The include server is not ready.
    # Make 'lzo' mode the default.
    if [ "$?" = 1 ]; then
      echo "__________Warning: distcc-pump include server failed;" \
        "running unpumped" 1>&2
    fi
    opts=",lzo"
  fi
  # When calculating final value of DISTCC_HOSTS, we get rid of
  # newlines (thanks to "`sed ...`").
  export DISTCC_HOSTS="--randomize `sed s/'$'/$opts/ $available_hosts_file`"
  if [ "$verbose" = 1 ]; then
    echo "__________Found" \
         "`wc -l < $available_hosts_file` available distcc servers"
  fi
}

StartIncludeServerAndDetermineHosts() {
  include_server_status=1 # ie, not running (yet)
  if [ -n "$DISTCC_POTENTIAL_HOSTS" ]; then
    # Probe the distcc servers. It may take up to a second. But
    # starting the include server takes about 150ms. Make these
    # activities happen in parallel. 

    # First, the lsdistcc command goes in the background.
    available_hosts_file=`MakeTmpFile "distcc-pump-hosts"`
    AvailableHosts "$available_hosts_file" &

    # Second, the include server goes in the foreground so variables can be set.
    StartIncludeServer && include_server_status=0

    # Await for AvailableHosts to finish.
    wait 
    ExportDISTCC_HOSTS "$include_server_status"
    rm -f "$available_hosts_file"
  else
    if [ -z "$DISTCC_HOSTS" ] &&
       [ -z "$DISTCC_DIR" -o ! -f "$DISTCC_DIR/hosts" ] &&
       [ -z "$HOME" -o ! -f "$HOME/.distcc/hosts" ] &&
       [ ! -f "/etc/distcc/hosts" ]
    then
      echo "$program_name:" \
              "expected environment variables \"DISTCC_HOSTS\" or" \
              "\"DISTCC_POTENTIAL_HOSTS\" to be set, or to find a" \
              "distcc hosts file in \"\$DISTCC_DIR/hosts\"," \
              "\"\$HOME/.distcc/hosts\", or \"/etc/distcc/hosts\"." \
              1>&2
      exit 1
    fi
    StartIncludeServer && include_server_status=0
  fi
  return "$include_server_status"
}

# Prints out environment variable settings, for the --startup option.
DumpEnvironmentVariables() {
  # Variables used by "pump --shutdown"
  echo export INCLUDE_SERVER_PID=\'$include_server_pid\'
  echo export INCLUDE_SERVER_DIR=\'$socket_dir\'
  # Variables used by the distcc client
  echo export INCLUDE_SERVER_PORT=\'$INCLUDE_SERVER_PORT\'
  if [ -n "$DISTCC_HOSTS" ]; then
    echo export DISTCC_HOSTS=\'$DISTCC_HOSTS\'
  fi
  echo export PATH=\'$distcc_location:$PATH\'
}

Main() {
  CheckUsage "$@"

  Initialize

  case "$*" in
    --startup)
      # Don't put ordinary progress messages on stdout,
      # because they interfere with the environment variable
      # settings that we print out.
      verbose=0
      # Redirect the include server's stdin/stdout/stderr
      redirect_io=1
      trap 'ShutDown' EXIT  # In case we get interrupted.
      StartIncludeServerAndDetermineHosts
      include_server_status=$?
      trap '' EXIT
      if [ "$include_server_status" = 0 ]; then
        DumpEnvironmentVariables
        exit 0
      else
        exit 1
      fi
      ;;
    --shutdown)
      include_server_pid="$INCLUDE_SERVER_PID"
      socket_dir="$INCLUDE_SERVER_DIR"
      include_server_stdout="$socket_dir/stdout"
      include_server_stderr="$socket_dir/stderr"
      ShutDown
      exit 0
      ;;
    *)
      trap 'ShutDown' EXIT
      Announce
      StartIncludeServerAndDetermineHosts || exit 1
      # Now execute the command that is the argument of 'pump'.
      PATH="$distcc_location:$PATH" \
        "$@"
      # When we exit, the ShutDown function will be called.
      ;;
  esac
}

Main "$@"