[#25105] YSQL: Clear pggate session state on query retry

Summary:
During the course of its execution, a query accumulates state in its pggate session. Some of the state is stored in:
 - The foreign key buffer which holds a collections of keys to be inspected, as well as a collection of keys thus far seen by the transaction.
 - The explicit row lock buffer which holds a collection of rows that are to be locked explicitly.
 - The insert on conflict buffer which holds a collection of keys to be inspected as well a collection of rows thus far modified by the query.

Currently, this state is cleared only at the transaction boundary by operations such as recreating, restarting and aborting a DML transaction.
Notably, this state is NOT cleared when a sub-transaction is rolled back. An example of this is when a query is retried (such as due to a kConflict or a kReadRestart error) in read committed mode.
This works fine with the row lock buffer which is largely idempotent (adding the same key/row to the buffer upon retry did not affect the state). However, it can cause issues with the foreign key buffer ([[ #25234 | GH-25234 ]]).

D39023 introduced a new insert-on-conflict buffer which is not idempotent.
One of the functions of this buffer is to ensure that a row is not modified twice within the same transaction.
Since two iterations of the same query, are likely to modify the same row(s), residual state can lead to the buffer incorrectly deducing that the same row has been modified twice.

This revision:
 - Introduces the distinction that some session state is sub-transaction-scoped ie. needs to be cleared at the sub-transaction boundary (such as for each retry/iteration of a query), while others are transaction-scoped ie. can be cleared at the transaction boundary only.
 - Explicitly clears sub-transaction-scoped session state when the subtransaction is rolled back.
 - Does not alter the behavior of the foreign key buffer.
Jira: DB-14251

Test Plan:
Run the following automated test:
```
./yb_build.sh --java-test 'org.yb.pgsql.TestPgRegressIsolationWithoutWaitQueues#testPgRegress'
./yb_build.sh --cxx-test pgwrapper_pg_libpq-test --gtest_filter 'BatchInsertOnConflictTest.InsertOnConflictWithQueryRestart'
```

Jenkins: urgent

Reviewers: patnaik.balivada, pjain, telgersma

Reviewed By: patnaik.balivada, pjain

Subscribers: smishra, yql

Tags: #jenkins-ready

Differential Revision: https://phorge.dev.yugabyte.com/D40460

Author: karthik-ramanathan-3006

src/yb/yql/pggate/pggate.cc

@@ -1844,12 +1844,12 @@ Status PgApiImpl::BeginTransaction(int64_t start_time) {
 }
 
 Status PgApiImpl::RecreateTransaction() {
-  ClearSessionState();
+  RollbackTransactionScopedSessionState();
   return pg_txn_manager_->RecreateTransaction();
 }
 
 Status PgApiImpl::RestartTransaction() {
-  ClearSessionState();
+  RollbackTransactionScopedSessionState();
   return pg_txn_manager_->RestartTransaction();
 }
 
@@ -1870,15 +1870,12 @@ bool PgApiImpl::IsRestartReadPointRequested() {
 }
 
 Status PgApiImpl::CommitPlainTransaction() {
-  DCHECK(pg_session_->explicit_row_lock_buffer().IsEmpty());
-  DCHECK(pg_session_->IsInsertOnConflictBufferEmpty());
-  pg_session_->InvalidateForeignKeyReferenceCache();
-  RETURN_NOT_OK(pg_session_->FlushBufferedOperations());
+  RETURN_NOT_OK(CommitTransactionScopedSessionState());
   return pg_txn_manager_->CommitPlainTransaction();
 }
 
 Status PgApiImpl::AbortPlainTransaction() {
-  ClearSessionState();
+  RollbackTransactionScopedSessionState();
   return pg_txn_manager_->AbortPlainTransaction();
 }
 
@@ -1931,7 +1928,7 @@ Status PgApiImpl::ExitSeparateDdlTxnMode(PgOid db_oid, bool is_silent_modificati
 }
 
 Status PgApiImpl::ClearSeparateDdlTxnMode() {
-  ClearSessionState();
+  RollbackTransactionScopedSessionState();
   return pg_txn_manager_->ExitSeparateDdlTxnModeWithAbort();
 }
 
@@ -1941,8 +1938,7 @@ Status PgApiImpl::SetActiveSubTransaction(SubTransactionId id) {
 }
 
 Status PgApiImpl::RollbackToSubTransaction(SubTransactionId id) {
-  pg_session_->DropBufferedOperations();
-  pg_session_->explicit_row_lock_buffer().Clear();
+  RollbackSubTransactionScopedSessionState();
   return pg_session_->RollbackToSubTransaction(id);
 }
 
@@ -2236,13 +2232,31 @@ Result<tserver::PgServersMetricsResponsePB> PgApiImpl::ServersMetrics() {
     return pg_session_->ServersMetrics();
 }
 
-void PgApiImpl::ClearSessionState() {
-  pg_session_->InvalidateForeignKeyReferenceCache();
+void PgApiImpl::RollbackSubTransactionScopedSessionState() {
   pg_session_->DropBufferedOperations();
   pg_session_->explicit_row_lock_buffer().Clear();
   pg_session_->ClearAllInsertOnConflictBuffers();
 }
 
+void PgApiImpl::RollbackTransactionScopedSessionState() {
+  RollbackSubTransactionScopedSessionState();
+  pg_session_->InvalidateForeignKeyReferenceCache();
+}
+
+Status PgApiImpl::CommitTransactionScopedSessionState() {
+  RSTATUS_DCHECK(
+      pg_session_->explicit_row_lock_buffer().IsEmpty(),
+      IllegalState,
+      "Expected row lock buffer to be empty");
+  RSTATUS_DCHECK(
+      pg_session_->IsInsertOnConflictBufferEmpty(),
+      IllegalState,
+      "Expected INSERT ... ON CONFLICT buffer to be empty");
+  RETURN_NOT_OK(pg_session_->FlushBufferedOperations());
+  pg_session_->InvalidateForeignKeyReferenceCache();
+  return Status::OK();
+}
+
 bool PgApiImpl::IsCronLeader() const { return tserver_shared_object_->IsCronLeader(); }
 
 Status PgApiImpl::SetCronLastMinute(int64_t last_minute) {

src/yb/yql/pggate/pggate.h

@@ -773,6 +773,10 @@ class PgApiImpl {
 
   void RestoreSessionState(const YBCPgSessionState& session_data);
 
+  void RollbackSubTransactionScopedSessionState();
+  void RollbackTransactionScopedSessionState();
+  Status CommitTransactionScopedSessionState();
+
   //------------------------------------------------------------------------------------------------
   // Replication Slots Functions.
 
@@ -836,8 +840,6 @@ class PgApiImpl {
   Status ReleaseAllAdvisoryLocks(uint32_t db_oid);
 
  private:
-  void ClearSessionState();
-
   Result<bool> RetrieveYbctidsImpl(
       PgDmlRead& dml_read, int natts, size_t max_mem_bytes, std::vector<Slice>& ybctids);
 

src/yb/yql/pgwrapper/pg_libpq-test.cc

@@ -4509,6 +4509,63 @@ TEST_F_EX(PgLibPqTest, BlockDangerousRoles, PgLibPqBlockDangerousRolesTest) {
                            YBPgErrorCode::YB_PG_RESERVED_NAME);
 }
 
+class BatchInsertOnConflictTest : public PgLibPqTestRF1 {
+ protected:
+  void RunConflictingIOCTxn(PGConn& conn, IsolationLevel isolation_level, int32_t expected_price) {
+    thread_holder_.AddThreadFunctor([this, &conn, isolation_level, expected_price]() -> void {
+      ASSERT_OK(conn.StartTransaction(isolation_level));
+      ASSERT_OK(conn.Execute(kIOCQuery));
+      const int32_t price = ASSERT_RESULT(conn.FetchRow<int32_t>(
+          "SELECT price FROM products WHERE id = 1"));
+
+      ASSERT_EQ(price, expected_price);
+      ASSERT_OK(conn.CommitTransaction());
+    });
+  }
+
+  void StartMainTxn(PGConn& conn, IsolationLevel isolation_level = READ_COMMITTED) {
+    ASSERT_OK(conn.StartTransaction(isolation_level));
+    ASSERT_OK(conn.Execute("UPDATE products SET price = price + 5 WHERE id = 1"));
+  }
+
+  void CommitMainTxnAfterWait(PGConn& conn) {
+    std::this_thread::sleep_for(RandomUniformInt(1, 10) * 100ms);
+    ASSERT_OK(conn.CommitTransaction());
+  }
+
+  constexpr static auto kInsertOnConflictBatchSize = 1024;
+  const std::vector<IsolationLevel> kIsolationLevels = {
+    READ_COMMITTED, SNAPSHOT_ISOLATION, SERIALIZABLE_ISOLATION
+  };
+  const std::string kIOCQuery = "INSERT INTO products VALUES (1, 'oats', 10) ON CONFLICT (id) "
+                                "DO UPDATE SET price = products.price + 5";
+  TestThreadHolder thread_holder_;
+};
+
+TEST_F(BatchInsertOnConflictTest, InsertOnConflictWithQueryRestart) {
+  PGConn conn1 = ASSERT_RESULT(Connect());
+  PGConn conn2 = ASSERT_RESULT(Connect());
+  int32_t expected_price = 10;
+
+  // Setup
+  ASSERT_OK(conn1.Execute("DROP TABLE IF EXISTS products"));
+  ASSERT_OK(conn1.Execute("CREATE TABLE products (id INT PRIMARY KEY, name TEXT, price INT)"));
+  ASSERT_OK(conn1.Execute("INSERT INTO products VALUES (1, 'oats', 10)"));
+  ASSERT_OK(conn2.Execute(
+      Format("SET yb_insert_on_conflict_read_batch_size TO $0", kInsertOnConflictBatchSize)));
+
+  // Test that a query retried due to a kConflict error correctly clears the insert on conflict
+  // buffer between retries. Not clearing the buffer would result in an incorrect "command cannot
+  // affect row a second time" error.
+  for (IsolationLevel isolation_level : kIsolationLevels) {
+    expected_price += 10;
+    StartMainTxn(conn1);
+    RunConflictingIOCTxn(conn2, isolation_level, expected_price);
+    CommitMainTxnAfterWait(conn1);
+    thread_holder_.WaitAndStop(2s * kTimeMultiplier);
+  }
+}
+
 // https://github.com/yugabyte/yugabyte-db/issues/24320.
 TEST_F(PgLibPqTest, TableRewriteOidCollision) {
   PGConn conn = ASSERT_RESULT(Connect());