From c6cc7eaa81923b374333e1252ac4a9c888374321 Mon Sep 17 00:00:00 2001
From: Samuel Merritt <sam@swiftstack.com>
Date: Mon, 18 Jun 2012 09:46:54 -0700
Subject: [PATCH] Make swift exit on ctrl-c.

The first time the user presses ctrl-c, all QueueFunctionThreads will
have abort=True set on them. This will cause them to finish the work
they're doing (e.g. finish uploading the current file) but then ignore
any further work and let the process exit. Also, a notification of
this is written to stderr so the user understands why the process
didn't exit immediately.

The second time the user presses ctrl-c, the process will exit
immediately. Any in-progress operations are abandoned.

Change-Id: Ie6927f78726ac1c7998e920cb608682ead10f25b
---
 bin/swift | 30 ++++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/bin/swift b/bin/swift
index 3ba8190..8e6724d 100755
--- a/bin/swift
+++ b/bin/swift
@@ -13,16 +13,17 @@
 # implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import signal
 import socket
 
 from errno import EEXIST, ENOENT
 from hashlib import md5
 from optparse import OptionParser
-from os import environ, listdir, makedirs, utime
+from os import environ, listdir, makedirs, utime, _exit as os_exit
 from os.path import basename, dirname, getmtime, getsize, isdir, join
 from Queue import Empty, Queue
 from sys import argv, exc_info, exit, stderr, stdout
-from threading import enumerate as threading_enumerate, Thread
+from threading import current_thread, enumerate as threading_enumerate, Thread
 from time import sleep
 from traceback import format_exception
 
@@ -67,6 +68,29 @@ def put_errors_from_threads(threads, error_queue):
     return was_error
 
 
+def attempt_graceful_exit(signum, frame):
+    """
+    Try to gracefully shut down. Sets abort=True on all non-main threads.
+
+    More importantly, installs the immediate_exit handler on the
+    signal that triggered this handler. If this function is installed
+    as a signal handler for SIGINT, then pressing Ctrl-C once will
+    cause this program to finish operations in progress, then exit.
+    Pressing it again will cause an immediate exit; no cleanup
+    handlers will get called.
+    """
+    stderr.write("Attempting graceful exit. "
+                 "Press Ctrl-C again to exit immediately.\n")
+    main_thread = current_thread()
+    for thread in [t for t in threading_enumerate() if t is not main_thread]:
+        thread.abort = True
+    signal.signal(signum, immediate_exit)
+
+
+def immediate_exit(signum, frame):
+    os_exit(2)
+
+
 class QueueFunctionThread(Thread):
 
     def __init__(self, queue, func, *args, **kwargs):
@@ -1031,6 +1055,8 @@ Example:
             exit('no such command: %s' % args[0])
         exit()
 
+    signal.signal(signal.SIGINT, attempt_graceful_exit)
+
     print_queue = Queue(10000)
 
     def _print(item):