diff --git a/tools/trace.py b/tools/trace.py
index 9c7cca71e879..2aa096fa3201 100755
--- a/tools/trace.py
+++ b/tools/trace.py
@@ -43,6 +43,7 @@ class Probe(object):
         build_id_enabled = False
         aggregate = False
         symcount = {}
+        done = False
 
         @classmethod
         def configure(cls, args):
@@ -635,7 +636,7 @@ def print_event(self, bpf, cpu, data, size):
                     if self.aggregate:
                         self.print_aggregate_events()
                     sys.stdout.flush()
-                    exit()
+                    Probe.done = True;
 
         def attach(self, bpf, verbose):
                 if len(self.library) == 0:
@@ -895,7 +896,7 @@ def _main_loop(self):
                       "-" if not all_probes_trivial else ""))
                 sys.stdout.flush()
 
-                while True:
+                while not Probe.done:
                         self.bpf.perf_buffer_poll()
 
         def run(self):