Add flags to bin/docker-worker-celery (#2621)

* add flags to docker-worker-celery and rename --with-beat to --with-scheduler

* Clean script up

Co-authored-by: Michael Matloka <dev@twixes.com>

Procfile

@@ -1,5 +1,5 @@
 release: REDIS_URL='redis://' python manage.py migrate
 web: gunicorn posthog.wsgi --log-file -
 worker: ./bin/docker-worker
-celeryworker: ./bin/docker-worker-celery --with-beat # optional
+celeryworker: ./bin/docker-worker-celery --with-scheduler # optional
 pluginworker: ./bin/plugin-server # optional

bin/docker-worker

@@ -15,5 +15,5 @@ then
   echo "⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️⚠️"
 else
   ./bin/plugin-server &
-  ./bin/docker-worker-celery --with-beat
+  ./bin/docker-worker-celery --with-scheduler
 fi

bin/docker-worker-celery

@@ -1,17 +1,81 @@
 #!/bin/bash
 set -e
 
-if [ "$1" == "--with-beat" ]; then
+help () {
+  echo "$0 - start PostHog's Celery worker"
+  echo
+  echo "$0 [options]"
+  echo
+  echo "Options:"
+  echo "  --help, -h            show this brief help"
+  echo "  --with-scheduler      start RedBeat, the Celery scheduler (deprecates --with-beat)"
+  echo "  --concurrency=<N>     start N workers (overrides env var WEB_CONCURRENCY)"
+  echo
+  echo "Advanced Celery options (disabled by default):"
+  echo "  --with-gossip         start Celery gossip (useful for Prometheus)"
+  echo "  --with-heartbeat      start Celery internal heartbeat (normally not useful)"
+  echo "  --with-mingle         start Celery mingle (normally not useful)"
+  exit 0
+}
+
+with_scheduler=false
+with_gossip=false
+with_heartbeat=false
+with_mingle=false
+
+while test $# -gt 0; do
+  case "$1" in
+    -h|--help)
+      help
+      ;;
+    --with-scheduler)
+      with_scheduler=true
+      shift
+      ;;
+    --with-beat) # Deprecated since the name is too similar to "heartbeat"
+      echo "⚠️ Using docker-worker-celery with --with-beat. This argument is deprecated. Use --with-scheduler instead!"
+      with_scheduler=true
+      shift
+      ;;
+    --with-gossip)
+      with_gossip=true
+      shift
+      ;;
+    --with-heartbeat)
+      with_heartbeat=true
+      shift
+      ;;
+    --with-mingle)
+      with_mingle=true
+      shift
+      ;;
+    --concurrency*)
+      export WEB_CONCURRENCY=`echo $1 | sed -e 's/^[^=]*=//g'`
+      shift
+      ;;
+    *)
+      break
+      ;;
+  esac
+done
+
+if [ "$with_scheduler" == "true" ]; then
   ./bin/docker-worker-beat &
 fi
 
-# On heroku $WEB_CONCURRENCY contains suggested nr of forks per dyno type
+FLAGS=()
+[ "$with_gossip" == "false" ]    && FLAGS+=("--without-gossip")
+[ "$with_mingle" == "false" ]    && FLAGS+=("--without-mingle")
+[ "$with_heartbeat" == "false" ] && FLAGS+=("--without-heartbeat")
+
+# On Heroku $WEB_CONCURRENCY contains suggested number of forks per dyno type
 # https://github.com/heroku/heroku-buildpack-python/blob/main/vendor/WEB_CONCURRENCY.sh
-if [[ -z "${WEB_CONCURRENCY}" ]]; then
-  celery -A posthog worker --without-heartbeat --without-gossip --without-mingle
-else
-  celery -A posthog worker --without-heartbeat --without-gossip --without-mingle --concurrency $WEB_CONCURRENCY
-fi
+[[ -n "${WEB_CONCURRENCY}" ]]    && FLAGS+=" --concurrency $WEB_CONCURRENCY"
+
+echo
+echo "celery -A posthog worker ${FLAGS[*]}"
+echo
+celery -A posthog worker ${FLAGS[*]}
 
 # Stop the beat!
 trap 'kill $(jobs -p)' EXIT