diff --git a/skywalking/agent/__init__.py b/skywalking/agent/__init__.py index 3a6b001a..e8544dc2 100644 --- a/skywalking/agent/__init__.py +++ b/skywalking/agent/__init__.py @@ -15,12 +15,13 @@ # limitations under the License. # -from skywalking.loggings import logger +import atexit from queue import Queue from threading import Thread, Event from typing import TYPE_CHECKING -from skywalking import config, plugins +from skywalking import config, plugins, loggings +from skywalking.loggings import logger from skywalking.agent.protocol import Protocol if TYPE_CHECKING: @@ -66,20 +67,27 @@ def __init(): plugins.install() +def __fini(): + __protocol.report(__queue, False) + __queue.join() + + def start(): global __started if __started: raise RuntimeError('the agent can only be started once') - from skywalking import loggings loggings.init() config.finalize() __started = True __init() __heartbeat_thread.start() __report_thread.start() + atexit.register(__fini) def stop(): + atexit.unregister(__fini) + __fini() __finished.set() diff --git a/skywalking/agent/protocol/__init__.py b/skywalking/agent/protocol/__init__.py index 8237f327..0f6e62e5 100644 --- a/skywalking/agent/protocol/__init__.py +++ b/skywalking/agent/protocol/__init__.py @@ -26,5 +26,5 @@ def connected(self): def heartbeat(self): raise NotImplementedError() - def report(self, queue: Queue): + def report(self, queue: Queue, block: bool = True): raise NotImplementedError() diff --git a/skywalking/agent/protocol/grpc.py b/skywalking/agent/protocol/grpc.py index 0a7da16f..f5468a44 100644 --- a/skywalking/agent/protocol/grpc.py +++ b/skywalking/agent/protocol/grpc.py @@ -67,10 +67,10 @@ def on_error(self): self.channel.unsubscribe(self._cb) self.channel.subscribe(self._cb, try_to_connect=True) - def report(self, queue: Queue): + def report(self, queue: Queue, block: bool = True): def generator(): while True: - segment = queue.get() # type: Segment + segment = queue.get(block=block) # type: Segment logger.debug('reporting segment %s', segment) diff --git a/skywalking/agent/protocol/http.py b/skywalking/agent/protocol/http.py index e97cb227..331f71a9 100644 --- a/skywalking/agent/protocol/http.py +++ b/skywalking/agent/protocol/http.py @@ -38,10 +38,10 @@ def heartbeat(self): def connected(self): return True - def report(self, queue: Queue): + def report(self, queue: Queue, block: bool = True): def generator(): while True: - segment = queue.get() # type: Segment + segment = queue.get(block=block) # type: Segment logger.debug('reporting segment %s', segment) diff --git a/skywalking/agent/protocol/kafka.py b/skywalking/agent/protocol/kafka.py index 405a7288..8e6e75af 100644 --- a/skywalking/agent/protocol/kafka.py +++ b/skywalking/agent/protocol/kafka.py @@ -42,10 +42,10 @@ def connected(self): def heartbeat(self): self.service_management.send_heart_beat() - def report(self, queue: Queue): + def report(self, queue: Queue, block: bool = True): def generator(): while True: - segment = queue.get() # type: Segment + segment = queue.get(block=block) # type: Segment logger.debug('reporting segment %s', segment)