[prev in list] [next in list] [prev in thread] [next in thread] 

List:       postgresql-general
Subject:    [HACKERS] [PATCH 2/6] pg_basebackup: factor out tar open/close code for reusability.
From:       Joshua Elsasser <josh () idealist ! org>
Date:       2015-09-29 22:16:24
Message-ID: 1443564988-17928-3-git-send-email-josh () idealist ! org
[Download RAW message or body]

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
[prev in list] [next in list] [prev in thread] [next in thread] 

Configure | About | News | Add a list | Sponsored by KoreLogic