Minor fix

maksm90 · akorotkov · commit cdbf9bdbd303 · 2019-05-09T07:10:00.000+03:00
diff --git a/Makefile b/Makefile
@@ -1,7 +1,7 @@
 # contrib/pg_wait_sampling/Makefile
 
 MODULE_big = pg_wait_sampling
-OBJS = pg_wait_sampling.o collector.o
+OBJS = pg_wait_sampling.o collector.o compat.o
 
 EXTENSION = pg_wait_sampling
 EXTVERSION = 1.1
diff --git a/collector.c b/collector.c
@@ -219,15 +219,25 @@ send_history(History *observations, shm_mq_handle *mqh)
 
 	mq_result = shm_mq_send(mqh, sizeof(count), &count, false);
 	if (mq_result == SHM_MQ_DETACHED)
+	{
+		ereport(WARNING,
+				(errmsg("pg_wait_sampling collector: "
+						"receiver of message queue have been detached")));
 		return;
+	}
 	for (i = 0; i < count; i++)
 	{
 		mq_result = shm_mq_send(mqh,
 								sizeof(HistoryItem),
 								&observations->items[i],
 								false);
 		if (mq_result == SHM_MQ_DETACHED)
+		{
+			ereport(WARNING,
+					(errmsg("pg_wait_sampling collector: "
+							"receiver of message queue have been detached")));
 			return;
+		}
 	}
 }
 
@@ -240,12 +250,27 @@ send_profile(HTAB *profile_hash, shm_mq_handle *mqh)
 	HASH_SEQ_STATUS	scan_status;
 	ProfileItem	   *item;
 	Size			count = hash_get_num_entries(profile_hash);
+	shm_mq_result	mq_result;
 
-	shm_mq_send(mqh, sizeof(count), &count, false);
+	mq_result = shm_mq_send(mqh, sizeof(count), &count, false);
+	if (mq_result == SHM_MQ_DETACHED)
+	{
+		ereport(WARNING,
+				(errmsg("pg_wait_sampling collector: "
+						"receiver of message queue have been detached")));
+		return;
+	}
 	hash_seq_init(&scan_status, profile_hash);
 	while ((item = (ProfileItem *) hash_seq_search(&scan_status)) != NULL)
 	{
-		shm_mq_send(mqh, sizeof(ProfileItem), item, false);
+		mq_result = shm_mq_send(mqh, sizeof(ProfileItem), item, false);
+		if (mq_result == SHM_MQ_DETACHED)
+		{
+			ereport(WARNING,
+					(errmsg("pg_wait_sampling collector: "
+							"receiver of message queue have been detached")));
+			return;
+		}
 	}
 }
 
@@ -428,11 +453,7 @@ collector_main(Datum main_arg)
 						send_profile(profile_hash, mqh);
 					}
 				}
-#if PG_VERSION_NUM >= 100000
-				shm_mq_detach(mqh);
-#else
-				shm_mq_detach(collector_mq);
-#endif
+				shm_mq_detach_compat(mqh, collector_mq);
 			}
 			else if (request == PROFILE_RESET)
 			{
diff --git a/compat.c b/compat.c
@@ -0,0 +1,13 @@
+#include "postgres.h"
+
+#include "pg_wait_sampling.h"
+
+inline void
+shm_mq_detach_compat(shm_mq_handle *mqh, shm_mq *mq)
+{
+#if PG_VERSION_NUM >= 100000
+	shm_mq_detach(mqh);
+#else
+	shm_mq_detach(mq);
+#endif
+}
diff --git a/pg_wait_sampling.c b/pg_wait_sampling.c
@@ -530,7 +530,10 @@ receive_array(SHMRequest request, Size item_size, Size *count)
 
 	res = shm_mq_receive(mqh, &len, &data, false);
 	if (res != SHM_MQ_SUCCESS || len != sizeof(*count))
+	{
+		shm_mq_detach_compat(mqh, mq);
 		elog(ERROR, "Error reading mq.");
+	}
 	memcpy(count, data, sizeof(*count));
 
 	result = palloc(item_size * (*count));
@@ -541,22 +544,14 @@ receive_array(SHMRequest request, Size item_size, Size *count)
 		res = shm_mq_receive(mqh, &len, &data, false);
 		if (res != SHM_MQ_SUCCESS || len != item_size)
 		{
-#if PG_VERSION_NUM >= 100000
-			shm_mq_detach(mqh);
-#else
-			shm_mq_detach(mq);
-#endif
+			shm_mq_detach_compat(mqh, mq);
 			elog(ERROR, "Error reading mq.");
 		}
 		memcpy(ptr, data, item_size);
 		ptr += item_size;
 	}
 
-#if PG_VERSION_NUM >= 100000
-	shm_mq_detach(mqh);
-#else
-	shm_mq_detach(mq);
-#endif
+	shm_mq_detach_compat(mqh, mq);
 
 	LockRelease(&queueTag, ExclusiveLock, false);
 
diff --git a/pg_wait_sampling.h b/pg_wait_sampling.h
@@ -83,4 +83,6 @@ extern void register_wait_collector(void);
 extern void alloc_history(History *, int);
 extern void collector_main(Datum main_arg);
 
+extern void shm_mq_detach_compat(shm_mq_handle *mqh, shm_mq *mq);
+
 #endif
