From 1ae9c50fbbe1fd4387659246e1d7d79d98c2bc01 Mon Sep 17 00:00:00 2001
From: gaogaotiantian <gaogaotiantian@hotmail.com>
Date: Mon, 17 Jul 2023 11:43:43 -0700
Subject: [PATCH] Fixed exception handling for vizviewer

---
 src/viztracer/viewer.py | 12 +++++++++---
 tests/test_viewer.py    |  3 +++
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/viztracer/viewer.py b/src/viztracer/viewer.py
index c4b96fea..15d8ea08 100644
--- a/src/viztracer/viewer.py
+++ b/src/viztracer/viewer.py
@@ -14,6 +14,7 @@
 import socketserver
 import subprocess
 import sys
+import traceback
 import threading
 import time
 import urllib.parse
@@ -298,9 +299,14 @@ def __init__(
         super().__init__(daemon=True)
 
     def run(self) -> None:
-        self.retcode = self.view()
-        # If it returns from view(), also set ready
-        self.ready.set()
+        try:
+            self.retcode = self.view()
+        except Exception as e:
+            self.retcode = 1
+            traceback.print_exception(e)
+        finally:
+            # If it returns from view(), also set ready
+            self.ready.set()
 
     def view(self) -> int:
         # Get file data
diff --git a/tests/test_viewer.py b/tests/test_viewer.py
index 50a1d5a8..cdebb0fc 100644
--- a/tests/test_viewer.py
+++ b/tests/test_viewer.py
@@ -457,6 +457,9 @@ def test_directory_max_port(self):
         finally:
             shutil.rmtree(tmp_dir)
 
+    def test_exception(self):
+        self.template(["vizviewer", "--port", "-3", "cmdline_test.json"], success=False, expected_output_file=None)
+
     def test_invalid(self):
         self.template(["vizviewer", "do_not_exist.json"], success=False, expected_output_file=None)
         self.template(["vizviewer", "README.md"], success=False, expected_output_file=None)