Fix MULTI_PROCESS parallel mode (QueueRead peer disconnect, lock groups)

Makefile

@@ -12,7 +12,7 @@ endif
 endif
 
 ifdef USE_PGXS
-PG_CONFIG = pg_config
+PG_CONFIG ?= pg_config
 PGXS := $(shell $(PG_CONFIG) --pgxs)
 include $(PGXS)
 else

bin/Makefile

@@ -21,7 +21,7 @@ endif
 endif
 
 ifdef USE_PGXS
-PG_CONFIG = pg_config
+PG_CONFIG ?= pg_config
 PGXS := $(shell $(PG_CONFIG) --pgxs)
 include $(PGXS)
 else

docs/pg_bulkload-ja.html

@@ -471,8 +471,9 @@ <h3>フォーマット共通の設定項目</h3>
 「WRITER=PARALLEL」と指定した場合は、MULTI_PROCESS は無視されます。
 なお、ロード先のデータベースに対してパスワード認証を必要とする場合には .pgpass を設定しなければなりません。
 詳細は<a href="#restrictions">使用上の注意と制約</a>を参照して下さい。
-pg_bulkloadのMULTI_PROCESSやPARALLELを有効にして実行する場合、他のPostgreSQLバックエンドプロセスがテーブルスキーマを変更しないようにしてください。
-データを読み取るプロセスとデータを書き出すプロセスで見られるテーブルスキーマが異なり、問題が発生する可能性があります。
+リーダーおよびライターのバックエンドが対象テーブルをどのようにロックするか、
+および他セッションがスキーマを変更してはいけない条件については、
+<a href="#parallel_locks">パラレルモードにおけるテーブルロック</a>を参照してください。
 </dd>
 
 </dl>
@@ -664,6 +665,36 @@ <h4>kill -9は使わない</h4>
 
 <h3>パラレルロードで使用する場合</h3>
 <p>パラレルロードで使用する場合(MULTI_PROCESS=YES または WRITER=PARALLEL)、以下のことに注意しなければなりません：</p>
+<p>
+pg_bulkload は 2 つの PostgreSQL バックエンドを使います。
+<em>リーダー</em>バックエンドがメインの <code>pg_bulkload()</code> を実行し（入力の読み取り・パース・検証）、
+<em>ライター</em>バックエンドが libpq 経由で起動され、TYPE=TUPLE の <code>pg_bulkload()</code> でテーブルへのダイレクトライトを行います。
+</p>
+<h4 id="parallel_locks">パラレルモードにおけるテーブルロック</h4>
+<p>
+PostgreSQL 9.6 以降では、リーダーバックエンドが PostgreSQL のロックグループのリーダーとなり、
+対象テーブルに <code>AccessShareLock</code> を保持したまま動作します。
+ライターバックエンドはテーブルにロックを取得する前にそのロックグループに参加し、
+ダイレクトライトのために <code>AccessExclusiveLock</code> を取得します。
+両バックエンドが同一ロックグループに属するため、これらのロックは互いにブロックせず、
+ライターが動作している間もリーダーは共有ロックを解放しません。
+ロードが進行中はライターの <code>AccessExclusiveLock</code> により、
+他セッションが対象テーブルに対する DDL を実行することもブロックされます。
+そのため、リーダーとライターが一貫したテーブル定義を参照するために、
+対象テーブルについて他セッションにスキーマ変更を避けるよう特別に注意する必要は通常ありません。
+</p>
+<p>
+PostgreSQL 9.6 より前のバージョン向けにビルドした場合は、
+リーダーバックエンドがライターバックエンドが <code>AccessExclusiveLock</code> を取得する前に
+<code>AccessShareLock</code> を解放します。その間に他セッションがテーブル定義を変更できるため、
+リーダーとライターで異なるテーブル定義を参照する可能性があります。
+ロード中は他セッションから対象テーブルに対する DDL などのスキーマ変更を行わないでください。
+</p>
+<p>
+PostgreSQL のバージョンにかかわらず、ロードが参照する他のオブジェクト
+（例: <a href="#FILTER">FILTER</a> 関数が参照する型やテーブル）の定義を、
+ロード中に変更しないでください。
+</p>
 <h4>認証における制約</h4>
 <p>MULTI_PROCESS=YESかつロード対象のデータベースにlocalhostから接続するのにパスワードが必要な場合、たとえパスワードを正しくプロンプトに入力しても、パスワード認証に失敗してしまいます。この問題を回避するには、以下のいずれかを設定してください。<br/>
 <ul>

