Handle SIGINT and SIGTERM in yielding listener

Signals INT and TERM were not correctly handled in the 'while True' loop
of the yielding listener function. Now, the loop breaks as expected,
terminating the listener thread and the application.
This commit is contained in:
Dominik Pataky 2020-08-01 10:46:35 +02:00
parent 5cdb514ffc
commit 81d57f3c4c

View file

@ -12,6 +12,7 @@ import gzip
import json import json
import logging import logging
import queue import queue
import signal
import socket import socket
import socketserver import socketserver
import threading import threading
@ -113,6 +114,7 @@ class ThreadedNetFlowListener(threading.Thread):
def run(self): def run(self):
# Process packets from the queue # Process packets from the queue
try: try:
# TODO: use per-client templates
templates = {"netflow": {}, "ipfix": {}} templates = {"netflow": {}, "ipfix": {}}
to_retry = [] to_retry = []
while not self._shutdown.is_set(): while not self._shutdown.is_set():
@ -171,11 +173,20 @@ class ThreadedNetFlowListener(threading.Thread):
def get_export_packets(host: str, port: int) -> ParsedPacket: def get_export_packets(host: str, port: int) -> ParsedPacket:
"""A threaded generator that will yield ExportPacket objects until it is killed """A threaded generator that will yield ExportPacket objects until it is killed
""" """
def handle_signal(s, f):
logger.debug("Received signal {}, raising StopIteration".format(s))
raise StopIteration
signal.signal(signal.SIGTERM, handle_signal)
signal.signal(signal.SIGINT, handle_signal)
listener = ThreadedNetFlowListener(host, port) listener = ThreadedNetFlowListener(host, port)
listener.start() listener.start()
try: try:
while True: while True:
yield listener.get() yield listener.get()
except StopIteration:
pass
finally: finally:
listener.stop() listener.stop()
listener.join() listener.join()