Skip to content

Commit e9a2225

Browse files
robertmhaasrobertmhaas
committed
Invent on_exit_nicely for pg_dump.
Per recent discussions on pgsql-hackers regarding parallel pg_dump.
1 parent 4bfe68d commit e9a2225

14 files changed

+207
-171
lines changed

‎src/bin/pg_dump/common.c

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -770,7 +770,7 @@ findParentsByOid(TableInfo *self,
770770
inhinfo[i].inhparent,
771771
self->dobj.name,
772772
oid);
773-
exit_nicely();
773+
exit_nicely(1);
774774
}
775775
self->parents[j++] = parent;
776776
}
@@ -809,7 +809,7 @@ parseOidArray(const char *str, Oid *array, int arraysize)
809809
if (argNum >= arraysize)
810810
{
811811
write_msg(NULL, "could not parse numeric array \"%s\": too many numbers\n", str);
812-
exit_nicely();
812+
exit_nicely(1);
813813
}
814814
temp[j] = '\0';
815815
array[argNum++] = atooid(temp);
@@ -824,7 +824,7 @@ parseOidArray(const char *str, Oid *array, int arraysize)
824824
j >= sizeof(temp) - 1)
825825
{
826826
write_msg(NULL, "could not parse numeric array \"%s\": invalid character in number\n", str);
827-
exit_nicely();
827+
exit_nicely(1);
828828
}
829829
temp[j++] = s;
830830
}

‎src/bin/pg_dump/compress_io.c

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454

5555
#include "compress_io.h"
5656
#include "dumpmem.h"
57+
#include "dumputils.h"
5758

5859
/*----------------------
5960
* Compressor API
@@ -109,8 +110,8 @@ ParseCompressionOption(int compression, CompressionAlgorithm *alg, int *level)
109110
*alg = COMPR_ALG_NONE;
110111
else
111112
{
112-
die_horribly(NULL, modulename, "Invalid compression code: %d\n",
113-
compression);
113+
exit_horribly(modulename, "Invalid compression code: %d\n",
114+
compression);
114115
*alg = COMPR_ALG_NONE; /* keep compiler quiet */
115116
}
116117

@@ -133,7 +134,7 @@ AllocateCompressor(int compression, WriteFunc writeF)
133134

134135
#ifndef HAVE_LIBZ
135136
if (alg == COMPR_ALG_LIBZ)
136-
die_horribly(NULL, modulename, "not built with zlib support\n");
137+
exit_horribly(modulename, "not built with zlib support\n");
137138
#endif
138139

139140
cs = (CompressorState *) pg_calloc(1, sizeof(CompressorState));
@@ -169,7 +170,7 @@ ReadDataFromArchive(ArchiveHandle *AH, int compression, ReadFunc readF)
169170
#ifdef HAVE_LIBZ
170171
ReadDataFromArchiveZlib(AH, readF);
171172
#else
172-
die_horribly(NULL, modulename, "not built with zlib support\n");
173+
exit_horribly(modulename, "not built with zlib support\n");
173174
#endif
174175
}
175176
}
@@ -187,7 +188,7 @@ WriteDataToArchive(ArchiveHandle *AH, CompressorState *cs,
187188
#ifdef HAVE_LIBZ
188189
return WriteDataToArchiveZlib(AH, cs, data, dLen);
189190
#else
190-
die_horribly(NULL, modulename, "not built with zlib support\n");
191+
exit_horribly(modulename, "not built with zlib support\n");
191192
#endif
192193
case COMPR_ALG_NONE:
193194
return WriteDataToArchiveNone(AH, cs, data, dLen);
@@ -234,9 +235,9 @@ InitCompressorZlib(CompressorState *cs, int level)
234235
cs->zlibOutSize = ZLIB_OUT_SIZE;
235236

236237
if (deflateInit(zp, level) != Z_OK)
237-
die_horribly(NULL, modulename,
238-
"could not initialize compression library: %s\n",
239-
zp->msg);
238+
exit_horribly(modulename,
239+
"could not initialize compression library: %s\n",
240+
zp->msg);
240241

241242
/* Just be paranoid - maybe End is called after Start, with no Write */
242243
zp->next_out = (void *) cs->zlibOut;
@@ -343,9 +344,9 @@ ReadDataFromArchiveZlib(ArchiveHandle *AH, ReadFunc readF)
343344
out = pg_malloc(ZLIB_OUT_SIZE + 1);
344345

345346
if (inflateInit(zp) != Z_OK)
346-
die_horribly(NULL, modulename,
347-
"could not initialize compression library: %s\n",
348-
zp->msg);
347+
exit_horribly(modulename,
348+
"could not initialize compression library: %s\n",
349+
zp->msg);
349350

350351
/* no minimal chunk size for zlib */
351352
while ((cnt = readF(AH, &buf, &buflen)))
@@ -514,7 +515,7 @@ cfopen_write(const char *path, const char *mode, int compression)
514515
fp = cfopen(fname, mode, 1);
515516
free(fname);
516517
#else
517-
die_horribly(NULL, modulename, "not built with zlib support\n");
518+
exit_horribly(modulename, "not built with zlib support\n");
518519
fp = NULL; /* keep compiler quiet */
519520
#endif
520521
}
@@ -541,7 +542,7 @@ cfopen(const char *path, const char *mode, int compression)
541542
fp = NULL;
542543
}
543544
#else
544-
die_horribly(NULL, modulename, "not built with zlib support\n");
545+
exit_horribly(modulename, "not built with zlib support\n");
545546
#endif
546547
}
547548
else