docs/pg_bulkload.html

@@ -495,10 +495,9 @@ <h3>Common</h3>
 If WRITER is PARALLEL, MULTI_PROCESS is ignored.
 If password authentication is configured to the database to load,
 you have to set up the password file. See <a href="#restrictions">Restrictions</a> for details.
-Please make sure when enabling MULTI_PROCESS that no other PostgreSQL backend
-process is trying to modify the schema of the table, because it may cause the
-schema of the table as seen by the reader and the writer processes to differ
-and cause problems.
+See <a href="#parallel_locks">table locks in parallel mode</a> for how the
+reader and writer backends lock the target table and when other sessions must
+not change its schema.
 </dd>
 
 </dl>
@@ -728,6 +727,42 @@ <h4>Do not use <code>kill -9</code></h4>
 PostgreSQL to continue.
 <p>
 
+<h3>On parallel loading</h3>
+<p>
+When using MULTI_PROCESS=YES or WRITER=PARALLEL, pg_bulkload uses two PostgreSQL
+backends: a <em>reader</em> backend that runs the main <code>pg_bulkload()</code>
+call (read and parse input, validate rows) and a <em>writer</em> backend started
+via libpq that runs <code>pg_bulkload()</code> with TYPE=TUPLE and performs the
+direct write to the table.
+</p>
+
+<h4 id="parallel_locks">Table locks in parallel mode</h4>
+<p>
+On PostgreSQL 9.6 and later, the reader backend becomes a PostgreSQL lock group
+leader and keeps <code>AccessShareLock</code> on the target table.  The writer
+backend joins that lock group before it acquires any lock on the table, then
+takes <code>AccessExclusiveLock</code> for the direct write.  Because both
+backends belong to the same lock group, these locks do not block each other and
+the reader does not release its share lock while the writer is running.
+While the load is active, the writer's <code>AccessExclusiveLock</code> also
+blocks other sessions from running DDL on that table, so you do not need an
+extra precaution against concurrent schema changes on the target table for
+the reader and writer to stay consistent.
+</p>
+<p>
+On PostgreSQL releases before 9.6 (when pg_bulkload is built against them), the
+reader backend releases <code>AccessShareLock</code> before the writer backend
+acquires <code>AccessExclusiveLock</code>.  Another session can modify the table
+definition in that interval, and the reader and writer backends may then see
+different table definitions.  Do not run DDL or other schema changes on the
+target table from other sessions for the duration of the load.
+</p>
+<p>
+On any PostgreSQL version, avoid changing the definitions of other objects the
+load depends on (for example types or tables referenced by a
+<a href="#FILTER">FILTER</a> function) while the load is running.
+</p>
+
 <h4>Authentication can fail when MULTI_PROCESS=YES</h4>
 <p> When MULTI_PROCESS=YES and password is required to connect from localhost
 to the database to load, the authentication will fail even if you enter

include/reader.h

@@ -75,6 +75,9 @@ extern Parser *CreateBinaryParser(void);
 extern Parser *CreateCSVParser(void);
 extern Parser *CreateTupleParser(void);
 extern Parser *CreateFunctionParser(void);
+#if PG_VERSION_NUM >= 90600
+extern void SetTupleParserQueue(Parser *parser, void *queue);
+#endif
 
 #define ParserInit(self, checker, infile, relid, multi_process, collation)		((self)->init((self), (checker), (infile), (relid), (multi_process), (collation)))
 #define ParserRead(self, checker)					((self)->read((self), (checker)))

lib/Makefile

@@ -54,7 +54,7 @@ endif
 endif
 
 ifdef USE_PGXS
-PG_CONFIG = pg_config
+PG_CONFIG ?= pg_config
 PGXS := $(shell $(PG_CONFIG) --pgxs)
 include $(PGXS)
 else

lib/parser_tuple.c

@@ -54,9 +54,6 @@ CreateTupleParser(void)
 static void
 TupleParserInit(TupleParser *self, Checker *checker, const char *infile, TupleDesc desc, bool multi_process, Oid collation)
 {
-	unsigned		key;
-	char			junk[2];
-
 	if (checker->check_constraints)
 		ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 						errmsg("does not support parameter \"CHECK_CONSTRAINTS\" in \"TYPE = TUPLE\"")));
