[Bugfix] allow pending data to send before exit (#98)

Author: tom-pytel

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()
 
 

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()

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)
 

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)
 

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)