‎src/bin/pg_dump/dumputils.c

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,15 @@
2626
int quote_all_identifiers = 0;
2727
const char *progname = NULL;
2828

29+
#define MAX_ON_EXIT_NICELY 20
30+
31+
static struct
32+
{
33+
on_exit_nicely_callback function;
34+
void *arg;
35+
} on_exit_nicely_list[MAX_ON_EXIT_NICELY];
36+
37+
static int on_exit_nicely_index;
2938

3039
#define supports_grant_options(version) ((version) >= 70400)
3140

@@ -1261,7 +1270,7 @@ exit_horribly(const char *modulename, const char *fmt,...)
12611270
vwrite_msg(modulename, fmt, ap);
12621271
va_end(ap);
12631272

1264-
exit(1);
1273+
exit_nicely(1);
12651274
}
12661275

12671276
/*
@@ -1289,6 +1298,27 @@ set_section (const char *arg, int *dumpSections)
12891298
progname, arg);
12901299
fprintf(stderr, _("Try \"%s --help\" for more information.\n"),
12911300
progname);
1292-
exit(1);
1301+
exit_nicely(1);
12931302
}
12941303
}
1304+
1305+
/* Register a callback to be run when exit_nicely is invoked. */
1306+
void
1307+
on_exit_nicely(on_exit_nicely_callback function, void *arg)
1308+
{
1309+
if (on_exit_nicely_index >= MAX_ON_EXIT_NICELY)
1310+
exit_horribly(NULL, "out of on_exit_nicely slots");
1311+
on_exit_nicely_list[on_exit_nicely_index].function = function;
1312+
on_exit_nicely_list[on_exit_nicely_index].arg = arg;
1313+
on_exit_nicely_index++;
1314+
}
1315+
1316+
/* Run accumulated on_exit_nicely callbacks and then exit quietly. */
1317+
void
1318+
exit_nicely(int code)
1319+
{
1320+
while (--on_exit_nicely_index >= 0)
1321+
(*on_exit_nicely_list[on_exit_nicely_index].function)(code,
1322+
on_exit_nicely_list[on_exit_nicely_index].arg);
1323+
exit(code);
1324+
}

‎src/bin/pg_dump/dumputils.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,4 +60,8 @@ extern void exit_horribly(const char *modulename, const char *fmt,...)
6060
__attribute__((format(PG_PRINTF_ATTRIBUTE, 2, 3), noreturn));
6161
extern void set_section (const char *arg, int *dumpSections);
6262

63+
typedef void (*on_exit_nicely_callback) (int code, void *arg);
64+
extern void on_exit_nicely(on_exit_nicely_callback function, void *arg);
65+
extern void exit_nicely(int code) __attribute__((noreturn));
66+
6367
#endif /* DUMPUTILS_H */

‎src/bin/pg_dump/pg_backup.h

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -159,15 +159,13 @@ typedef struct _restoreOptions
159159
* Main archiver interface.
160160
*/
161161

