aio: Fix crash potential for pg_aios views due to late state update

anarazel · anarazel · commit 0d9114b7040d · 2025-04-25T13:31:13.000-04:00
pgaio_io_reclaim() reset the fields in PgAioHandle before updating the state to IDLE or incrementing the generation. For most things that's OK, but for pg_get_aios() it is not - if it copied the PgAioHandle while fields were being reset, we wouldn't detect that and could call pgaio_io_get_target_description() with ioh->target == PGAIO_TID_INVALID, leading to a crash. Fix this issue by incrementing the generation and state earlier, before resetting. Also add an assertion to pgaio_io_get_target_description() for the target to be valid - that'd have made this case a bit easier to debug. While at it, add/update a few related assertions. Author: Alexander Lakhin <exclusion@gmail.com> Discussion: https://postgr.es/m/062daca9-dfad-4750-9da8-b13388301ad9@gmail.com
diff --git a/src/backend/storage/aio/aio.c b/src/backend/storage/aio/aio.c
@@ -670,6 +670,21 @@ pgaio_io_reclaim(PgAioHandle *ioh)
 
 	Assert(!ioh->resowner);
 
+	/*
+	 * Update generation & state first, before resetting the IO's fields,
+	 * otherwise a concurrent "viewer" could think the fields are valid, even
+	 * though they are being reset.  Increment the generation first, so that
+	 * we can assert elsewhere that we never wait for an IDLE IO.  While it's
+	 * a bit weird for the state to go backwards for a generation, it's OK
+	 * here, as there cannot be references to the "reborn" IO yet.  Can't
+	 * update both at once, so something has to give.
+	 */
+	ioh->generation++;
+	pgaio_io_update_state(ioh, PGAIO_HS_IDLE);
+
+	/* ensure the state update is visible before we reset fields */
+	pg_write_barrier();
+
 	ioh->op = PGAIO_OP_INVALID;
 	ioh->target = PGAIO_TID_INVALID;
 	ioh->flags = 0;
@@ -679,12 +694,6 @@ pgaio_io_reclaim(PgAioHandle *ioh)
 	ioh->result = 0;
 	ioh->distilled_result.status = PGAIO_RS_UNKNOWN;
 
-	/* XXX: the barrier is probably superfluous */
-	pg_write_barrier();
-	ioh->generation++;
-
-	pgaio_io_update_state(ioh, PGAIO_HS_IDLE);
-
 	/*
 	 * We push the IO to the head of the idle IO list, that seems more cache
 	 * efficient in cases where only a few IOs are used.
diff --git a/src/backend/storage/aio/aio_target.c b/src/backend/storage/aio/aio_target.c
@@ -49,7 +49,8 @@ pgaio_io_has_target(PgAioHandle *ioh)
 const char *
 pgaio_io_get_target_name(PgAioHandle *ioh)
 {
-	Assert(ioh->target >= 0 && ioh->target < PGAIO_TID_COUNT);
+	/* explicitly allow INVALID here, function used by debug messages */
+	Assert(ioh->target >= PGAIO_TID_INVALID && ioh->target < PGAIO_TID_COUNT);
 
 	return pgaio_target_info[ioh->target]->name;
 }
@@ -82,6 +83,9 @@ pgaio_io_get_target_data(PgAioHandle *ioh)
 char *
 pgaio_io_get_target_description(PgAioHandle *ioh)
 {
+	/* disallow INVALID, there wouldn't be a description */
+	Assert(ioh->target > PGAIO_TID_INVALID && ioh->target < PGAIO_TID_COUNT);
+
 	return pgaio_target_info[ioh->target]->describe_identity(&ioh->target_data);
 }
 
@@ -98,6 +102,8 @@ pgaio_io_get_target_description(PgAioHandle *ioh)
 bool
 pgaio_io_can_reopen(PgAioHandle *ioh)
 {
+	Assert(ioh->target > PGAIO_TID_INVALID && ioh->target < PGAIO_TID_COUNT);
+
 	return pgaio_target_info[ioh->target]->reopen != NULL;
 }
 
@@ -109,8 +115,8 @@ pgaio_io_can_reopen(PgAioHandle *ioh)
 void
 pgaio_io_reopen(PgAioHandle *ioh)
 {
-	Assert(ioh->target >= 0 && ioh->target < PGAIO_TID_COUNT);
-	Assert(ioh->op >= 0 && ioh->op < PGAIO_OP_COUNT);
+	Assert(ioh->target > PGAIO_TID_INVALID && ioh->target < PGAIO_TID_COUNT);
+	Assert(ioh->op > PGAIO_OP_INVALID && ioh->op < PGAIO_OP_COUNT);
 
 	pgaio_target_info[ioh->target]->reopen(ioh);
 }

Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,8 @@ pgaio_io_has_target(PgAioHandle *ioh)`
`49`	`49`	`const char *`
`50`	`50`	`pgaio_io_get_target_name(PgAioHandle *ioh)`
`51`	`51`	`{`
`52`		`- Assert(ioh->target >= 0 && ioh->target < PGAIO_TID_COUNT);`
	`52`	`+ /* explicitly allow INVALID here, function used by debug messages */`
	`53`	`+ Assert(ioh->target >= PGAIO_TID_INVALID && ioh->target < PGAIO_TID_COUNT);`
`53`	`54`
`54`	`55`	`return pgaio_target_info[ioh->target]->name;`
`55`	`56`	`}`
`@@ -82,6 +83,9 @@ pgaio_io_get_target_data(PgAioHandle *ioh)`
`82`	`83`	`char *`
`83`	`84`	`pgaio_io_get_target_description(PgAioHandle *ioh)`
`84`	`85`	`{`
	`86`	`+ /* disallow INVALID, there wouldn't be a description */`
	`87`	`+ Assert(ioh->target > PGAIO_TID_INVALID && ioh->target < PGAIO_TID_COUNT);`
	`88`	`+`
`85`	`89`	`return pgaio_target_info[ioh->target]->describe_identity(&ioh->target_data);`
`86`	`90`	`}`
`87`	`91`
`@@ -98,6 +102,8 @@ pgaio_io_get_target_description(PgAioHandle *ioh)`
`98`	`102`	`bool`
`99`	`103`	`pgaio_io_can_reopen(PgAioHandle *ioh)`
`100`	`104`	`{`
	`105`	`+ Assert(ioh->target > PGAIO_TID_INVALID && ioh->target < PGAIO_TID_COUNT);`
	`106`	`+`
`101`	`107`	`return pgaio_target_info[ioh->target]->reopen != NULL;`
`102`	`108`	`}`
`103`	`109`
`@@ -109,8 +115,8 @@ pgaio_io_can_reopen(PgAioHandle *ioh)`
`109`	`115`	`void`
`110`	`116`	`pgaio_io_reopen(PgAioHandle *ioh)`
`111`	`117`	`{`
`112`		`- Assert(ioh->target >= 0 && ioh->target < PGAIO_TID_COUNT);`
`113`		`- Assert(ioh->op >= 0 && ioh->op < PGAIO_OP_COUNT);`
	`118`	`+ Assert(ioh->target > PGAIO_TID_INVALID && ioh->target < PGAIO_TID_COUNT);`
	`119`	`+ Assert(ioh->op > PGAIO_OP_INVALID && ioh->op < PGAIO_OP_COUNT);`
`114`	`120`
`115`	`121`	`pgaio_target_info[ioh->target]->reopen(ioh);`
`116`	`122`	`}`