This adds a simple struct and open and close functions to abstract and isolate the zlib vs. stdio output logic and allow it to be reused. --- src/bin/pg_basebackup/pg_basebackup.c | 300 +++++++++++++++++----------------- 1 file changed, 154 insertions(+), 146 deletions(-)
diff --git a/src/bin/pg_basebackup/pg_basebackup.c b/src/bin/pg_basebackup/pg_basebackup.c index 80de882..fa942ab 100644 --- a/src/bin/pg_basebackup/pg_basebackup.c +++ b/src/bin/pg_basebackup/pg_basebackup.c @@ -51,6 +51,15 @@ typedef struct TablespaceList TablespaceListCell *tail; } TablespaceList; +typedef struct TarStream { + char path[MAXPGPATH]; + FILE *file; +#ifdef HAVE_LIBZ + gzFile zfile; +#endif + bool keepopen; +} TarStream; + /* Global options */ static char *basedir = NULL; static TablespaceList tablespace_dirs = {NULL, NULL}; @@ -100,6 +109,8 @@ static void disconnect_and_exit(int code); static void verify_dir_is_empty_or_create(char *dirname); static void progress_report(int tablespacenum, const char *filename, bool force); +static void OpenTarFile(TarStream *tarfile, const char *path); +static void CloseTarFile(TarStream *tarfile); static void ReceiveTarFile(PGconn *conn, PGresult *res, int rownum); static void ReceiveAndUnpackTarFile(PGconn *conn, PGresult *res, int rownum); static void GenerateRecoveryConf(PGconn *conn); @@ -725,169 +736,194 @@ parse_max_rate(char *src) * Write a piece of tar data */ static void -writeTarData( -#ifdef HAVE_LIBZ - gzFile ztarfile, -#endif - FILE *tarfile, char *buf, int r, char *current_file) +writeTarData(TarStream *stream, char *buf, int r) { #ifdef HAVE_LIBZ - if (ztarfile != NULL) + if (stream->zfile != NULL) { - if (gzwrite(ztarfile, buf, r) != r) + if (gzwrite(stream->zfile, buf, r) != r) { fprintf(stderr, _("%s: could not write to compressed file \"%s\": %s\n"), - progname, current_file, get_gz_error(ztarfile)); + progname, stream->path, get_gz_error(stream->zfile)); disconnect_and_exit(1); } } else #endif { - if (fwrite(buf, r, 1, tarfile) != 1) + if (fwrite(buf, r, 1, stream->file) != 1) { fprintf(stderr, _("%s: could not write to file \"%s\": %s\n"), - progname, current_file, strerror(errno)); + progname, stream->path, strerror(errno)); disconnect_and_exit(1); } } } -#ifdef HAVE_LIBZ -#define WRITE_TAR_DATA(buf, sz) writeTarData(ztarfile, tarfile, buf, sz, filename) -#else -#define WRITE_TAR_DATA(buf, sz) writeTarData(tarfile, buf, sz, filename) -#endif /* - * Receive a tar format file from the connection to the server, and write - * the data from this file directly into a tar file. If compression is - * enabled, the data will be compressed while written to the file. - * - * The file will be named base.tar[.gz] if it's for the main data directory - * or <tablespaceoid>.tar[.gz] if it's for another tablespace. - * - * No attempt to inspect or validate the contents of the file is done. + * Open a (possibly zlib-compressed) tar file for writing. A filename of - + * will cause stdout to be used. */ static void -ReceiveTarFile(PGconn *conn, PGresult *res, int rownum) +OpenTarFile(TarStream *tarfile, const char *path) { - char filename[MAXPGPATH]; - char *copybuf = NULL; - FILE *tarfile = NULL; - char tarhdr[512]; - bool basetablespace = PQgetisnull(res, rownum, 0); - bool in_tarhdr = true; - bool skip_file = false; - size_t tarhdrsz = 0; - size_t filesz = 0; + bool use_stdout; -#ifdef HAVE_LIBZ - gzFile ztarfile = NULL; -#endif + MemSet(tarfile, 0, sizeof(*tarfile)); + use_stdout = (strcmp(path, "-") == 0); + strlcpy(tarfile->path, path, sizeof(tarfile->path)); - if (basetablespace) +#ifdef HAVE_LIBZ + if (compresslevel != 0) { - /* - * Base tablespaces - */ - if (strcmp(basedir, "-") == 0) + /* Compression is in use */ + if (use_stdout) + tarfile->zfile = gzdopen(dup(fileno(stdout)), "wb"); + else + tarfile->zfile = gzopen(tarfile->path, "wb"); + if (!tarfile->zfile) { -#ifdef HAVE_LIBZ - if (compresslevel != 0) - { - ztarfile = gzdopen(dup(fileno(stdout)), "wb"); - if (gzsetparams(ztarfile, compresslevel, - Z_DEFAULT_STRATEGY) != Z_OK) - { - fprintf(stderr, - _("%s: could not set compression level %d: %s\n"), - progname, compresslevel, get_gz_error(ztarfile)); - disconnect_and_exit(1); - } - } - else -#endif - tarfile = stdout; - strcpy(filename, "-"); + fprintf(stderr, + _("%s: could not create compressed file \"%s\": %s\n"), + progname, tarfile->path, strerror(errno)); + disconnect_and_exit(1); } - else + if (gzsetparams(tarfile->zfile, compresslevel, + Z_DEFAULT_STRATEGY) != Z_OK) { -#ifdef HAVE_LIBZ - if (compresslevel != 0) - { - snprintf(filename, sizeof(filename), "%s/base.tar.gz", basedir); - ztarfile = gzopen(filename, "wb"); - if (gzsetparams(ztarfile, compresslevel, - Z_DEFAULT_STRATEGY) != Z_OK) - { - fprintf(stderr, - _("%s: could not set compression level %d: %s\n"), - progname, compresslevel, get_gz_error(ztarfile)); - disconnect_and_exit(1); - } - } - else -#endif - { - snprintf(filename, sizeof(filename), "%s/base.tar", basedir); - tarfile = fopen(filename, "wb"); - } + fprintf(stderr, + _("%s: could not set compression level %d: %s\n"), + progname, compresslevel, get_gz_error(tarfile->zfile)); + disconnect_and_exit(1); } } + else +#endif { - /* - * Specific tablespace - */ -#ifdef HAVE_LIBZ - if (compresslevel != 0) - { - snprintf(filename, sizeof(filename), "%s/%s.tar.gz", basedir, - PQgetvalue(res, rownum, 0)); - ztarfile = gzopen(filename, "wb"); - if (gzsetparams(ztarfile, compresslevel, - Z_DEFAULT_STRATEGY) != Z_OK) - { - fprintf(stderr, - _("%s: could not set compression level %d: %s\n"), - progname, compresslevel, get_gz_error(ztarfile)); - disconnect_and_exit(1); - } + /* Either no zlib support, or zlib support but compresslevel = 0 */ + if (use_stdout) { + tarfile->file = stdout; + tarfile->keepopen = true; } else -#endif + tarfile->file = fopen(tarfile->path, "wb"); + if (!tarfile->file) + { + fprintf(stderr, _("%s: could not create file \"%s\": %s\n"), + progname, tarfile->path, strerror(errno)); + disconnect_and_exit(1); + } + } +} + + +/* + * Close a tar file opened by OpenTarFile. The end-of-archive marker of two + * zero blocks will be written first. + */ +static void +CloseTarFile(TarStream *tarfile) +{ + /* 2 * 512 bytes empty data */ + static char zerobuf[1024]; + + /* + * Write two completely empty blocks at the end of the tar file, as + * required by some tar programs. + */ + writeTarData(tarfile, zerobuf, sizeof(zerobuf)); + + if (tarfile->keepopen) + { + if (fflush(tarfile->file) != 0) { - snprintf(filename, sizeof(filename), "%s/%s.tar", basedir, - PQgetvalue(res, rownum, 0)); - tarfile = fopen(filename, "wb"); + fprintf(stderr, + _("%s: could not write to file \"%s\": %s\n"), + progname, tarfile->path, strerror(errno)); + disconnect_and_exit(1); } + return; } #ifdef HAVE_LIBZ - if (compresslevel != 0) + if (tarfile->zfile != NULL) { - if (!ztarfile) + if (gzclose(tarfile->zfile) != 0) { - /* Compression is in use */ fprintf(stderr, - _("%s: could not create compressed file \"%s\": %s\n"), - progname, filename, get_gz_error(ztarfile)); + _("%s: could not close compressed file \"%s\": %s\n"), + progname, tarfile->path, get_gz_error(tarfile->zfile)); disconnect_and_exit(1); } } else #endif { - /* Either no zlib support, or zlib support but compresslevel = 0 */ - if (!tarfile) + if (fclose(tarfile->file) != 0) { - fprintf(stderr, _("%s: could not create file \"%s\": %s\n"), - progname, filename, strerror(errno)); + fprintf(stderr, + _("%s: could not close file \"%s\": %s\n"), + progname, tarfile->path, strerror(errno)); disconnect_and_exit(1); } } +} + + +/* + * Open a (possibly zlib-compressed) tar file for writing. The filebase + * argument should be the desired filename relative to basedir, without a .tar + * or .tar.gz file extension. If the user specified a basedir of - then stdout + * will be used. + */ +static void +openRelativeTarFile(TarStream *tarfile, const char *filebase) +{ + char path[MAXPGPATH]; + + if (strcmp(basedir, "-") == 0) + strlcpy(path, "-", sizeof(path)); +#ifdef HAVE_LIBZ + else if (compresslevel != 0) + snprintf(path, sizeof(path), "%s/%s.tar.gz", basedir, filebase); +#endif + else + snprintf(path, sizeof(path), "%s/%s.tar", basedir, filebase); + OpenTarFile(tarfile, path); +} + + +/* + * Receive a tar format file from the connection to the server, and write + * the data from this file directly into a tar file. If compression is + * enabled, the data will be compressed while written to the file. + * + * The file will be named base.tar[.gz] if it's for the main data directory + * or <tablespaceoid>.tar[.gz] if it's for another tablespace. + * + * No attempt to inspect or validate the contents of the file is done. + */ +static void +ReceiveTarFile(PGconn *conn, PGresult *res, int rownum) +{ + TarStream stream; + char *copybuf = NULL; + char tarhdr[512]; + bool basetablespace = PQgetisnull(res, rownum, 0); + bool in_tarhdr = true; + bool skip_file = false; + size_t tarhdrsz = 0; + size_t filesz = 0; + + if (basetablespace) + /* Base tablespace */ + openRelativeTarFile(&stream, "base"); + else + /* Specific tablespace */ + openRelativeTarFile(&stream, PQgetvalue(res, rownum, 0)); /* * Get the COPY data stream @@ -937,41 +973,13 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum) padding = ((recoveryconfcontents->len + 511) & ~511) - recoveryconfcontents->len; - WRITE_TAR_DATA(header, sizeof(header)); - WRITE_TAR_DATA(recoveryconfcontents->data, recoveryconfcontents->len); + writeTarData(&stream, header, sizeof(header)); + writeTarData(&stream, recoveryconfcontents->data, recoveryconfcontents->len); if (padding) - WRITE_TAR_DATA(zerobuf, padding); - } - - /* 2 * 512 bytes empty data at end of file */ - WRITE_TAR_DATA(zerobuf, sizeof(zerobuf)); - -#ifdef HAVE_LIBZ - if (ztarfile != NULL) - { - if (gzclose(ztarfile) != 0) - { - fprintf(stderr, - _("%s: could not close compressed file \"%s\": %s\n"), - progname, filename, get_gz_error(ztarfile)); - disconnect_and_exit(1); - } - } - else -#endif - { - if (strcmp(basedir, "-") != 0) - { - if (fclose(tarfile) != 0) - { - fprintf(stderr, - _("%s: could not close file \"%s\": %s\n"), - progname, filename, strerror(errno)); - disconnect_and_exit(1); - } - } + writeTarData(&stream, zerobuf, padding); } + CloseTarFile(&stream); break; } else if (r == -2) @@ -988,7 +996,7 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum) * tablespace, we never have to look for an existing recovery.conf * file in the stream. */ - WRITE_TAR_DATA(copybuf, r); + writeTarData(&stream, copybuf, r); } else { @@ -1059,7 +1067,7 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum) * header unmodified. */ if (!skip_file) - WRITE_TAR_DATA(tarhdr, 512); + writeTarData(&stream, tarhdr, 512); } } else @@ -1077,7 +1085,7 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum) bytes2write = (filesz > rr ? rr : filesz); if (!skip_file) - WRITE_TAR_DATA(copybuf + pos, bytes2write); + writeTarData(&stream, copybuf + pos, bytes2write); rr -= bytes2write; pos += bytes2write; @@ -1098,9 +1106,9 @@ ReceiveTarFile(PGconn *conn, PGresult *res, int rownum) } } totaldone += r; - progress_report(rownum, filename, false); + progress_report(rownum, stream.path, false); } /* while (1) */ - progress_report(rownum, filename, true); + progress_report(rownum, stream.path, true); if (copybuf != NULL) PQfreemem(copybuf); -- 2.3.0 -- Sent via pgsql-hackers mailing list (pgsql-hackers@postgresql.org) To make changes to your subscription: http://www.postgresql.org/mailpref/pgsql-hackers