162-
163-
/* Lets the archive know we have a DB connection to shutdown if it dies */
164-
165-
PGconn *ConnectDatabase(Archive *AH,
162+
extern PGconn *ConnectDatabase(Archive *AH,
166163
const char *dbname,
167164
const char *pghost,
168165
const char *pgport,
169166
const char *username,
170167
enum trivalue prompt_password);
168+
extern void DisconnectDatabase(Archive *AHX);
171169

172170
/* Called to add a TOC entry */
173171
extern void ArchiveEntry(Archive *AHX,

‎src/bin/pg_dump/pg_backup_archiver.c

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -459,10 +459,7 @@ RestoreArchive(Archive *AHX, RestoreOptions *ropt)
459459
RestoreOutput(AH, sav);
460460

461461
if (ropt->useDB)
462-
{
463-
PQfinish(AH->connection);
464-
AH->connection = NULL;
465-
}
462+
DisconnectDatabase(&AH->public);
466463
}
467464

468465
/*
@@ -1435,11 +1432,10 @@ vdie_horribly(ArchiveHandle *AH, const char *modulename,
14351432
{
14361433
if (AH->public.verbose)
14371434
write_msg(NULL, "*** aborted because of error\n");
1438-
if (AH->connection)
1439-
PQfinish(AH->connection);
1435+
DisconnectDatabase(&AH->public);
14401436
}
14411437

1442-
exit(1);
1438+
exit_nicely(1);
14431439
}
14441440

14451441
/* As above, but with variable arg list */
@@ -3332,8 +3328,7 @@ restore_toc_entries_parallel(ArchiveHandle *AH)
33323328
* mainly to ensure that we don't exceed the specified number of parallel
33333329
* connections.
33343330
*/
3335-
PQfinish(AH->connection);
3336-
AH->connection = NULL;
3331+
DisconnectDatabase(&AH->public);
33373332

33383333
/* blow away any transient state from the old connection */
33393334
if (AH->currUser)
@@ -3795,8 +3790,7 @@ parallel_restore(RestoreArgs *args)
37953790
retval = restore_toc_entry(AH, te, ropt, true);
37963791

37973792
/* And clean up */
3798-
PQfinish(AH->connection);
3799-
AH->connection = NULL;
3793+
DisconnectDatabase((Archive *) AH);
38003794

38013795
/* If we reopened the file, we are done with it, so close it now */
38023796
if (te->section == SECTION_DATA)

‎src/bin/pg_dump/pg_backup_db.c

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -310,6 +310,15 @@ ConnectDatabase(Archive *AHX,
310310
return AH->connection;
311311
}
312312

313+
void
314+
DisconnectDatabase(Archive *AHX)
315+
{
316+
ArchiveHandle *AH = (ArchiveHandle *) AHX;
317+
318+
PQfinish(AH->connection); /* noop if AH->connection is NULL */
319+
AH->connection = NULL;
320+
}
321+
313322

314323
static void
315324
notice_processor(void *arg, const char *message)

‎src/bin/pg_dump/pg_backup_directory.c

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
#include "compress_io.h"
3737
#include "dumpmem.h"
38+
#include "dumputils.h"
3839

3940
#include <dirent.h>
4041
#include <sys/stat.h>
@@ -633,13 +634,13 @@ createDirectory(const char *dir)
633634
if (stat(dir, &st) == 0)
634635
{
635636
if (S_ISDIR(st.st_mode))
636-
die_horribly(NULL, modulename,
637-
"cannot create directory %s, it exists already\n",
638-
dir);
637+
exit_horribly(modulename,
638+
"cannot create directory %s, it exists already\n",
639+
dir);
639640
else
640-
die_horribly(NULL, modulename,
641-
"cannot create directory %s, a file with this name "
642-
"exists already\n", dir);
641+
exit_horribly(modulename,
642+
"cannot create directory %s, a file with this name "
643+
"exists already\n", dir);
643644
}
644645

645646
/*
@@ -648,8 +649,8 @@ createDirectory(const char *dir)
648649
* between our two calls.
649650
*/
650651
if (mkdir(dir, 0700) < 0)
651-
die_horribly(NULL, modulename, "could not create directory %s: %s",
652-
dir, strerror(errno));
652+
exit_horribly(modulename, "could not create directory %s: %s",
653+
dir, strerror(errno));
653654
}
654655

655656

‎src/bin/pg_dump/pg_backup_files.c

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -127,15 +127,15 @@ InitArchiveFmt_Files(ArchiveHandle *AH)
127127
{
128128
AH->FH = fopen(AH->fSpec, PG_BINARY_W);
129129
if (AH->FH == NULL)
130-
die_horribly(NULL, modulename, "could not open output file \"%s\": %s\n",
131-
AH->fSpec, strerror(errno));
130+
exit_horribly(modulename, "could not open output file \"%s\": %s\n",
131+
AH->fSpec, strerror(errno));
132132
}
133133
else
134134
{
135135
AH->FH = stdout;
136136
if (AH->FH == NULL)
137-
die_horribly(NULL, modulename, "could not open output file: %s\n",
138-
strerror(errno));
137+
exit_horribly(modulename, "could not open output file: %s\n",
138+
strerror(errno));
139139
}
140140

141141
ctx->hasSeek = checkSeek(AH->FH);
@@ -152,15 +152,15 @@ InitArchiveFmt_Files(ArchiveHandle *AH)
152152
{
153153
AH->FH = fopen(AH->fSpec, PG_BINARY_R);
154154
if (AH->FH == NULL)
155-
die_horribly(NULL, modulename, "could not open input file \"%s\": %s\n",
156-
AH->fSpec, strerror(errno));
155+
exit_horribly(modulename, "could not open input file \"%s\": %s\n",
156+
AH->fSpec, strerror(errno));
157157
}
158158
else
159159
{
160160
AH->FH = stdin;
161161
if (AH->FH == NULL)
162-
die_horribly(NULL, modulename, "could not open input file: %s\n",
163-
strerror(errno));
162+
exit_horribly(modulename, "could not open input file: %s\n",
163+
strerror(errno));
164164
}
165165

166166
ctx->hasSeek = checkSeek(AH->FH);

0 commit comments

Comments
 (0)