Attempt at cleaning up getting the next bucket

The code for getting the next bucket was unnecessarily complex and hard
to read.

Author: jeltz

src/hash_query.c

@@ -137,8 +137,8 @@ pgsm_startup(void)
 static void
 InitializeSharedState(pgsmSharedState *pgsm)
 {
-	pg_atomic_init_u64(&pgsm->current_wbucket, 0);
-	pg_atomic_init_u64(&pgsm->prev_bucket_sec, 0);
+	pg_atomic_init_u64(&pgsm->current_bucket, 0);
+	pg_atomic_init_u64(&pgsm->current_bucket_sec, 0);
 }
 
 /*

src/hash_query.h

@@ -231,8 +231,8 @@ typedef struct pgsmEntry
 typedef struct pgsmSharedState
 {
 	LWLock	   *lock;			/* protects hashtable search/modification */
-	pg_atomic_uint64 current_wbucket;
-	pg_atomic_uint64 prev_bucket_sec;
+	pg_atomic_uint64 current_bucket;
+	pg_atomic_uint64 current_bucket_sec;
 	void	   *raw_dsa_area;	/* DSA area pointer to store query texts */
 	HTAB	   *hash_handle;
 

src/pg_stat_monitor.c

@@ -1712,7 +1712,7 @@ pgsm_store(pgsmEntry *entry)
 
 	pgsm = pgsm_get_ss();
 
-	prev_bucket_id = pg_atomic_read_u64(&pgsm->current_wbucket);
+	prev_bucket_id = pg_atomic_read_u64(&pgsm->current_bucket);
 	bucketid = get_next_wbucket(pgsm);
 
 	if (bucketid != prev_bucket_id)
@@ -2405,7 +2405,7 @@ pg_stat_monitor_internal(FunctionCallInfo fcinfo,
 		values[i++] = BoolGetDatum(toplevel);
 
 		/* bucket_done at column number 72 */
-		values[i++] = BoolGetDatum(pg_atomic_read_u64(&pgsm->current_wbucket) != bucketid);
+		values[i++] = BoolGetDatum(pg_atomic_read_u64(&pgsm->current_bucket) != bucketid);
 
 		/* clean up and return the tuplestore */
 		tuplestore_putvalues(tupstore, tupdesc, values, nulls);
@@ -2423,64 +2423,45 @@ static uint64
 get_next_wbucket(pgsmSharedState *pgsm)
 {
 	struct timeval tv;
-	uint64		current_bucket_sec;
-	bool		update_bucket = false;
+	uint64		new_bucket_start;
+	uint64		new_bucket_id;
 
 	gettimeofday(&tv, NULL);
-	current_bucket_sec = pg_atomic_read_u64(&pgsm->prev_bucket_sec);
+
+	new_bucket_start = (tv.tv_sec / pgsm_bucket_time) * (uint64) pgsm_bucket_time;
 
 	/*
-	 * If current bucket expired we loop attempting to update prev_bucket_sec.
-	 *
-	 * pg_atomic_compare_exchange_u64 may fail in two possible ways: 1.
-	 * Another thread/process updated the variable before us. 2. A spurious
-	 * failure / hardware event.
+	 * If current bucket expired we loop attempting to update current_bucket_sec.
 	 *
-	 * In both failure cases we read prev_bucket_sec from memory again, if it
-	 * was a spurious failure then the value of prev_bucket_sec must be the
-	 * same as before, which will cause the while loop to execute again.
+	 * If another thread updated current_bucket_sec, then use the new bucket
+	 * created by that thread.
 	 *
-	 * If another thread updated prev_bucket_sec, then its current value will
-	 * definitely make the while condition to fail, we can stop the loop as
-	 * another thread has already updated prev_bucket_sec.
+	 * XXX: This code looks potentially racy, especially if we skip buckets.
 	 */
-	while ((tv.tv_sec - (uint) current_bucket_sec) >= ((uint) pgsm_bucket_time))
-	{
-		if (pg_atomic_compare_exchange_u64(&pgsm->prev_bucket_sec, &current_bucket_sec, (uint64) tv.tv_sec))
-		{
-			update_bucket = true;
-			break;
-		}
-
-		current_bucket_sec = pg_atomic_read_u64(&pgsm->prev_bucket_sec);
-	}
-
-	if (update_bucket)
+	while (1)
 	{
-		uint64		new_bucket_id;
+		uint64		current_bucket_sec = pg_atomic_read_u64(&pgsm->current_bucket_sec);
 
-		new_bucket_id = (tv.tv_sec / pgsm_bucket_time) % pgsm_max_buckets;
+		if (current_bucket_sec + (uint64) pgsm_bucket_time > (uint64) tv.tv_sec)
+			return pg_atomic_read_u64(&pgsm->current_bucket);
 
-		/* Update bucket id and retrieve the previous one. */
-		pg_atomic_exchange_u64(&pgsm->current_wbucket, new_bucket_id);
-
-		pgsm_lock_aquire(pgsm, LW_EXCLUSIVE);
-		hash_entry_dealloc(new_bucket_id);
+		if (pg_atomic_compare_exchange_u64(&pgsm->current_bucket_sec, &current_bucket_sec, new_bucket_start))
+			break;
+	}
 
-		pgsm_lock_release(pgsm);
+	new_bucket_id = (new_bucket_start / pgsm_bucket_time) % pgsm_max_buckets;
 
-		/* Align the value in prev_bucket_sec to the bucket start time */
-		tv.tv_sec = (tv.tv_sec) - (tv.tv_sec % pgsm_bucket_time);
+	/* Update bucket id and retrieve the previous one. */
+	pg_atomic_exchange_u64(&pgsm->current_bucket, new_bucket_id);
 
-		pg_atomic_exchange_u64(&pgsm->prev_bucket_sec, (uint64) tv.tv_sec);
+	pgsm_lock_aquire(pgsm, LW_EXCLUSIVE);
+	hash_entry_dealloc(new_bucket_id);
+	pgsm_lock_release(pgsm);
 
-		pgsm->bucket_start_time[new_bucket_id] = (TimestampTz) tv.tv_sec -
-			((POSTGRES_EPOCH_JDATE - UNIX_EPOCH_JDATE) * SECS_PER_DAY);
-		pgsm->bucket_start_time[new_bucket_id] = pgsm->bucket_start_time[new_bucket_id] * USECS_PER_SEC;
-		return new_bucket_id;
-	}
+	pgsm->bucket_start_time[new_bucket_id] = (TimestampTz) (new_bucket_start -
+		((POSTGRES_EPOCH_JDATE - UNIX_EPOCH_JDATE) * SECS_PER_DAY)) * USECS_PER_SEC;
 
-	return pg_atomic_read_u64(&pgsm->current_wbucket);
+	return new_bucket_id;
 }
 
 /*