Preserve thread context when waiting for segment generation in RTG (#117148)

The multi-get counterpart of
#114623.

Author: pxsalehi

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: []

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))