Invent on_exit_nicely for pg_dump.

Per recent discussions on pgsql-hackers regarding parallel pg_dump.

src/bin/pg_dump/common.c

@@ -770,7 +770,7 @@ findParentsByOid(TableInfo *self,
 							  inhinfo[i].inhparent,
 							  self->dobj.name,
 							  oid);
-					exit_nicely();
+					exit_nicely(1);
 				}
 				self->parents[j++] = parent;
 			}
@@ -809,7 +809,7 @@ parseOidArray(const char *str, Oid *array, int arraysize)
 				if (argNum >= arraysize)
 				{
 					write_msg(NULL, "could not parse numeric array \"%s\": too many numbers\n", str);
-					exit_nicely();
+					exit_nicely(1);
 				}
 				temp[j] = '\0';
 				array[argNum++] = atooid(temp);
@@ -824,7 +824,7 @@ parseOidArray(const char *str, Oid *array, int arraysize)
 				j >= sizeof(temp) - 1)
 			{
 				write_msg(NULL, "could not parse numeric array \"%s\": invalid character in number\n", str);
-				exit_nicely();
+				exit_nicely(1);
 			}
 			temp[j++] = s;
 		}

src/bin/pg_dump/compress_io.c

@@ -54,6 +54,7 @@
 
 #include "compress_io.h"
 #include "dumpmem.h"
+#include "dumputils.h"
 
 /*----------------------
  * Compressor API
@@ -109,7 +110,7 @@ ParseCompressionOption(int compression, CompressionAlgorithm *alg, int *level)
 		*alg = COMPR_ALG_NONE;
 	else
 	{
-		die_horribly(NULL, modulename, "Invalid compression code: %d\n",
+		exit_horribly(modulename, "Invalid compression code: %d\n",
 					  compression);
 		*alg = COMPR_ALG_NONE;	/* keep compiler quiet */
 	}
@@ -133,7 +134,7 @@ AllocateCompressor(int compression, WriteFunc writeF)
 
 #ifndef HAVE_LIBZ
 	if (alg == COMPR_ALG_LIBZ)
-		die_horribly(NULL, modulename, "not built with zlib support\n");
+		exit_horribly(modulename, "not built with zlib support\n");
 #endif
 
 	cs = (CompressorState *) pg_calloc(1, sizeof(CompressorState));
@@ -169,7 +170,7 @@ ReadDataFromArchive(ArchiveHandle *AH, int compression, ReadFunc readF)
 #ifdef HAVE_LIBZ
 		ReadDataFromArchiveZlib(AH, readF);
 #else
-		die_horribly(NULL, modulename, "not built with zlib support\n");
+		exit_horribly(modulename, "not built with zlib support\n");
 #endif
 	}
 }
@@ -187,7 +188,7 @@ WriteDataToArchive(ArchiveHandle *AH, CompressorState *cs,
 #ifdef HAVE_LIBZ
 			return WriteDataToArchiveZlib(AH, cs, data, dLen);
 #else
-			die_horribly(NULL, modulename, "not built with zlib support\n");
+			exit_horribly(modulename, "not built with zlib support\n");
 #endif
 		case COMPR_ALG_NONE:
 			return WriteDataToArchiveNone(AH, cs, data, dLen);
@@ -234,7 +235,7 @@ InitCompressorZlib(CompressorState *cs, int level)
 	cs->zlibOutSize = ZLIB_OUT_SIZE;
 
 	if (deflateInit(zp, level) != Z_OK)
-		die_horribly(NULL, modulename,
+		exit_horribly(modulename,
 					  "could not initialize compression library: %s\n",
 					  zp->msg);
 
@@ -343,7 +344,7 @@ ReadDataFromArchiveZlib(ArchiveHandle *AH, ReadFunc readF)
 	out = pg_malloc(ZLIB_OUT_SIZE + 1);
 
 	if (inflateInit(zp) != Z_OK)
-		die_horribly(NULL, modulename,
+		exit_horribly(modulename,
 					  "could not initialize compression library: %s\n",
 					  zp->msg);
 
@@ -514,7 +515,7 @@ cfopen_write(const char *path, const char *mode, int compression)
 		fp = cfopen(fname, mode, 1);
 		free(fname);
 #else
-		die_horribly(NULL, modulename, "not built with zlib support\n");
+		exit_horribly(modulename, "not built with zlib support\n");
 		fp = NULL;				/* keep compiler quiet */
 #endif
 	}
@@ -541,7 +542,7 @@ cfopen(const char *path, const char *mode, int compression)
 			fp = NULL;
 		}
 #else
-		die_horribly(NULL, modulename, "not built with zlib support\n");
+		exit_horribly(modulename, "not built with zlib support\n");
 #endif
 	}
 	else

