fix: grpc timeout segment data loss (#116)

Author: tom-pytel

skywalking/agent/__init__.py

@@ -16,7 +16,7 @@
 #
 
 import atexit
-from queue import Queue
+from queue import Queue, Full
 from threading import Thread, Event
 from typing import TYPE_CHECKING
 
@@ -109,8 +109,7 @@ def connected():
 
 
 def archive(segment: 'Segment'):
-    if __queue.full():
+    try:  # unlike checking __queue.full() then inserting, this is atomic
+        __queue.put(segment, block=False)
+    except Full:
         logger.warning('the queue is full, the segment will be abandoned')
-        return
-
-    __queue.put(segment)

skywalking/agent/protocol/grpc.py

@@ -18,7 +18,8 @@
 import logging
 from skywalking.loggings import logger
 import traceback
-from queue import Queue, Empty
+from queue import Queue, Empty, Full
+from time import time
 
 import grpc
 
@@ -68,10 +69,16 @@ def on_error(self):
         self.channel.subscribe(self._cb, try_to_connect=True)
 
     def report(self, queue: Queue, block: bool = True):
+        start = time()
+        segment = None
+
         def generator():
+            nonlocal segment
+
             while True:
                 try:
-                    segment = queue.get(block=block)  # type: Segment
+                    timeout = max(0, config.QUEUE_TIMEOUT - int(time() - start))  # type: int
+                    segment = queue.get(block=block, timeout=timeout)  # type: Segment
                 except Empty:
                     return
 
@@ -120,5 +127,12 @@ def generator():
 
         try:
             self.traces_reporter.report(generator())
+
         except grpc.RpcError:
             self.on_error()
+
+            if segment:
+                try:
+                    queue.put(segment, block=False)
+                except Full:
+                    pass

skywalking/client/grpc.py

@@ -55,4 +55,4 @@ def __init__(self, channel: grpc.Channel):
         self.report_stub = TraceSegmentReportServiceStub(channel)
 
     def report(self, generator):
-        self.report_stub.collect(generator)
+        self.report_stub.collect(generator, timeout=config.GRPC_TIMEOUT)

skywalking/config.py

@@ -23,6 +23,11 @@
 if TYPE_CHECKING:
     from typing import List
 
+# In order to prevent timeouts and possible segment loss make sure QUEUE_TIMEOUT is always at least few seconds lower
+# than GRPC_TIMEOUT.
+GRPC_TIMEOUT = 300  # type: int
+QUEUE_TIMEOUT = 240  # type: int
+
 RE_IGNORE_PATH = re.compile('^$')  # type: re.Pattern
 
 service_name = os.getenv('SW_AGENT_NAME') or 'Python Service Name'  # type: str