Scope Datastore transaction stacks to ApiProxy environment lifecycle.

In container environments where worker threads are pooled and reused across incoming HTTP requests (such as Jetty 12 QueuedThreadPool and virtual threads), if a Datastore transaction terminates exceptionally or fails to roll back cleanly, the transaction instance can remain on the thread-local transaction stack (TransactionStackImpl.STACK).

When BeginTransaction fails asynchronously, the future (SettableFuture) caches the DatastoreFailureException in memory. Any subsequent operation on that thread calling getCurrentTransaction() inherits the un-popped, poisoned transaction and immediately throws DatastoreFailureException ("Internal Datastore Error") locally without making remote network calls to the Appserver. This explains why backend RPC logs show no Datastore server errors during these error storms (b/494621464, b/537604854).

To fix this cleanly while preserving per-thread isolation for multi-threaded requests (e.g. ThreadManager or parallel execution):
- Updated TransactionStackImpl.StaticMember to track the current ApiProxy.Environment per thread via a ThreadLocal CURRENT_ENV.
- Whenever a thread transitions to a different ApiProxy.Environment (or leaving/entering a request boundary), STACK.get().clear() automatically purges leftover transaction state from the thread-local stack.

PiperOrigin-RevId: 963026907
Change-Id: Iefa6154f861036ba520b065f5ada3aa43aaab246

Author: ludoch

api/src/main/java/com/google/appengine/api/datastore/TransactionStackImpl.java

@@ -16,7 +16,9 @@
 
 package com.google.appengine.api.datastore;
 
+import com.google.apphosting.api.ApiProxy;
 import com.google.common.collect.Lists;
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
@@ -177,6 +179,8 @@ interface ThreadLocalTransactionStack {
     TransactionDataMap get();
 
     class StaticMember implements ThreadLocalTransactionStack {
+      private static final ThreadLocal<WeakReference<ApiProxy.Environment>> CURRENT_ENV =
+          new ThreadLocal<>();
       private static final ThreadLocal<TransactionDataMap> STACK =
           new ThreadLocal<TransactionDataMap>() {
             @Override
@@ -187,6 +191,13 @@ protected TransactionDataMap initialValue() {
 
       @Override
       public TransactionDataMap get() {
+        ApiProxy.Environment env = ApiProxy.getCurrentEnvironment();
+        WeakReference<ApiProxy.Environment> weakEnv = CURRENT_ENV.get();
+        ApiProxy.Environment prevEnv = weakEnv != null ? weakEnv.get() : null;
+        if (env != prevEnv) {
+          STACK.get().clear();
+          CURRENT_ENV.set(env != null ? new WeakReference<>(env) : null);
+        }
         return STACK.get();
       }
     }

api_dev/src/test/java/com/google/appengine/api/datastore/TransactionStackImplTest.java

@@ -21,6 +21,8 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.google.apphosting.api.ApiProxy;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -79,34 +81,17 @@ public void testLifecycle_Nested_Txns() {
     assertThat(stack.getFutures(txn3)).isEmpty();
 
     assertThat(stack.peek()).isSameInstanceAs(txn3);
-    assertThat(stack.getFutures(txn1)).isEmpty();
-    assertThat(stack.getFutures(txn2)).isEmpty();
-    assertThat(stack.getFutures(txn3)).isEmpty();
 
     assertThat(stack.pop()).isSameInstanceAs(txn3);
-    assertThat(stack.getFutures(txn1)).isEmpty();
-    assertThat(stack.getFutures(txn2)).isEmpty();
-    assertEmptyFutures(txn3);
-
+    assertThat(stack.getFutures(txn3)).isEmpty();
     assertThat(stack.peek()).isSameInstanceAs(txn2);
-    assertThat(stack.getFutures(txn1)).isEmpty();
-    assertThat(stack.getFutures(txn2)).isEmpty();
-    assertEmptyFutures(txn3);
 
     assertThat(stack.pop()).isSameInstanceAs(txn2);
-    assertThat(stack.getFutures(txn1)).isEmpty();
-    assertEmptyFutures(txn2);
-    assertEmptyFutures(txn3);
-
+    assertThat(stack.getFutures(txn2)).isEmpty();
     assertThat(stack.peek()).isSameInstanceAs(txn1);
-    assertThat(stack.getFutures(txn1)).isEmpty();
-    assertEmptyFutures(txn2);
-    assertEmptyFutures(txn3);
 
     assertThat(stack.pop()).isSameInstanceAs(txn1);
-    assertEmptyFutures(txn1);
-    assertEmptyFutures(txn2);
-    assertEmptyFutures(txn3);
+    assertThat(stack.getFutures(txn1)).isEmpty();
 
     // once we've deregistered, everything else should fail
     assertThat(stack.peek(null)).isNull();
@@ -209,4 +194,65 @@ public void testPeekOnEmptyStack() {
   public void testPushNull() {
     assertThrows(NullPointerException.class, () -> stack.push(null));
   }
+
+  @Test
+  public void testStaticMember_scopedToApiProxyEnvironment() {
+    TransactionStackImpl.ThreadLocalTransactionStack staticMember =
+        new TransactionStackImpl.ThreadLocalTransactionStack.StaticMember();
+    ApiProxy.Environment env1 = mock(ApiProxy.Environment.class);
+
+    ApiProxy.setEnvironmentForCurrentThread(env1);
+    try {
+      TransactionStackImpl.TransactionDataMap map1 = staticMember.get();
+      assertThat(map1).isNotNull();
+      Transaction txn1 = mock(Transaction.class);
+      when(txn1.getId()).thenReturn("123");
+      map1.txns.push(txn1);
+
+      // Subsequent call in same environment returns same stack with txn1
+      assertThat(staticMember.get().txns).containsExactly(txn1);
+
+      // Switching to a new environment (simulating request completion + thread reuse)
+      ApiProxy.Environment env2 = mock(ApiProxy.Environment.class);
+      ApiProxy.setEnvironmentForCurrentThread(env2);
+
+      // New environment triggers clear() on the thread-local stack
+      assertThat(staticMember.get().txns).isEmpty();
+    } finally {
+      ApiProxy.clearEnvironmentForCurrentThread();
+    }
+  }
+
+  @Test
+  public void testStaticMember_threadIsolation() throws Exception {
+    TransactionStackImpl.ThreadLocalTransactionStack staticMember =
+        new TransactionStackImpl.ThreadLocalTransactionStack.StaticMember();
+    ApiProxy.Environment env = mock(ApiProxy.Environment.class);
+    ApiProxy.setEnvironmentForCurrentThread(env);
+
+    try {
+      TransactionStackImpl.TransactionDataMap mainThreadMap = staticMember.get();
+      AtomicReference<TransactionStackImpl.TransactionDataMap> otherThreadMap =
+          new AtomicReference<>();
+
+      Thread otherThread =
+          new Thread(
+              () -> {
+                ApiProxy.setEnvironmentForCurrentThread(env);
+                try {
+                  otherThreadMap.set(staticMember.get());
+                } finally {
+                  ApiProxy.clearEnvironmentForCurrentThread();
+                }
+              });
+      otherThread.start();
+      otherThread.join();
+
+      // Even though both threads share the same ApiProxy.Environment, each thread has its own stack
+      assertThat(otherThreadMap.get()).isNotNull();
+      assertThat(otherThreadMap.get()).isNotSameInstanceAs(mainThreadMap);
+    } finally {
+      ApiProxy.clearEnvironmentForCurrentThread();
+    }
+  }
 }