[PATCH 2/6] pg_basebackup: factor out tar open/close code for reusability. - Mailing list pgsql-hackers

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);staticvoid 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);staticvoid 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
notwrite 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
lookfor 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