11#! /usr/bin/env bash
2+ set -euo pipefail
3+ IFS=$' \n\t '
4+
25
36server_args=${SERVER_ARGS:- }
47worker_args=${WORKER_ARGS:- }
58server_addr=${SERVER_ADDR:- 127.0.0.1: 27183}
69spawn_server=${SPAWN_SERVER:- true}
710spawn_worker=${SPAWN_WORKER:- true}
11+ tm_loglevel=${TM_LOGLEVEL:- info}
12+
13+ # default: nproc - 1
14+ default_nworkers=$(( $(nproc) - 1 ))
15+ nworkers=$default_nworkers
16+ help_flag=false
17+
18+ # ################### helpers
19+ function check_posint() {
20+ local re=' ^[0-9]+$'
21+ local mynum=" $1 "
22+ local option=" $2 "
23+
24+ if ! [[ " $mynum " =~ $re ]] ; then
25+ (echo -n " Error in option '$option ': " >&2 )
26+ (echo " must be a positive integer, got $mynum ." >&2 )
27+ exit 1
28+ fi
29+
30+ if ! [ " $mynum " -gt 0 ] ; then
31+ (echo " Error in option '$option ': must be positive, got $mynum ." >&2 )
32+ exit 1
33+ fi
34+ }
35+ # ################### end: helpers
36+
37+ # ################### help
38+ function short_usage() {
39+ (>&2 echo \
40+ " Usage:
41+ $( basename " $0 " ) [options]" )
42+ }
43+
44+ function usage() {
45+ (>&2 short_usage )
46+ (>&2 echo \
47+ "
48+ Entrypoint for task-maker-rust Docker image.
49+
50+ Options:
51+ -j, --jobs NWORKERS Number of workers to launch [default: <nproc>-1].
52+ -h, --help Show this help and exits.
53+ " )
54+ }
55+
56+ # the leading ":" lets us handle errors ourselves
57+ while getopts " :hj:-:" OPT; do
58+ if [ " $OPT " = " -" ]; then
59+ # Long option handling
60+ longopt=" ${OPTARG%% =* } "
61+ if [ " $longopt " != " $OPTARG " ]; then
62+ # do not use
63+ (echo " Error. Invalid option: --$OPTARG . " \
64+ " Use --opt OPT instead of --opt=OPT." >&2 )
65+ exit 1
66+ else
67+ # form: --long VALUE (value may be in next argv)
68+ OPTARG=
69+ fi
70+
71+ case " $longopt " in
72+ help)
73+ OPT=" h"
74+ ;;
75+
76+ jobs)
77+ OPT=" j"
78+ # If no "=VALUE", try to take the next argv as the argument
79+ next=" ${! OPTIND} "
80+ if [ -n " $next " ] && [[ " $next " != -* ]]; then
81+ OPTARG=" $next "
82+ OPTIND=$(( OPTIND + 1 ))
83+ fi
84+ ;;
85+
86+ * )
87+ echo " Error. Bad long option --$longopt " >&2
88+ exit 1
89+ ;;
90+ esac
91+ fi
92+
93+ case " $OPT " in
94+ h)
95+ help_flag=true
96+ ;;
97+
98+ j)
99+ check_posint " $OPTARG " ' -j/--jobs'
100+ nworkers=" $OPTARG "
101+ ;;
8102
9- export RUST_LOG=info
10- export RUST_BACKTRACE=1
103+ \? )
104+ echo " Error. Invalid option: -$OPTARG " >&2
105+ exit 1
106+ ;;
107+
108+ :)
109+ echo " Error. Option $OPTARG requires an argument." >&2
110+ exit 1
111+ ;;
112+ esac
113+ done
114+
115+ if $help_flag ; then
116+ usage
117+ exit 0
118+ fi
119+ # ################### end: help
11120
12121server_store=$( mktemp -d tmserver.XXXXXXX -p /tmp)
13- worker_store=$( mktemp -d tmworker.XXXXXXX -p /tmp)
14122
15- function spawn_server() {
16- task-maker-rust --store-dir " $server_store " $server_args --server
123+ # create nworkers worker stores directories
124+ worker_base=$( mktemp -u tmworker.XXXXXXX -p /tmp)
125+ declare -a worker_stores=()
126+ for (( i= 1 ; i<= nworkers; i++ )) ; do
127+ idx=$( printf " %02d" " $i " )
128+ dir=" ${worker_base} -${idx} "
129+ mkdir -p " $dir "
130+ worker_stores+=(" $dir " )
131+ done
132+
133+ # cleanup exit trap
134+ function cleanup {
135+ # remove server and worker stores
136+ rm -rf /tmp/tmserver.*
137+ for wstore in $( find /tmp -type d -name ' tmworker.*' ) ; do
138+ rm -rf $wstore
139+ done
17140}
18- function spawn_worker() {
19- task-maker-rust --store-dir " $worker_store " $worker_args --worker $server_addr
141+ trap cleanup EXIT
142+
143+ function loglevel_verbosity_flag() {
144+ case " $1 " in
145+ error)
146+ echo " "
147+ ;;
148+ warn)
149+ echo " -v"
150+ ;;
151+ info)
152+ echo " -vv"
153+ ;;
154+ debug)
155+ echo " -vvv"
156+ ;;
157+ * )
158+ (echo " Unknown level: $1 . Default to 'error'." >&2 );
159+ return 1
160+ ;;
161+ esac
20162}
21163
164+ function spawn_tmserver() {
165+ local v=" $1 "
166+
167+ if [[ " ${tm_loglevel} " == ' debug' ]]; then
168+ set -x
169+ fi
170+
171+ task-maker-tools ${v: +$v } \
172+ server \
173+ --store-dir " $server_store " \
174+ $server_args
175+ }
176+
177+ function spawn_tmworker() {
178+ local v=" $1 "
179+ local store_dir=" $2 "
180+
181+ if [[ " ${tm_loglevel} " == ' debug' ]]; then
182+ set -x
183+ fi
184+
185+ task-maker-tools ${v: +$v } \
186+ worker \
187+ --store-dir " $store_dir " \
188+ $worker_args \
189+ " $server_addr "
190+ }
191+
192+ verbosity_flag=$( loglevel_verbosity_flag " ${tm_loglevel} " )
193+
22194# worker only
23- if [[ $spawn_server != true && $spawn_worker == true ]]; then
24- spawn_worker
195+ if [[ $spawn_server == false && $spawn_worker == true ]]; then
196+ for (( i= 0 ; i< ${# worker_stores[@]} ; i++ )) ; do
197+ spawn_tmworker " $verbosity_flag " " ${worker_stores[$i]} " &
198+ done
199+ wait
200+
25201# server only
26- elif [[ $spawn_server == true && $spawn_worker != true ]]; then
27- spawn_server
202+ elif [[ $spawn_server == true && $spawn_worker == false ]]; then
203+ spawn_tmserver " $verbosity_flag "
204+
28205# server+worker
29206elif [[ $spawn_server == true && $spawn_worker == true ]]; then
30207 # run the workers in background, but wait for the server
31- ( sleep 2s && spawn_worker ) &
32- spawn_server
208+ (
209+ sleep 2s
210+ for (( i= 0 ; i< ${# worker_stores[@]} ; i++ )) ; do
211+ spawn_tmworker " $verbosity_flag " " ${worker_stores[$i]} " &
212+ done
213+ wait
214+ ) &
215+ spawn_tmserver " $verbosity_flag "
216+
33217# nothing to spawn
34218else
35219 bash
36- fi
220+ fi
221+
222+ exit 0
0 commit comments