1 files changed, 42 insertions, 11 deletions

diff --git a/lib/iomaster.php b/lib/iomaster.php
index 94abdeefc..1f6c31ee7 100644
--- a/lib/iomaster.php
+++ b/lib/iomaster.php
@@ -38,6 +38,9 @@ abstract class IoMaster
     protected $pollTimeouts = array();
     protected $lastPoll = array();
 
+    public $shutdown = false; // Did we do a graceful shutdown?
+    public $respawn = true; // Should we respawn after shutdown?
+
     /**
      * @param string $id process ID to use in logging/monitoring
      */
@@ -148,7 +151,7 @@ abstract class IoMaster
         $this->logState('init');
         $this->start();
 
-        while (true) {
+        while (!$this->shutdown) {
             $timeouts = array_values($this->pollTimeouts);
             $timeouts[] = 60; // default max timeout
 
@@ -200,16 +203,7 @@ abstract class IoMaster
             $this->logState('idle');
             $this->idle();
 
-            $memoryLimit = $this->softMemoryLimit();
-            if ($memoryLimit > 0) {
-                $usage = memory_get_usage();
-                if ($usage > $memoryLimit) {
-                    common_log(LOG_INFO, "Queue thread hit soft memory limit ($usage > $memoryLimit); gracefully restarting.");
-                    break;
-                } else if (common_config('queue', 'debug_memory')) {
-                    common_log(LOG_DEBUG, "Memory usage $usage");
-                }
-            }
+            $this->checkMemory();
         }
 
         $this->logState('shutdown');
@@ -217,6 +211,24 @@ abstract class IoMaster
     }
 
     /**
+     * Check runtime memory usage, possibly triggering a graceful shutdown
+     * and thread respawn if we've crossed the soft limit.
+     */
+    protected function checkMemory()
+    {
+        $memoryLimit = $this->softMemoryLimit();
+        if ($memoryLimit > 0) {
+            $usage = memory_get_usage();
+            if ($usage > $memoryLimit) {
+                common_log(LOG_INFO, "Queue thread hit soft memory limit ($usage > $memoryLimit); gracefully restarting.");
+                $this->requestRestart();
+            } else if (common_config('queue', 'debug_memory')) {
+                common_log(LOG_DEBUG, "Memory usage $usage");
+            }
+        }
+    }
+
+    /**
      * Return fully-parsed soft memory limit in bytes.
      * @return intval 0 or -1 if not set
      */
@@ -358,5 +370,24 @@ abstract class IoMaster
         $owners[] = "thread:" . $this->id;
         $this->monitor->stats($key, $owners);
     }
+
+    /**
+     * For IoManagers to request a graceful shutdown at end of event loop.
+     */
+    public function requestShutdown()
+    {
+        $this->shutdown = true;
+        $this->respawn = false;
+    }
+
+    /**
+     * For IoManagers to request a graceful restart at end of event loop.
+     */
+    public function requestRestart()
+    {
+        $this->shutdown = true;
+        $this->respawn = true;
+    }
+
 }