src/bin/pg_dump/dumputils.c

@@ -26,6 +26,15 @@
 int			quote_all_identifiers = 0;
 const char *progname = NULL;
 
+#define MAX_ON_EXIT_NICELY				20
+
+static struct
+{
+	on_exit_nicely_callback	function;
+	void	   *arg;
+} on_exit_nicely_list[MAX_ON_EXIT_NICELY];
+
+static int on_exit_nicely_index;
 
 #define supports_grant_options(version) ((version) >= 70400)
 
@@ -1261,7 +1270,7 @@ exit_horribly(const char *modulename, const char *fmt,...)
 	vwrite_msg(modulename, fmt, ap);
 	va_end(ap);
 
-	exit(1);
+	exit_nicely(1);
 }
 
 /*
@@ -1289,6 +1298,27 @@ set_section (const char *arg, int *dumpSections)
 				progname, arg);
 		fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
 				progname);
-		exit(1);
+		exit_nicely(1);
 	}
 }
+
+/* Register a callback to be run when exit_nicely is invoked. */
+void
+on_exit_nicely(on_exit_nicely_callback function, void *arg)
+{
+	if (on_exit_nicely_index >= MAX_ON_EXIT_NICELY)
+		exit_horribly(NULL, "out of on_exit_nicely slots");
+	on_exit_nicely_list[on_exit_nicely_index].function = function;
+	on_exit_nicely_list[on_exit_nicely_index].arg = arg;
+	on_exit_nicely_index++;
+}
+
+/* Run accumulated on_exit_nicely callbacks and then exit quietly. */
+void
+exit_nicely(int code)
+{
+	while (--on_exit_nicely_index >= 0)
+		(*on_exit_nicely_list[on_exit_nicely_index].function)(code,
+			on_exit_nicely_list[on_exit_nicely_index].arg);
+	exit(code);
+}

src/bin/pg_dump/dumputils.h

@@ -60,4 +60,8 @@ extern void exit_horribly(const char *modulename, const char *fmt,...)
 				__attribute__((format(PG_PRINTF_ATTRIBUTE, 2, 3), noreturn));
 extern void set_section (const char *arg, int *dumpSections);
 
+typedef void (*on_exit_nicely_callback) (int code, void *arg);
+extern void on_exit_nicely(on_exit_nicely_callback function, void *arg);
+extern void exit_nicely(int code) __attribute__((noreturn));
+
 #endif   /* DUMPUTILS_H */

src/bin/pg_dump/pg_backup.h

@@ -159,15 +159,13 @@ typedef struct _restoreOptions
  * Main archiver interface.
  */
 
