Modify result immediately

2 files changed, 6 insertions, 3 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/federation/FederationResult.java b/container-search/src/main/java/com/yahoo/search/federation/FederationResult.java
index e26bf153eb3..5f1cfccf549 100644
--- a/container-search/src/main/java/com/yahoo/search/federation/FederationResult.java
+++ b/container-search/src/main/java/com/yahoo/search/federation/FederationResult.java
@@ -94,6 +94,7 @@ class FederationResult {
         public Optional<Result> getIfAvailable(long timeout) {
             if (availableResult.isPresent()) return availableResult;
             availableResult = futureResult.getIfAvailable(timeout, TimeUnit.MILLISECONDS);
+            availableResult.ifPresent(result -> target.modifyTargetResult(result));
             return availableResult;
         }
         
diff --git a/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java b/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java
index ec1836fd07b..c5573dc8fee 100644
--- a/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java
+++ b/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java
@@ -226,12 +226,15 @@ public class FederationSearcher extends ForkingSearcher {
         if (timeout <= 0) return Optional.empty();
 
         Execution newExecution = new Execution(target.getChain(), execution.context());
+        Result result;
         if (strictSearchchain) {
             query.resetTimeout();
-            return Optional.of(newExecution.search(createFederationQuery(query, query, Window.from(query), timeout, target)));
+            result = newExecution.search(createFederationQuery(query, query, Window.from(query), timeout, target));
         } else {
-            return Optional.of(newExecution.search(cloneFederationQuery(query, Window.from(query), timeout, target)));
+            result = newExecution.search(cloneFederationQuery(query, Window.from(query), timeout, target));
         }
+        target.modifyTargetResult(result);
+        return Optional.of(result);
     }
 
     private FederationResult search(Query query, Execution execution, Collection<Target> targets) {
@@ -558,7 +561,6 @@ public class FederationSearcher extends ForkingSearcher {
     }
 
     private void mergeResult(Query query, Target target, Result mergedResults, Result result) {
-        target.modifyTargetResult(result);
         ComponentId searchChainId = target.getId();
         Chain<Searcher> searchChain = target.getChain();