77from logging import basicConfig , getLevelName , getLogger
88from multiprocessing import Process
99from pathlib import Path
10+ from queue import Queue
1011from time import sleep
1112from typing import Any , Generator , List
1213
14+ from watchdog .observers import Observer
15+
1316from taskiq .abc .broker import AsyncBroker
1417from taskiq .cli .args import TaskiqArgs
1518from taskiq .cli .async_task_runner import async_listen_messages
19+ from taskiq .cli .watcher import FileWatcher
1620
1721try :
1822 import uvloop # noqa: WPS433
2529
2630restart_workers = True
2731worker_processes : List [Process ] = []
32+ observer = Observer ()
33+ reload_queue : "Queue[bool]" = Queue (- 1 )
2834
2935
3036def signal_handler (_signal : int , _frame : Any ) -> None :
@@ -45,9 +51,31 @@ def signal_handler(_signal: int, _frame: Any) -> None:
4551 # This is how we kill children,
4652 # by sending SIGINT to child processes.
4753 if process .pid is None :
48- process . kill ()
49- else :
54+ continue
55+ try :
5056 os .kill (process .pid , signal .SIGINT )
57+ except ProcessLookupError :
58+ continue
59+ process .join ()
60+ if observer .is_alive ():
61+ observer .stop ()
62+ observer .join ()
63+
64+
65+ def schedule_workers_reload () -> None :
66+ """
67+ Function to schedule workers to restart.
68+
69+ This function adds worker ids to the queue.
70+
71+ This queue is later read in watcher loop.
72+ """
73+ global worker_processes # noqa: WPS420
74+ global reload_queue # noqa: WPS420
75+
76+ reload_queue .put (True )
77+ logger .info ("Scheduled workers reload." )
78+ reload_queue .join ()
5179
5280
5381@contextmanager
@@ -212,13 +240,16 @@ def interrupt_handler(_signum: int, _frame: Any) -> None:
212240 loop .run_until_complete (shutdown_broker (broker , args .shutdown_timeout ))
213241
214242
215- def watch_workers_restarts (args : TaskiqArgs ) -> None :
243+ def watcher_loop (args : TaskiqArgs ) -> None : # noqa: C901, WPS213
216244 """
217245 Infinate loop for main process.
218246
219247 This loop restarts worker processes
220248 if they exit with error returncodes.
221249
250+ Also it reads process ids from reload_queue
251+ and reloads workers if they were scheduled to reload.
252+
222253 :param args: cli arguements.
223254 """
224255 global worker_processes # noqa: WPS420
@@ -228,6 +259,21 @@ def watch_workers_restarts(args: TaskiqArgs) -> None:
228259 # List of processes to remove.
229260 sleep (1 )
230261 process_to_remove = []
262+ if not reload_queue .empty ():
263+ while not reload_queue .empty ():
264+ reload_queue .get ()
265+ reload_queue .task_done ()
266+
267+ for worker_id , worker in enumerate (worker_processes ):
268+ worker .terminate ()
269+ worker .join ()
270+ worker_processes [worker_id ] = Process (
271+ target = start_listen ,
272+ kwargs = {"args" : args },
273+ name = f"worker-{ worker_id } " ,
274+ )
275+ worker_processes [worker_id ].start ()
276+
231277 for worker_id , worker in enumerate (worker_processes ):
232278 if worker .is_alive ():
233279 continue
@@ -241,14 +287,13 @@ def watch_workers_restarts(args: TaskiqArgs) -> None:
241287 worker_processes [worker_id ].start ()
242288 else :
243289 logger .info ("Worker-%s terminated." , worker_id )
244- worker .join ()
245290 process_to_remove .append (worker )
246291
247292 for dead_process in process_to_remove :
248293 worker_processes .remove (dead_process )
249294
250295
251- def run_worker (args : TaskiqArgs ) -> None :
296+ def run_worker (args : TaskiqArgs ) -> None : # noqa: WPS213
252297 """
253298 This function starts worker processes.
254299
@@ -279,7 +324,17 @@ def run_worker(args: TaskiqArgs) -> None:
279324 )
280325 worker_processes .append (work_proc )
281326
327+ if args .reload :
328+ observer .schedule (
329+ FileWatcher (
330+ callback = schedule_workers_reload ,
331+ use_gitignore = not args .no_gitignore ,
332+ ),
333+ path = "." ,
334+ recursive = True ,
335+ )
336+ observer .start ()
282337 signal .signal (signal .SIGINT , signal_handler )
283338 signal .signal (signal .SIGTERM , signal_handler )
284339
285- watch_workers_restarts (args = args )
340+ watcher_loop (args = args )
0 commit comments