[PATCH 2/6] pg_basebackup: factor out tar open/close code for reusability.

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