-
-/* Lets the archive know we have a DB connection to shutdown if it dies */
-
-PGconn *ConnectDatabase(Archive *AH,
+extern PGconn *ConnectDatabase(Archive *AH,
 				const char *dbname,
 				const char *pghost,
 				const char *pgport,
 				const char *username,
 				enum trivalue prompt_password);
+extern void DisconnectDatabase(Archive *AHX);
 
 /* Called to add a TOC entry */
 extern void ArchiveEntry(Archive *AHX,

src/bin/pg_dump/pg_backup_archiver.c

@@ -459,10 +459,7 @@ RestoreArchive(Archive *AHX, RestoreOptions *ropt)
 		RestoreOutput(AH, sav);
 
 	if (ropt->useDB)
-	{
-		PQfinish(AH->connection);
-		AH->connection = NULL;
-	}
+		DisconnectDatabase(&AH->public);
 }
 
 /*
@@ -1435,11 +1432,10 @@ vdie_horribly(ArchiveHandle *AH, const char *modulename,
 	{
 		if (AH->public.verbose)
 			write_msg(NULL, "*** aborted because of error\n");
-		if (AH->connection)
-			PQfinish(AH->connection);
+		DisconnectDatabase(&AH->public);
 	}
 
-	exit(1);
+	exit_nicely(1);
 }
 
 /* As above, but with variable arg list */
@@ -3332,8 +3328,7 @@ restore_toc_entries_parallel(ArchiveHandle *AH)
 	 * mainly to ensure that we don't exceed the specified number of parallel
 	 * connections.
 	 */
-	PQfinish(AH->connection);
-	AH->connection = NULL;
+	DisconnectDatabase(&AH->public);
 
 	/* blow away any transient state from the old connection */
 	if (AH->currUser)
@@ -3795,8 +3790,7 @@ parallel_restore(RestoreArgs *args)
 	retval = restore_toc_entry(AH, te, ropt, true);
 
 	/* And clean up */
-	PQfinish(AH->connection);
-	AH->connection = NULL;
+	DisconnectDatabase((Archive *) AH);
 
 	/* If we reopened the file, we are done with it, so close it now */
 	if (te->section == SECTION_DATA)

src/bin/pg_dump/pg_backup_db.c

@@ -310,6 +310,15 @@ ConnectDatabase(Archive *AHX,
 	return AH->connection;
 }
 
+void
+DisconnectDatabase(Archive *AHX)
+{
+	ArchiveHandle *AH = (ArchiveHandle *) AHX;
+
+	PQfinish(AH->connection);		/* noop if AH->connection is NULL */
+	AH->connection = NULL;
+}
+
 
 static void
 notice_processor(void *arg, const char *message)

src/bin/pg_dump/pg_backup_directory.c

@@ -35,6 +35,7 @@
 
 #include "compress_io.h"
 #include "dumpmem.h"
+#include "dumputils.h"
 
 #include <dirent.h>
 #include <sys/stat.h>
@@ -633,11 +634,11 @@ createDirectory(const char *dir)
 	if (stat(dir, &st) == 0)
 	{
 		if (S_ISDIR(st.st_mode))
-			die_horribly(NULL, modulename,
+			exit_horribly(modulename,
 						  "cannot create directory %s, it exists already\n",
 						  dir);
 		else
-			die_horribly(NULL, modulename,
+			exit_horribly(modulename,
 						  "cannot create directory %s, a file with this name "
 						  "exists already\n", dir);
 	}
@@ -648,7 +649,7 @@ createDirectory(const char *dir)
 	 * between our two calls.
 	 */
 	if (mkdir(dir, 0700) < 0)
-		die_horribly(NULL, modulename, "could not create directory %s: %s",
+		exit_horribly(modulename, "could not create directory %s: %s",
 					  dir, strerror(errno));
 }
 

src/bin/pg_dump/pg_backup_files.c

@@ -127,14 +127,14 @@ InitArchiveFmt_Files(ArchiveHandle *AH)
 		{
 			AH->FH = fopen(AH->fSpec, PG_BINARY_W);
 			if (AH->FH == NULL)
-				die_horribly(NULL, modulename, "could not open output file \"%s\": %s\n",
+				exit_horribly(modulename, "could not open output file \"%s\": %s\n",
 							  AH->fSpec, strerror(errno));
 		}
 		else
 		{
 			AH->FH = stdout;
 			if (AH->FH == NULL)
-				die_horribly(NULL, modulename, "could not open output file: %s\n",
+				exit_horribly(modulename, "could not open output file: %s\n",
 							  strerror(errno));
 		}
 
@@ -152,14 +152,14 @@ InitArchiveFmt_Files(ArchiveHandle *AH)
 		{
 			AH->FH = fopen(AH->fSpec, PG_BINARY_R);
 			if (AH->FH == NULL)
-				die_horribly(NULL, modulename, "could not open input file \"%s\": %s\n",
+				exit_horribly(modulename, "could not open input file \"%s\": %s\n",
 							  AH->fSpec, strerror(errno));
 		}
 		else
 		{
 			AH->FH = stdin;
 			if (AH->FH == NULL)
-				die_horribly(NULL, modulename, "could not open input file: %s\n",
+				exit_horribly(modulename, "could not open input file: %s\n",
 							  strerror(errno));
 		}
 

src/bin/pg_dump/pg_backup_tar.c

@@ -29,6 +29,7 @@
 #include "pg_backup_archiver.h"
 #include "pg_backup_tar.h"
 #include "dumpmem.h"
+#include "dumputils.h"
 
 #include <sys/stat.h>
 #include <ctype.h>
@@ -178,7 +179,7 @@ InitArchiveFmt_Tar(ArchiveHandle *AH)
 		{
 			ctx->tarFH = fopen(AH->fSpec, PG_BINARY_W);
 			if (ctx->tarFH == NULL)
-				die_horribly(NULL, modulename,
+				exit_horribly(modulename,
 							  "could not open TOC file \"%s\" for output: %s\n",
 							  AH->fSpec, strerror(errno));
 		}
@@ -186,7 +187,7 @@ InitArchiveFmt_Tar(ArchiveHandle *AH)
 		{
 			ctx->tarFH = stdout;
 			if (ctx->tarFH == NULL)
-				die_horribly(NULL, modulename,
+				exit_horribly(modulename,
 							  "could not open TOC file for output: %s\n",
 							  strerror(errno));
 		}
@@ -214,7 +215,8 @@ InitArchiveFmt_Tar(ArchiveHandle *AH)
 		 * positioning.
 		 */
 		if (AH->compression != 0)
-			die_horribly(NULL, modulename, "compression is not supported by tar archive format\n");
+			exit_horribly(modulename,
+						  "compression is not supported by tar archive format\n");
 	}
 	else
 	{							/* Read Mode */
@@ -222,14 +224,14 @@ InitArchiveFmt_Tar(ArchiveHandle *AH)
 		{
 			ctx->tarFH = fopen(AH->fSpec, PG_BINARY_R);
 			if (ctx->tarFH == NULL)
-				die_horribly(NULL, modulename, "could not open TOC file \"%s\" for input: %s\n",
+				exit_horribly(modulename, "could not open TOC file \"%s\" for input: %s\n",
 							  AH->fSpec, strerror(errno));
 		}
 		else
 		{
 			ctx->tarFH = stdin;
 			if (ctx->tarFH == NULL)
-				die_horribly(NULL, modulename, "could not open TOC file for input: %s\n",
+				exit_horribly(modulename, "could not open TOC file for input: %s\n",
 							  strerror(errno));
 		}
 

src/bin/pg_dump/pg_dump.h

@@ -518,8 +518,6 @@ extern void simple_string_list_append(SimpleStringList *list, const char *val);
 extern bool simple_oid_list_member(SimpleOidList *list, Oid val);
 extern bool simple_string_list_member(SimpleStringList *list, const char *val);
 
-extern void exit_nicely(void) __attribute__((noreturn));
-
 extern void parseOidArray(const char *str, Oid *array, int arraysize);
 
 extern void sortDumpableObjects(DumpableObject **objs, int numObjs);

src/bin/pg_dump/pg_dumpall.c

@@ -151,12 +151,12 @@ main(int argc, char *argv[])
 		if (strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-?") == 0)
 		{
 			help();
-			exit(0);
+			exit_nicely(0);
 		}
 		if (strcmp(argv[1], "--version") == 0 || strcmp(argv[1], "-V") == 0)
 		{
 			puts("pg_dumpall (PostgreSQL) " PG_VERSION);
-			exit(0);
+			exit_nicely(0);
 		}
 	}
 
@@ -181,7 +181,7 @@ main(int argc, char *argv[])
 					  "but was not the same version as %s.\n"
 					  "Check your installation.\n"),
 					full_path, progname);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	pgdumpopts = createPQExpBuffer();
@@ -296,7 +296,7 @@ main(int argc, char *argv[])
 
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
-				exit(1);
+				exit_nicely(1);
 		}
 	}
 
@@ -307,7 +307,7 @@ main(int argc, char *argv[])
 				progname, argv[optind]);
 		fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
 				progname);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	/* Make sure the user hasn't specified a mix of globals-only options */
@@ -317,7 +317,7 @@ main(int argc, char *argv[])
 				progname);
 		fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
 				progname);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	if (globals_only && tablespaces_only)
@@ -326,7 +326,7 @@ main(int argc, char *argv[])
 				progname);
 		fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
 				progname);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	if (roles_only && tablespaces_only)
@@ -335,7 +335,7 @@ main(int argc, char *argv[])
 				progname);
 		fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
 				progname);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	/* Add long options to the pg_dump argument list */
