Save pid and wait_event_info from PGPROC to avoid race condition

When using a pointer to PGPROC structure it can be changed somewhere else
and cause a race condition, which leads to an incorrect entry in
pg_wait_sampling. Saving pid and wait_event_info at the start
guarantees consistent data

Reviewed and amended by Sergey Shinderuk

Author: Oleg Tselebrovskiy

Diff for: collector.c

@@ -164,13 +164,9 @@ probe_waits(History *observations, HTAB *profile_hash,
 					   *observation;
 		PGPROC		   *proc = &ProcGlobal->allProcs[i];
 
-		if (!pgws_should_sample_proc(proc))
+		if (!pgws_should_sample_proc(proc, &item.pid, &item.wait_event_info))
 			continue;
 
-		/* Collect next wait event sample */
-		item.pid = proc->pid;
-		item.wait_event_info = proc->wait_event_info;
-
 		if (pgws_profileQueries)
 			item.queryId = pgws_proc_queryids[i];
 		else

Diff for: pg_wait_sampling.c

@@ -451,9 +451,15 @@ search_proc(int pid)
  * views.
  */
 bool
-pgws_should_sample_proc(PGPROC *proc)
+pgws_should_sample_proc(PGPROC *proc, int *pid_p, uint32 *wait_event_info_p)
 {
-	if (proc->wait_event_info == 0 && !pgws_sampleCpu)
+	int			pid = proc->pid;
+	uint32		wait_event_info = proc->wait_event_info;
+
+	*pid_p = pid;
+	*wait_event_info_p = wait_event_info;
+
+	if (wait_event_info == 0 && !pgws_sampleCpu)
 		return false;
 
 	/*
@@ -462,7 +468,7 @@ pgws_should_sample_proc(PGPROC *proc)
 	 * null wait events. So instead we make use of DisownLatch() resetting
 	 * owner_pid during ProcKill().
 	 */
-	if (proc->pid == 0 || proc->procLatch.owner_pid == 0 || proc->pid == MyProcPid)
+	if (pid == 0 || proc->procLatch.owner_pid == 0 || pid == MyProcPid)
 		return false;
 
 	return true;
@@ -533,7 +539,9 @@ pg_wait_sampling_get_current(PG_FUNCTION_ARGS)
 			{
 				PGPROC *proc = &ProcGlobal->allProcs[i];
 
-				if (!pgws_should_sample_proc(proc))
+				if (!pgws_should_sample_proc(proc,
+											 &params->items[j].pid,
+											 &params->items[j].wait_event_info))
 					continue;
 
 				params->items[j].pid = proc->pid;

Diff for: pg_wait_sampling.h

@@ -24,15 +24,15 @@
 
 typedef struct
 {
-	uint32			pid;
+	int				pid;
 	uint32			wait_event_info;
 	uint64			queryId;
 	uint64			count;
 } ProfileItem;
 
 typedef struct
 {
-	uint32			pid;
+	int				pid;
 	uint32			wait_event_info;
 	uint64			queryId;
 	TimestampTz		ts;
@@ -73,7 +73,7 @@ extern CollectorShmqHeader *pgws_collector_hdr;
 extern shm_mq			   *pgws_collector_mq;
 extern uint64			   *pgws_proc_queryids;
 extern void pgws_init_lock_tag(LOCKTAG *tag, uint32 lock);
-extern bool pgws_should_sample_proc(PGPROC *proc);
+extern bool pgws_should_sample_proc(PGPROC *proc, int *pid_p, uint32 *wait_event_info_p);
 
 /* collector.c */
 extern void pgws_register_wait_collector(void);