diff --git a/src/backend/distributed/executor/adaptive_executor.c b/src/backend/distributed/executor/adaptive_executor.c index 5d46bbf7d..a3b707230 100644 --- a/src/backend/distributed/executor/adaptive_executor.c +++ b/src/backend/distributed/executor/adaptive_executor.c @@ -2102,6 +2102,13 @@ RunDistributedExecution(DistributedExecution *execution) PG_TRY(); { + /* Preemptively step state machines in case of immediate errors */ + WorkerSession *session = NULL; + foreach_ptr(session, execution->sessionList) + { + ConnectionStateMachine(session); + } + bool cancellationReceived = false; int eventSetSize = GetEventSetSize(execution->sessionList); diff --git a/src/test/regress/expected/failure_savepoints.out b/src/test/regress/expected/failure_savepoints.out index a2a113e9a..873eaca5c 100644 --- a/src/test/regress/expected/failure_savepoints.out +++ b/src/test/regress/expected/failure_savepoints.out @@ -342,6 +342,39 @@ WHERE shardstate = 3 AND shardid IN ( (0 rows) TRUNCATE researchers; +-- test that we don't mark reference placements unhealthy +CREATE TABLE ref(a int, b int); +SELECT create_reference_table('ref'); + create_reference_table +--------------------------------------------------------------------- + +(1 row) + +SELECT citus.mitmproxy('conn.onQuery(query="^ROLLBACK").kill()'); + mitmproxy +--------------------------------------------------------------------- + +(1 row) + +BEGIN; +SAVEPOINT start; +INSERT INTO ref VALUES (1001,2); +SELECT * FROM ref; + a | b +--------------------------------------------------------------------- + 1001 | 2 +(1 row) + +ROLLBACK TO SAVEPOINT start; +WARNING: connection not open +WARNING: connection not open +SELECT * FROM ref; +WARNING: connection not open +WARNING: connection error: localhost:xxxxx +WARNING: connection not open +WARNING: connection not open +ERROR: connection error: localhost:xxxxx +END; -- clean up SELECT citus.mitmproxy('conn.allow()'); mitmproxy @@ -351,3 +384,4 @@ SELECT citus.mitmproxy('conn.allow()'); DROP TABLE artists; DROP TABLE researchers; +DROP TABLE ref; diff --git a/src/test/regress/sql/failure_savepoints.sql b/src/test/regress/sql/failure_savepoints.sql index a2b854215..29f74badf 100644 --- a/src/test/regress/sql/failure_savepoints.sql +++ b/src/test/regress/sql/failure_savepoints.sql @@ -188,7 +188,21 @@ WHERE shardstate = 3 AND shardid IN ( ) RETURNING placementid; TRUNCATE researchers; +-- test that we don't mark reference placements unhealthy +CREATE TABLE ref(a int, b int); +SELECT create_reference_table('ref'); + +SELECT citus.mitmproxy('conn.onQuery(query="^ROLLBACK").kill()'); +BEGIN; +SAVEPOINT start; +INSERT INTO ref VALUES (1001,2); +SELECT * FROM ref; +ROLLBACK TO SAVEPOINT start; +SELECT * FROM ref; +END; + -- clean up SELECT citus.mitmproxy('conn.allow()'); DROP TABLE artists; DROP TABLE researchers; +DROP TABLE ref;