Preserve thread context when waiting for segment generation in RTG (e…

…lastic#114623) Closes ES-9778
georgewallace · Oct 25, 2024 · c7a5da4 · c7a5da4
1 parent 9420620
commit c7a5da4
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/docs/changelog/114623.yaml b/docs/changelog/114623.yaml
@@ -0,0 +1,5 @@
+pr: 114623
+summary: Preserve thread context when waiting for segment generation in RTG
+area: CRUD
+type: bug
+issues: []
diff --git a/server/src/main/java/org/elasticsearch/action/get/TransportGetAction.java b/server/src/main/java/org/elasticsearch/action/get/TransportGetAction.java
@@ -20,6 +20,7 @@
 import org.elasticsearch.action.NoShardAvailableActionException;
 import org.elasticsearch.action.admin.indices.refresh.TransportShardRefreshAction;
 import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.action.support.replication.BasicReplicationRequest;
 import org.elasticsearch.action.support.single.shard.TransportSingleShardAction;
 import org.elasticsearch.client.internal.node.NodeClient;
@@ -284,11 +285,11 @@ private void tryGetFromTranslog(GetRequest request, IndexShard indexShard, Disco
                     } else {
                         assert r.segmentGeneration() > -1L;
                         assert r.primaryTerm() > Engine.UNKNOWN_PRIMARY_TERM;
-                        indexShard.waitForPrimaryTermAndGeneration(
-                            r.primaryTerm(),
-                            r.segmentGeneration(),
-                            listener.delegateFailureAndWrap((ll, aLong) -> super.asyncShardOperation(request, shardId, ll))
+                        final ActionListener<Long> termAndGenerationListener = ContextPreservingActionListener.wrapPreservingContext(
+                            listener.delegateFailureAndWrap((ll, aLong) -> super.asyncShardOperation(request, shardId, ll)),
+                            threadPool.getThreadContext()
                         );
+                        indexShard.waitForPrimaryTermAndGeneration(r.primaryTerm(), r.segmentGeneration(), termAndGenerationListener);
                     }
                 }
             }), TransportGetFromTranslogAction.Response::new, getExecutor(request, shardId))