@@ -67,10 +64,19 @@ TupleParserInit(TupleParser *self, Checker *checker, const char *infile, TupleDe
 
 	checker->tchecker = NULL;
 
-	if (sscanf(infile, ":%u%1s", &key, junk) != 1)
-		elog(ERROR, "invalid shmem key format: %s", infile);
+#if PG_VERSION_NUM >= 90600
+	Assert(self->queue);
+#else
+	{
+		unsigned		key;
+		char			junk[2];
+
+		if (sscanf(infile, ":%u%1s", &key, junk) != 1)
+			elog(ERROR, "invalid shmem key format: %s", infile);
 
-	self->queue = QueueOpen(key);
+		self->queue = QueueOpen(key);
+	}
+#endif
 	self->buflen = BLCKSZ;
 	self->buffer = palloc(self->buflen);
 }
@@ -131,3 +137,11 @@ TupleParserDumpRecord(TupleParser *self, FILE *fp, char *filename)
 {
 	/* parse error does not happen in TupleParser. */
 }
+
+#if PG_VERSION_NUM >= 90600
+void
+SetTupleParserQueue(Parser *parser, void *queue)
+{
+	((TupleParser *) parser)->queue = queue;
+}
+#endif

lib/pg_bulkload.c

@@ -15,6 +15,7 @@
 
 #include "access/heapam.h"
 #include "access/reloptions.h"
+#include "c.h"
 #include "catalog/objectaddress.h"
 #if PG_VERSION_NUM >= 120000
 #include "catalog/pg_am.h"
@@ -41,6 +42,10 @@
 #include "pg_profile.h"
 #include "pg_strutil.h"
 #include "pgut/pgut-be.h"
+#include "pgut/pgut-ipc.h"
+#if PG_VERSION_NUM >= 90600
+#include "storage/proc.h"
+#endif
 
 PG_MODULE_MAGIC;
 
@@ -485,6 +490,11 @@ ParseOptions(Datum options, Reader **rd, Writer **wt, time_t tm)
 	char		   *type = NULL;
 	char		   *writer = NULL;
 	bool			multi_process = false;
+#if PG_VERSION_NUM >= 90600
+	char			*infile = NULL;
+	bool			parallel_writer = false;
+	Queue			*queue = NULL;
+#endif
 
 	Assert(*rd == NULL);
 	Assert(*wt == NULL);
@@ -505,6 +515,10 @@ ParseOptions(Datum options, Reader **rd, Writer **wt, time_t tm)
 		{
 			ASSERT_ONCE(type == NULL);
 			type = value;
+#if PG_VERSION_NUM >= 90600
+			if (!pg_strcasecmp(value, "TUPLE"))
+				parallel_writer = true;
+#endif
 		}
 		else if (CompareKeyword(keyword, "WRITER") ||
 				 CompareKeyword(keyword, "LOADER"))
@@ -516,16 +530,66 @@ ParseOptions(Datum options, Reader **rd, Writer **wt, time_t tm)
 		{
 			multi_process = ParseBoolean(value);
 		}
+#if PG_VERSION_NUM >= 90600
+		else if (CompareKeyword(keyword, "INPUT"))
+		{
+			ASSERT_ONCE(infile == NULL);
+			infile = pstrdup(value);
+		}
+#endif
 		else
 		{
 			rest_defs = lappend(rest_defs, opt);
 			continue;
 		}
 	}
 
+#if PG_VERSION_NUM >= 90600
+	/*
+	 * When running as the writer backend of a MULTI_PROCESS load
+	 * (TYPE=TUPLE, INPUT=:), join the reader's lock group before
+	 * DirectWriterParam() acquires a lock via the RangeVarGetRelidExtended() call.
+	 * This is necessary for the lock group mechanism to work. The process must
+	 * join the lock group before the first table lock acquisition to prevent
+	 * conflicts with locks acquired on the table by other group members.
+	 */
+	if (parallel_writer)
+	{
+		unsigned		key;
+		char			junk[2];
+		void		   *leader_pgproc;
+		int				leader_pid;
+		Assert(infile);
+
+		if (sscanf(infile, ":%u%1s", &key, junk) != 1)
+			elog(ERROR, "invalid shmem key format: %s", infile);
+
+		queue = QueueOpen(key);
+		QueueGetLockGroupInfo(queue, &leader_pgproc, &leader_pid);
+		if (leader_pgproc != NULL && leader_pid != 0)
+		{
+			if (!BecomeLockGroupMember((PGPROC *) leader_pgproc, leader_pid))
+				ereport(ERROR,
+						(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+						 errmsg("could not join lock group of parallel writer leader (pid %d)",
+								leader_pid)));
+		}
+	}
+#endif
+
 	*wt = WriterCreate(writer, multi_process);
 	*rd = ReaderCreate(type);
 
