From 593deb92e8ebd4e77e482c567d97b6ee496ac378 Mon Sep 17 00:00:00 2001
From: Eric Wong <normalperson@yhbt.net>
Date: Thu, 16 Jun 2011 23:57:31 +0000
Subject: ensure at_exit handlers run on graceful shutdown

rescuing from SystemExit and exit()-ing again is ugly, but
changes made to lower stack depth positively affect _everyone_
so we'll tolerate some ugliness here.

We'll need to disable graceful exit for some tests, too...
---
 lib/unicorn/http_server.rb |  6 ++----
 t/t0020-at_exit-handler.sh | 49 ++++++++++++++++++++++++++++++++++++++++++++++
 test/unit/test_server.rb   |  2 +-
 test/unit/test_signals.rb  |  2 +-
 test/unit/test_upload.rb   |  2 +-
 5 files changed, 54 insertions(+), 7 deletions(-)
 create mode 100644 t/t0020-at_exit-handler.sh

diff --git a/lib/unicorn/http_server.rb b/lib/unicorn/http_server.rb
index 0a9af86..1489495 100644
--- a/lib/unicorn/http_server.rb
+++ b/lib/unicorn/http_server.rb
@@ -474,14 +474,12 @@ class Unicorn::HttpServer
       else
         after_fork_internal
         worker_loop(worker)
-        exit(0)
+        exit
       end
     end
-    rescue SystemExit => e
-      exit!(e.status)
     rescue => e
       @logger.error(e) rescue nil
-      exit!(1)
+      exit!
   end
 
   def maintain_worker_count
diff --git a/t/t0020-at_exit-handler.sh b/t/t0020-at_exit-handler.sh
new file mode 100644
index 0000000..fda14b5
--- /dev/null
+++ b/t/t0020-at_exit-handler.sh
@@ -0,0 +1,49 @@
+#!/bin/sh
+. ./test-lib.sh
+
+t_plan 5 "at_exit/END handlers work as expected"
+
+t_begin "setup and startup" && {
+	unicorn_setup
+	cat >> $unicorn_config <<EOF
+at_exit { \$stdout.syswrite("#{Process.pid} BOTH\\n") }
+END { \$stdout.syswrite("#{Process.pid} END BOTH\\n") }
+after_fork do |_,_|
+  at_exit { \$stdout.syswrite("#{Process.pid} WORKER ONLY\\n") }
+  END { \$stdout.syswrite("#{Process.pid} END WORKER ONLY\\n") }
+end
+EOF
+
+	unicorn -D pid.ru -c $unicorn_config
+	unicorn_wait_start
+}
+
+t_begin "read worker PID" && {
+	worker_pid=$(curl -sSf http://$listen/)
+	t_info "worker_pid=$worker_pid"
+}
+
+t_begin "issue graceful shutdown (SIGQUIT) and wait for termination" && {
+	kill -QUIT $unicorn_pid
+
+	while kill -0 $unicorn_pid >/dev/null 2>&1
+	do
+		sleep 1
+	done
+}
+
+t_begin "check stderr" && check_stderr
+
+dbgcat r_err
+dbgcat r_out
+
+t_begin "all at_exit handlers ran" && {
+	grep "$worker_pid BOTH" $r_out
+	grep "$unicorn_pid BOTH" $r_out
+	grep "$worker_pid END BOTH" $r_out
+	grep "$unicorn_pid END BOTH" $r_out
+	grep "$worker_pid WORKER ONLY" $r_out
+	grep "$worker_pid END WORKER ONLY" $r_out
+}
+
+t_done
diff --git a/test/unit/test_server.rb b/test/unit/test_server.rb
index 88d7aba..403bd1f 100644
--- a/test/unit/test_server.rb
+++ b/test/unit/test_server.rb
@@ -39,7 +39,7 @@ class WebServerTest < Test::Unit::TestCase
     redirect_test_io do
       wait_workers_ready("test_stderr.#$$.log", 1)
       File.truncate("test_stderr.#$$.log", 0)
-      @server.stop(true)
+      @server.stop(false)
     end
     reset_sig_handlers
   end
diff --git a/test/unit/test_signals.rb b/test/unit/test_signals.rb
index 71cf8f4..5a3b519 100644
--- a/test/unit/test_signals.rb
+++ b/test/unit/test_signals.rb
@@ -124,7 +124,7 @@ class SignalsTest < Test::Unit::TestCase
     assert diff > 1.0, "diff was #{diff.inspect}"
     assert diff < 60.0
     ensure
-      Process.kill(:QUIT, pid) rescue nil
+      Process.kill(:TERM, pid) rescue nil
   end
 
   def test_response_write
diff --git a/test/unit/test_upload.rb b/test/unit/test_upload.rb
index e2c103a..7777e00 100644
--- a/test/unit/test_upload.rb
+++ b/test/unit/test_upload.rb
@@ -53,7 +53,7 @@ class UploadTest < Test::Unit::TestCase
   end
 
   def teardown
-    redirect_test_io { @server.stop(true) } if @server
+    redirect_test_io { @server.stop(false) } if @server
     @random.close
     reset_sig_handlers
   end
-- 
cgit v1.2.3-24-ge0c7