From 8360c54e7095a7dfd960c94aa757b1edc0422d6c Mon Sep 17 00:00:00 2001
From: Xin Li <137219293+xinlili-statsig@users.noreply.github.com>
Date: Tue, 10 Sep 2024 17:40:53 -0700
Subject: [PATCH] wrap error (#311)

Wrap error from interrupt streaming
---
 .github/workflows/kong.yml       |  2 +-
 statsig/grpc_websocket_worker.py | 24 ++++++++++++++++++++----
 statsig/statsig_network.py       |  1 -
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/kong.yml b/.github/workflows/kong.yml
index d696074..b206068 100644
--- a/.github/workflows/kong.yml
+++ b/.github/workflows/kong.yml
@@ -44,7 +44,7 @@ jobs:
 
       - name: Upload base benchmark score
         if: github.ref == 'refs/heads/main'
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: python-perf
           path: /tmp/perf/python_perf_score.txt
diff --git a/statsig/grpc_websocket_worker.py b/statsig/grpc_websocket_worker.py
index f8357b0..8c98ca5 100644
--- a/statsig/grpc_websocket_worker.py
+++ b/statsig/grpc_websocket_worker.py
@@ -149,10 +149,7 @@ def _start_listening(self, dcs_data_cb: Callable, since_time=0, timeout=None):
     def _listen_for_dcs(self, since_time=0):
         try:
             if self.dcs_stream is not None:
-                initial_metadata = self.dcs_stream.initial_metadata()
-                for metadata in initial_metadata:
-                    if metadata.key == "x-sfp-hostname":
-                        self.server_host_name = metadata.value
+                self.get_stream_metadata()
                 for response in self.dcs_stream:
                     if self.retrying:
                         self.retrying = False
@@ -279,6 +276,25 @@ def on_reconnect(self):
             if self._backup_callbacks:
                 self._backup_callbacks.cancel_backup()
 
+    def get_stream_metadata(self):
+        try:
+            if self.dcs_stream is not None:
+                initial_metadata = self.dcs_stream.initial_metadata()
+                for metadata in initial_metadata:
+                    if metadata.key == "x-sfp-hostname":
+                        self.server_host_name = metadata.value
+        except Exception as error:
+            self.error_boundary.log_exception(
+                "grpcWebSocket: connection error",
+                error,
+                {
+                    "retryAttempt": self.retry_limit - self.remaining_retry,
+                    "hostName": socket.gethostname(),
+                    "sfpHostName": self.server_host_name,
+                },
+                True,
+            )
+
     def shutdown(self) -> None:
         self.is_shutting_down = True
         if self._backup_callbacks:
diff --git a/statsig/statsig_network.py b/statsig/statsig_network.py
index ff4fc8b..dff9552 100644
--- a/statsig/statsig_network.py
+++ b/statsig/statsig_network.py
@@ -84,7 +84,6 @@ def __init__(
         self.http_worker: IStatsigNetworkWorker = defaultHttpWorker
         for endpoint, config in options.proxy_configs.items():
             protocol = config.protocol
-            worker = defaultHttpWorker
             if protocol == NetworkProtocol.GRPC:
                 self.load_grpc_worker(endpoint, config)
             elif protocol == NetworkProtocol.GRPC_WEBSOCKET: