sdks/python/scripts/run_expansion_services.sh

#!/bin/bash
#
#    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.
#

read -r -d '' USAGE <<END
Usage: run_expansion_services.sh (start|stop) [options]
Options:
  --group_id [unique id for stop services later]
  --java_expansion_service_jar [path to expansion service jar]
  --python_virtualenv_dir [path to virtualenv root dir]
  --python_expansion_service_module [name of expansion service module]
  --java_port [number]
  --python_port [number]
END

while [[ $# -gt 0 ]]; do
  key="$1"
  case $key in
    --group_id)
      GROUP_ID="$2"
      shift
      shift
      ;;
    --java_expansion_service_jar)
      JAVA_EXPANSION_SERVICE_JAR="$2"
      shift
      shift
      ;;
    --java_expansion_service_allowlist_file)
      JAVA_EXPANSION_SERVICE_ALLOWLIST_FILE="$2"
      shift
      shift
      ;;
    --python_expansion_service_allowlist_glob)
      PYTHON_EXPANSION_SERVICE_ALLOWLIST_GLOB="$2"
      shift
      shift
      ;;
    --python_virtualenv_dir)
      PYTHON_VIRTUALENV_DIR="$2"
      shift
      shift
      ;;
    --python_expansion_service_module)
      PYTHON_EXPANSION_SERVICE_MODULE="$2"
      shift
      shift
      ;;
    --java_port)
      JAVA_PORT="$2"
      shift
      shift
      ;;
    --python_port)
      PYTHON_PORT="$2"
      shift
      shift
      ;;
    start)
      STARTSTOP="$1"
      shift
      ;;
    stop)
      STARTSTOP="$1"
      shift
      ;;
    *)
      echo "Unknown option: $1"
      echo "$USAGE"
      exit 1
      ;;
  esac
done

FILE_BASE="beam-expansion-service"
if [ -n "$GROUP_ID" ]; then
  FILE_BASE="$FILE_BASE-$GROUP_ID"
fi

TEMP_DIR=/tmp
pid=$TEMP_DIR/$FILE_BASE.pid
lock=$TEMP_DIR/$FILE_BASE.lock

# Check whether flock exists since some OS distributions (like MacOS)
# don't have it by default
command -v flock >/dev/null 2>&1
CHECK_FLOCK=$?

if [[ $CHECK_FLOCK -eq 0 ]]; then
  exec 200>$lock
  if ! flock -n 200; then
    echo "script already running."
    exit 0
  fi
fi

wait_for_port() {
  local port=$1
  local srv_pid=$2
  local name=$3
  local logfile=$4
  local elapsed=0
  while [ $elapsed -lt 60 ]; do
    if ! kill -0 "$srv_pid" >/dev/null 2>&1; then
      echo "$name expansion service (pid $srv_pid) exited prematurely. Logs:"
      cat "$logfile" || true
      return 1
    fi
    if python -c "import socket, sys; s = socket.socket(); res = s.connect_ex(('localhost', int(sys.argv[1]))); s.close(); sys.exit(0 if res == 0 else 1)" "$port" >/dev/null 2>&1; then
      echo "$name expansion service is listening @ $port"
      return 0
    fi
    sleep 1
    elapsed=$((elapsed + 1))
  done
  echo "Timed out waiting 60s for $name expansion service @ $port. Logs:"
  cat "$logfile" || true
  return 1
}

case $STARTSTOP in
  start)
    if [ -f "$pid" ]; then
      echo "services already running."
      exit 0
    fi

    source $PYTHON_VIRTUALENV_DIR/bin/activate

    echo "Launching Java expansion service @ $JAVA_PORT"
    if [[ $CHECK_FLOCK -eq 0 ]]; then
      java -jar $JAVA_EXPANSION_SERVICE_JAR $JAVA_PORT --javaClassLookupAllowlistFile="$JAVA_EXPANSION_SERVICE_ALLOWLIST_FILE" >$TEMP_DIR/$FILE_BASE-java.log 2>&1 </dev/null 200>&- &
    else
      java -jar $JAVA_EXPANSION_SERVICE_JAR $JAVA_PORT --javaClassLookupAllowlistFile="$JAVA_EXPANSION_SERVICE_ALLOWLIST_FILE" >$TEMP_DIR/$FILE_BASE-java.log 2>&1 </dev/null &
    fi
    java_pid=$!
    if kill -0 $java_pid >/dev/null 2>&1; then
      echo $java_pid >> $pid
    else
      echo "Can't start Java expansion service."
      exit 1
    fi

    echo "Launching Python expansion service @ $PYTHON_PORT"
    if [[ $CHECK_FLOCK -eq 0 ]]; then
      python -m $PYTHON_EXPANSION_SERVICE_MODULE -p $PYTHON_PORT --fully_qualified_name_glob "$PYTHON_EXPANSION_SERVICE_ALLOWLIST_GLOB" >$TEMP_DIR/$FILE_BASE-python.log 2>&1 </dev/null 200>&- &
    else
      python -m $PYTHON_EXPANSION_SERVICE_MODULE -p $PYTHON_PORT --fully_qualified_name_glob "$PYTHON_EXPANSION_SERVICE_ALLOWLIST_GLOB" >$TEMP_DIR/$FILE_BASE-python.log 2>&1 </dev/null &
    fi
    python_pid=$!
    if kill -0 $python_pid >/dev/null 2>&1; then
      echo $python_pid >> $pid
    else
      echo "Can't start Python expansion service."
      exit 1
    fi

    wait_for_port "$JAVA_PORT" "$java_pid" "Java" "$TEMP_DIR/$FILE_BASE-java.log" || exit 1
    wait_for_port "$PYTHON_PORT" "$python_pid" "Python" "$TEMP_DIR/$FILE_BASE-python.log" || exit 1
    ;;
  stop)
    if [ -f "$pid" ]; then
      while read stop_pid; do
        if kill -0 $stop_pid >/dev/null 2>&1; then
          echo "Stopping expansion service pid: $stop_pid."
          kill $stop_pid
        else
          echo "Skipping invalid pid: $stop_pid."
        fi
      done < $pid
      rm $pid
    fi
    ;;
esac

if [[ $CHECK_FLOCK -eq 0 ]]; then
  flock -u 200
fi