+#if PG_VERSION_NUM >= 90600
+	if (parallel_writer)
+	{
+		Assert(queue);
+		SetTupleParserQueue((*rd)->parser, queue);
+	}
+
+	(*rd)->infile = infile;
+#endif
+
 	foreach (cell, rest_defs)
 	{
 		opt = lfirst(cell);

lib/pgut/pgut-ipc.c

@@ -99,6 +99,10 @@ typedef struct QueueHeader
 	uint32		begin;		/* position that begins to read on data */
 	uint32		end;		/* position that begins to write on data */
 	slock_t		mutex;		/* locks shared variables begin, end and data */
+#if PG_VERSION_NUM >= 90600
+	void		*lock_group_leader;		/* PGPROC* of the lock group leader */
+	int			lock_group_leader_pid;	/* PID of the lock group leader */
+#endif
 	char		data[1];	/* VARIABLE LENGTH ARRAY - MUST BE LAST */
 } QueueHeader;
 
@@ -197,6 +201,10 @@ QueueCreate(unsigned *key, uint32 size)
 	header->size = size;
 	header->begin = header->end = 0;
 	SpinLockInit(&header->mutex);
+#if PG_VERSION_NUM >= 90600
+	header->lock_group_leader = NULL;
+	header->lock_group_leader_pid = 0;
+#endif
 
 	self = palloc(sizeof(Queue));
 	self->handle = handle;
@@ -370,6 +378,24 @@ QueueRead(Queue *self, void *buffer, uint32 len, bool need_lock)
 	CHECK_FOR_INTERRUPTS();
 	pg_usleep(SPIN_SLEEP_MSEC * 1000);
 
+#ifndef WIN32
+	/*
+	 * Detect if the writer detached without sending the
+	 * terminator (e.g. due to an error on the writer side).
+	 * When shm_nattch drops to 1, only we are still attached;
+	 * the writer is gone and no more data will ever arrive.
+	 */
+	{
+		struct shmid_ds	ds;
+
+		if (shmctl(self->handle, IPC_STAT, &ds) < 0 ||
+			ds.shm_nattch <= 1)
+			ereport(ERROR,
+					(errcode(ERRCODE_CONNECTION_FAILURE),
+					 errmsg("parallel writer has disconnected")));
+	}
+#endif
+
 	goto retry;
 }
 
@@ -509,3 +535,49 @@ QueueWrite(Queue *self, const struct iovec iov[], int count, uint32 timeout_msec
 	sleep_msec += SPIN_SLEEP_MSEC;
 	goto retry;
 }
+
+#if PG_VERSION_NUM >= 90600
+/**
+ * @brief Store lock group leader information in the queue header.
+ *
+ * The reader (lock group leader) calls this after QueueCreate() so
+ * that the writer, upon opening the same queue, can retrieve the
+ * information and join the lock group with BecomeLockGroupMember().
+ *
+ * @param self          [in] Queue handle returned by QueueCreate().
+ * @param leader_pgproc [in] PGPROC pointer of the lock group leader
+ *                           (typically MyProc of the calling backend).
+ * @param leader_pid    [in] PID of the lock group leader
+ *                           (typically MyProcPid of the calling backend).
+ */
+void
+QueueSetLockGroupInfo(Queue *self, void *leader_pgproc, int leader_pid)
+{
+	volatile QueueHeader *header = self->header;
+
+	header->lock_group_leader = leader_pgproc;
+	header->lock_group_leader_pid = leader_pid;
+}
+
+/**
+ * @brief Retrieve lock group leader information from the queue header.
+ *
+ * The writer calls this after QueueOpen() to obtain the PGPROC pointer
+ * and PID of the lock group leader so it can join the group with
+ * BecomeLockGroupMember() before acquiring any heavyweight locks.
+ *
+ * @param self          [in]  Queue handle returned by QueueOpen().
+ * @param leader_pgproc [out] Receives the PGPROC pointer of the leader,
+ *                            or NULL if no lock group info was stored.
+ * @param leader_pid    [out] Receives the PID of the leader,
+ *                            or 0 if no lock group info was stored.
+ */
+void
+QueueGetLockGroupInfo(Queue *self, void **leader_pgproc, int *leader_pid)
+{
+	volatile QueueHeader *header = self->header;
+
+	*leader_pgproc = header->lock_group_leader;
+	*leader_pid = header->lock_group_leader_pid;
+}
+#endif