notmuch new/tag: Flush all changes to database when interrupted.

By installing a signal handler for SIGINT we can ensure that no work
that is already complete will be lost if the user interrupts a
"notmuch new" run with Control-C.

2 files changed, 52 insertions, 18 deletions

diff --git a/notmuch-new.c b/notmuch-new.c
index 6247088a..83a05ba6 100644
--- a/notmuch-new.c
+++ b/notmuch-new.c
@@ -28,6 +28,16 @@ handle_sigalrm (unused (int signal))
     do_add_files_print_progress = 1;
 }
 
+static volatile sig_atomic_t interrupted;
+
+static void
+handle_sigint (unused (int sig))
+{
+    static char msg[] = "Stopping...         \n";
+    write(2, msg, sizeof(msg)-1);
+    interrupted = 1;
+}
+
 static void
 tag_inbox_and_unread (notmuch_message_t *message)
 {
@@ -125,7 +135,7 @@ add_files_recursive (notmuch_database_t *notmuch,
 	pathconf (path, _PC_NAME_MAX) + 1;
     entry = malloc (entry_length);
 
-    while (1) {
+    while (!interrupted) {
 	err = readdir_r (dir, entry, &e);
 	if (err) {
 	    fprintf (stderr, "Error reading directory: %s\n",
@@ -319,7 +329,7 @@ count_files (const char *path, int *count)
 	pathconf (path, _PC_NAME_MAX) + 1;
     entry = malloc (entry_length);
 
-    while (1) {
+    while (!interrupted) {
 	err = readdir_r (dir, entry, &e);
 	if (err) {
 	    fprintf (stderr, "Error reading directory: %s\n",
@@ -385,7 +395,14 @@ notmuch_new_command (void *ctx,
     struct stat st;
     const char *db_path;
     char *dot_notmuch_path;
-    int new_database = 0;
+    struct sigaction action;
+
+    /* Setup our handler for SIGINT */
+    memset (&action, 0, sizeof (struct sigaction));
+    action.sa_handler = handle_sigint;
+    sigemptyset (&action.sa_mask);
+    action.sa_flags = SA_RESTART;
+    sigaction (SIGINT, &action, NULL);
 
     config = notmuch_config_open (ctx, NULL, NULL);
     if (config == NULL)
@@ -396,10 +413,20 @@ notmuch_new_command (void *ctx,
     dot_notmuch_path = talloc_asprintf (ctx, "%s/%s", db_path, ".notmuch");
 
     if (stat (dot_notmuch_path, &st)) {
-	new_database = 1;
+	int count;
+
+	count = 0;
+	count_files (db_path, &count);
+	if (interrupted)
+	    return 1;
+
 	notmuch = notmuch_database_create (db_path);
+	add_files_state.ignore_read_only_directories = FALSE;
+	add_files_state.total_files = count;
     } else {
 	notmuch = notmuch_database_open (db_path);
+	add_files_state.ignore_read_only_directories = TRUE;
+	add_files_state.total_files = 0;
     }
 
     if (notmuch == NULL)
@@ -408,17 +435,6 @@ notmuch_new_command (void *ctx,
     talloc_free (dot_notmuch_path);
     dot_notmuch_path = NULL;
 
-    if (new_database) {
-	int count;
-	count = 0;
-	count_files (db_path, &count);
-	add_files_state.ignore_read_only_directories = FALSE;
-	add_files_state.total_files = count;
-    } else {
-	add_files_state.ignore_read_only_directories = TRUE;
-	add_files_state.total_files = 0;
-    }
-
     add_files_state.saw_read_only_directory = FALSE;
     add_files_state.total_files = 0;
     add_files_state.processed_files = 0;
@@ -465,5 +481,5 @@ notmuch_new_command (void *ctx,
 
     notmuch_database_close (notmuch);
 
-    return ret;
+    return ret || interrupted;
 }
diff --git a/notmuch-tag.c b/notmuch-tag.c
index 80582acf..7d92ec48 100644
--- a/notmuch-tag.c
+++ b/notmuch-tag.c
@@ -20,6 +20,16 @@
 
 #include "notmuch-client.h"
 
+static volatile sig_atomic_t interrupted;
+
+static void
+handle_sigint (unused (int sig))
+{
+    static char msg[] = "Stopping...         \n";
+    write(2, msg, sizeof(msg)-1);
+    interrupted = 1;
+}
+
 int
 notmuch_tag_command (void *ctx, unused (int argc), unused (char *argv[]))
 {
@@ -32,8 +42,16 @@ notmuch_tag_command (void *ctx, unused (int argc), unused (char *argv[]))
     notmuch_query_t *query;
     notmuch_messages_t *messages;
     notmuch_message_t *message;
+    struct sigaction action;
     int i;
 
+    /* Setup our handler for SIGINT */
+    memset (&action, 0, sizeof (struct sigaction));
+    action.sa_handler = handle_sigint;
+    sigemptyset (&action.sa_mask);
+    action.sa_flags = SA_RESTART;
+    sigaction (SIGINT, &action, NULL);
+
     add_tags = talloc_size (ctx, argc * sizeof (int));
     if (add_tags == NULL) {
 	fprintf (stderr, "Out of memory.\n");
@@ -87,7 +105,7 @@ notmuch_tag_command (void *ctx, unused (int argc), unused (char *argv[]))
     }
 
     for (messages = notmuch_query_search_messages (query, 0, -1);
-	 notmuch_messages_has_more (messages);
+	 notmuch_messages_has_more (messages) && !interrupted;
 	 notmuch_messages_advance (messages))
     {
 	message = notmuch_messages_get (messages);
@@ -109,5 +127,5 @@ notmuch_tag_command (void *ctx, unused (int argc), unused (char *argv[]))
     notmuch_query_destroy (query);
     notmuch_database_close (notmuch);
 
-    return 0;
+    return interrupted;
 }