Add single-shard modification failure tests

I'm pretty sure a lot of this test functionality may be covered in some
of our existing regression tests, but I've included them to ensure we
put all failure-based tests under our new testing method for that kind
of test.

Didn't include lower replication factor, as (for a single-shard mod.),
it's indistinguishable from modifying a reference table. So these all
test modifications which hit a single, replicated shard.

src/test/regress/expected/failure_single_mod.out

@@ -0,0 +1,134 @@
+SELECT citus.mitmproxy('conn.allow()');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+SELECT citus.clear_network_traffic();
+ clear_network_traffic 
+-----------------------
+ 
+(1 row)
+
+SET citus.shard_count = 2;
+SET citus.shard_replication_factor = 2;
+CREATE TABLE mod_test (key int, value text);
+SELECT create_distributed_table('mod_test', 'key');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+-- verify behavior of single INSERT; should mark shard as failed
+SELECT citus.mitmproxy('conn.onQuery(query="^INSERT").kill()');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+INSERT INTO mod_test VALUES (2, 6);
+WARNING:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  while executing command on localhost:9060
+SELECT COUNT(*) FROM mod_test WHERE key=2;
+ count 
+-------
+     1
+(1 row)
+
+-- some clean up
+UPDATE pg_dist_shard_placement SET shardstate = 1
+WHERE shardid IN (
+  SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'mod_test'::regclass
+) AND shardstate = 3 RETURNING placementid;
+ placementid 
+-------------
+         137
+(1 row)
+
+TRUNCATE mod_test;
+-- verify behavior of UPDATE ... RETURNING; should mark as failed
+SELECT citus.mitmproxy('conn.allow()');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+INSERT INTO mod_test VALUES (2, 6);
+SELECT citus.mitmproxy('conn.onQuery(query="^UPDATE").kill()');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+UPDATE mod_test SET value='ok' WHERE key=2 RETURNING key;
+WARNING:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  while executing command on localhost:9060
+ key 
+-----
+   2
+(1 row)
+
+SELECT COUNT(*) FROM mod_test WHERE value='ok';
+ count 
+-------
+     1
+(1 row)
+
+-- some clean up
+UPDATE pg_dist_shard_placement SET shardstate = 1
+WHERE shardid IN (
+  SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'mod_test'::regclass
+) AND shardstate = 3 RETURNING placementid;
+ placementid 
+-------------
+         137
+(1 row)
+
+TRUNCATE mod_test;
+-- verify behavior of multi-statement modifications to a single shard
+-- should succeed but mark a placement as failed
+SELECT citus.mitmproxy('conn.onQuery(query="^UPDATE").kill()');
+ mitmproxy 
+-----------
+ 
+(1 row)
+
+BEGIN;
+INSERT INTO mod_test VALUES (2, 6);
+INSERT INTO mod_test VALUES (2, 7);
+DELETE FROM mod_test WHERE key=2 AND value = '7';
+UPDATE mod_test SET value='ok' WHERE key=2;
+WARNING:  server closed the connection unexpectedly
+	This probably means the server terminated abnormally
+	before or while processing the request.
+CONTEXT:  while executing command on localhost:9060
+COMMIT;
+WARNING:  connection not open
+CONTEXT:  while executing command on localhost:9060
+WARNING:  connection not open
+CONTEXT:  while executing command on localhost:9060
+WARNING:  connection not open
+CONTEXT:  while executing command on localhost:9060
+SELECT COUNT(*) FROM mod_test WHERE key=2;
+ count 
+-------
+     1
+(1 row)
+
+-- some clean up
+UPDATE pg_dist_shard_placement SET shardstate = 1
+WHERE shardid IN (
+  SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'mod_test'::regclass
+) AND shardstate = 3 RETURNING placementid;
+ placementid 
+-------------
+         137
+(1 row)
+
+TRUNCATE mod_test;
+-- ==== Clean up, we're done here ====
+DROP TABLE mod_test;

src/test/regress/failure_schedule

@@ -26,3 +26,4 @@ test: failure_single_select
 test: failure_ref_tables
 test: failure_real_time_select
 test: failure_insert_select_pushdown
+test: failure_single_mod

src/test/regress/sql/failure_single_mod.sql

@@ -0,0 +1,61 @@
+SELECT citus.mitmproxy('conn.allow()');
+SELECT citus.clear_network_traffic();
+
+SET citus.shard_count = 2;
+SET citus.shard_replication_factor = 2;
+
+CREATE TABLE mod_test (key int, value text);
+SELECT create_distributed_table('mod_test', 'key');
+
+-- verify behavior of single INSERT; should mark shard as failed
+SELECT citus.mitmproxy('conn.onQuery(query="^INSERT").kill()');
+INSERT INTO mod_test VALUES (2, 6);
+
+SELECT COUNT(*) FROM mod_test WHERE key=2;
+
+-- some clean up
+UPDATE pg_dist_shard_placement SET shardstate = 1
+WHERE shardid IN (
+  SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'mod_test'::regclass
+) AND shardstate = 3 RETURNING placementid;
+TRUNCATE mod_test;
+
+-- verify behavior of UPDATE ... RETURNING; should mark as failed
+SELECT citus.mitmproxy('conn.allow()');
+INSERT INTO mod_test VALUES (2, 6);
+
+SELECT citus.mitmproxy('conn.onQuery(query="^UPDATE").kill()');
+UPDATE mod_test SET value='ok' WHERE key=2 RETURNING key;
+
+SELECT COUNT(*) FROM mod_test WHERE value='ok';
+
+-- some clean up
+UPDATE pg_dist_shard_placement SET shardstate = 1
+WHERE shardid IN (
+  SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'mod_test'::regclass
+) AND shardstate = 3 RETURNING placementid;
+TRUNCATE mod_test;
+
+-- verify behavior of multi-statement modifications to a single shard
+-- should succeed but mark a placement as failed
+SELECT citus.mitmproxy('conn.onQuery(query="^UPDATE").kill()');
+
+BEGIN;
+INSERT INTO mod_test VALUES (2, 6);
+INSERT INTO mod_test VALUES (2, 7);
+DELETE FROM mod_test WHERE key=2 AND value = '7';
+UPDATE mod_test SET value='ok' WHERE key=2;
+COMMIT;
+
+SELECT COUNT(*) FROM mod_test WHERE key=2;
+
+-- some clean up
+UPDATE pg_dist_shard_placement SET shardstate = 1
+WHERE shardid IN (
+  SELECT shardid FROM pg_dist_shard WHERE logicalrelid = 'mod_test'::regclass
+) AND shardstate = 3 RETURNING placementid;
+TRUNCATE mod_test;
+
+-- ==== Clean up, we're done here ====
+
+DROP TABLE mod_test;