@@ -375,7 +375,7 @@ main(int argc, char *argv[])
 		{
 			fprintf(stderr, _("%s: could not connect to database \"%s\"\n"),
 					progname, pgdb);
-			exit(1);
+			exit_nicely(1);
 		}
 	}
 	else
@@ -393,7 +393,7 @@ main(int argc, char *argv[])
 					progname);
 			fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
 					progname);
-			exit(1);
+			exit_nicely(1);
 		}
 	}
 
@@ -407,7 +407,7 @@ main(int argc, char *argv[])
 		{
 			fprintf(stderr, _("%s: could not open the output file \"%s\": %s\n"),
 					progname, filename, strerror(errno));
-			exit(1);
+			exit_nicely(1);
 		}
 	}
 	else
@@ -525,7 +525,7 @@ main(int argc, char *argv[])
 	if (filename)
 		fclose(OPF);
 
-	exit(0);
+	exit_nicely(0);
 }
 
 
@@ -1068,7 +1068,7 @@ dumpTablespaces(PGconn *conn)
 			fprintf(stderr, _("%s: could not parse ACL list (%s) for tablespace \"%s\"\n"),
 					progname, spcacl, fspcname);
 			PQfinish(conn);
-			exit(1);
+			exit_nicely(1);
 		}
 
 		if (spccomment && strlen(spccomment))
