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

…lastic#117148) The multi-get counterpart of elastic#114623.
DaveCTurner · Nov 22, 2024 · bcb29b7 · bcb29b7
1 parent a1247b3
commit bcb29b7
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/docs/changelog/117148.yaml b/docs/changelog/117148.yaml
@@ -0,0 +1,5 @@
+pr: 117148
+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/TransportShardMultiGetAction.java b/server/src/main/java/org/elasticsearch/action/get/TransportShardMultiGetAction.java
@@ -19,6 +19,7 @@
 import org.elasticsearch.action.ActionType;
 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.TransportActions;
 import org.elasticsearch.action.support.replication.BasicReplicationRequest;
 import org.elasticsearch.action.support.single.shard.TransportSingleShardAction;
@@ -280,15 +281,15 @@ private void tryShardMultiGetFromTranslog(
                     } else {
                         assert r.segmentGeneration() > -1L;
                         assert r.primaryTerm() > Engine.UNKNOWN_PRIMARY_TERM;
-                        indexShard.waitForPrimaryTermAndGeneration(
-                            r.primaryTerm(),
-                            r.segmentGeneration(),
+                        final ActionListener<Long> termAndGenerationListener = ContextPreservingActionListener.wrapPreservingContext(
                             listener.delegateFailureAndWrap(
                                 (ll, aLong) -> getExecutor(request, shardId).execute(
                                     ActionRunnable.supply(ll, () -> handleLocalGets(request, r.multiGetShardResponse(), shardId))
                                 )
-                            )
+                            ),
+                            threadPool.getThreadContext()
                         );
+                        indexShard.waitForPrimaryTermAndGeneration(r.primaryTerm(), r.segmentGeneration(), termAndGenerationListener);
                     }
                 }
             }), TransportShardMultiGetFomTranslogAction.Response::new, getExecutor(request, shardId))