From 6d5c32ad4b6078a72b73b695981faec36ef917c0 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Tue, 15 Mar 2011 17:21:35 +1300
Subject: [PATCH] Fix hang on shutdown.

---
 libmproxy/console.py    |  2 +-
 libmproxy/controller.py | 19 ++++++++++++-------
 test/test_flow.py       |  4 ++--
 test/tutils.py          |  1 +
 4 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 9dd80c7ea..d36d87dae 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -1277,7 +1277,7 @@ class ConsoleMaster(flow.FlowMaster):
         slave = controller.Slave(q, self.server)
         slave.start()
         try:
-            while not self._shutdown:
+            while not controller.exit:
                 self.statusbar.redraw()
                 size = self.drawscreen()
                 self.tick(q)
diff --git a/libmproxy/controller.py b/libmproxy/controller.py
index 82e934c4d..52b6ac6b5 100644
--- a/libmproxy/controller.py
+++ b/libmproxy/controller.py
@@ -16,6 +16,8 @@
 import sys
 import Queue, threading
 
+exit = False
+
 #begin nocover
 
 class Msg:
@@ -34,7 +36,12 @@ class Msg:
         self.acked = False
         try:
             masterq.put(self, timeout=3)
-            return self.q.get()
+            while not exit:
+                try:
+                    g = self.q.get(timeout=0.5)
+                except Queue.Empty:
+                    continue
+                return g
         except (Queue.Empty, Queue.Full):
             return None
 
@@ -52,7 +59,6 @@ class Slave(threading.Thread):
 class Master:
     def __init__(self, server):
         self.server = server
-        self._shutdown = False
         self.masterq = None
 
     def tick(self, q):
@@ -73,7 +79,7 @@ class Master:
         self.masterq = q
         slave = Slave(q, self.server)
         slave.start()
-        while not self._shutdown:
+        while not exit:
             self.tick(q)
         self.shutdown()
 
@@ -86,9 +92,8 @@ class Master:
             msg.ack()
 
     def shutdown(self):
-        if not self._shutdown:
-            self._shutdown = True
+        global exit
+        if not exit:
+            exit = True
             if self.server:
                 self.server.shutdown()
-
-
diff --git a/test/test_flow.py b/test/test_flow.py
index 91e3f5f40..a3fa48d7b 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -1,6 +1,6 @@
 import Queue
 from cStringIO import StringIO
-from libmproxy import console, proxy, filt, flow
+from libmproxy import console, proxy, filt, flow, controller
 import tutils
 import libpry
 
@@ -442,7 +442,7 @@ class uFlowMaster(libpry.AutoTree):
         assert fm.do_server_playback(tutils.tflow())
         q = Queue.Queue()
         fm.tick(q)
-        assert fm._shutdown
+        assert controller.exit
 
     def test_stickycookie(self):
         s = flow.State()
diff --git a/test/tutils.py b/test/tutils.py
index ec2c71a80..ae9dea271 100644
--- a/test/tutils.py
+++ b/test/tutils.py
@@ -59,6 +59,7 @@ class TestMaster(controller.Master):
 class ProxyThread(threading.Thread):
     def __init__(self, port, testq):
         self.tmaster = TestMaster(port, testq)
+        controller.exit = False
         threading.Thread.__init__(self)
 
     def run(self):