@@ -1372,7 +1372,7 @@ dumpCreateDB(PGconn *conn)
 			fprintf(stderr, _("%s: could not parse ACL list (%s) for database \"%s\"\n"),
 					progname, dbacl, fdbname);
 			PQfinish(conn);
-			exit(1);
+			exit_nicely(1);
 		}
 
 		fprintf(OPF, "%s", buf->data);
@@ -1587,7 +1587,7 @@ dumpDatabases(PGconn *conn)
 		if (ret != 0)
 		{
 			fprintf(stderr, _("%s: pg_dump failed on database \"%s\", exiting\n"), progname, dbname);
-			exit(1);
+			exit_nicely(1);
 		}
 
 		if (filename)
@@ -1597,7 +1597,7 @@ dumpDatabases(PGconn *conn)
 			{
 				fprintf(stderr, _("%s: could not re-open the output file \"%s\": %s\n"),
 						progname, filename, strerror(errno));
-				exit(1);
+				exit_nicely(1);
 			}
 		}
 
@@ -1724,7 +1724,7 @@ connectDatabase(const char *dbname, const char *pghost, const char *pgport,
 		{
 			fprintf(stderr, _("%s: could not connect to database \"%s\"\n"),
 					progname, dbname);
-			exit(1);
+			exit_nicely(1);
 		}
 
 		if (PQstatus(conn) == CONNECTION_BAD &&
@@ -1746,7 +1746,7 @@ connectDatabase(const char *dbname, const char *pghost, const char *pgport,
 			fprintf(stderr,
 					_("%s: could not connect to database \"%s\": %s\n"),
 					progname, dbname, PQerrorMessage(conn));
-			exit(1);
+			exit_nicely(1);
 		}
 		else
 		{
@@ -1759,14 +1759,14 @@ connectDatabase(const char *dbname, const char *pghost, const char *pgport,
 	if (!remoteversion_str)
 	{
 		fprintf(stderr, _("%s: could not get server version\n"), progname);
-		exit(1);
+		exit_nicely(1);
 	}
 	server_version = parse_version(remoteversion_str);
 	if (server_version < 0)
 	{
 		fprintf(stderr, _("%s: could not parse server version \"%s\"\n"),
 				progname, remoteversion_str);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	my_version = parse_version(PG_VERSION);
@@ -1774,7 +1774,7 @@ connectDatabase(const char *dbname, const char *pghost, const char *pgport,
 	{
 		fprintf(stderr, _("%s: could not parse version \"%s\"\n"),
 				progname, PG_VERSION);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	/*
@@ -1788,7 +1788,7 @@ connectDatabase(const char *dbname, const char *pghost, const char *pgport,
 		fprintf(stderr, _("server version: %s; %s version: %s\n"),
 				remoteversion_str, progname, PG_VERSION);
 		fprintf(stderr, _("aborting because of server version mismatch\n"));
-		exit(1);
+		exit_nicely(1);
 	}
 
 	/*
@@ -1822,7 +1822,7 @@ executeQuery(PGconn *conn, const char *query)
 		fprintf(stderr, _("%s: query was: %s\n"),
 				progname, query);
 		PQfinish(conn);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	return res;
@@ -1848,7 +1848,7 @@ executeCommand(PGconn *conn, const char *query)
 		fprintf(stderr, _("%s: query was: %s\n"),
 				progname, query);
 		PQfinish(conn);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	PQclear(res);

src/bin/pg_dump/pg_restore.c

@@ -138,12 +138,12 @@ main(int argc, char **argv)
 		if (strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-?") == 0)
 		{
 			usage(progname);
-			exit(0);
+			exit_nicely(0);
 		}
 		if (strcmp(argv[1], "--version") == 0 || strcmp(argv[1], "-V") == 0)
 		{
 			puts("pg_restore (PostgreSQL) " PG_VERSION);
-			exit(0);
+			exit_nicely(0);
 		}
 	}
 
@@ -279,7 +279,7 @@ main(int argc, char **argv)
 
 			default:
 				fprintf(stderr, _("Try \"%s --help\" for more information.\n"), progname);
-				exit(1);
+				exit_nicely(1);
 		}
 	}
 
@@ -296,21 +296,21 @@ main(int argc, char **argv)
 				progname, argv[optind]);
 		fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
 				progname);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	if (opts->dataOnly && opts->schemaOnly)
 	{
 		fprintf(stderr, _("%s: options -s/--schema-only and -a/--data-only cannot be used together\n"),
 			progname);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	if ((opts->dataOnly || opts->schemaOnly) && (opts->dumpSections != DUMP_UNSECTIONED))
 	{
 		fprintf(stderr, _("%s: options -s/--schema-only and -a/--data-only cannot be used with --section\n"),
 			progname);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	if (opts->dataOnly)
@@ -332,7 +332,7 @@ main(int argc, char **argv)
 					progname);
 			fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
 					progname);
-			exit(1);
+			exit_nicely(1);
 		}
 		opts->useDB = 1;
 	}
@@ -342,7 +342,7 @@ main(int argc, char **argv)
 	{
 		fprintf(stderr, _("%s: cannot specify both --single-transaction and multiple jobs\n"),
 				progname);
-		exit(1);
+		exit_nicely(1);
 	}
 
 	opts->disable_triggers = disable_triggers;
@@ -378,7 +378,7 @@ main(int argc, char **argv)
 			default:
 				write_msg(NULL, "unrecognized archive format \"%s\"; please specify \"c\", \"d\", or \"t\"\n",
 						  opts->formatName);
-				exit(1);
+				exit_nicely(1);
 		}
 	}