Added strict timeout (timeAllowed ) for shard tolerant requests

solr/core/src/java/org/apache/solr/handler/component/HttpShardHandler.java

@@ -223,7 +223,12 @@ protected ShardResponse transfomResponse(
    */
   @Override
   public ShardResponse takeCompletedIncludingErrors() {
-    return take(false);
+    return take(false, -1);
+  }
+
+  @Override
+  public ShardResponse takeCompletedIncludingErrorsWithTimeout(long maxAllowedTimeInMillis) {
+    return take(false, maxAllowedTimeInMillis);
   }
 
   /**
@@ -232,13 +237,22 @@ public ShardResponse takeCompletedIncludingErrors() {
    */
   @Override
   public ShardResponse takeCompletedOrError() {
-    return take(true);
+    return take(true, -1);
   }
 
-  private ShardResponse take(boolean bailOnError) {
+  private ShardResponse take(boolean bailOnError, long maxAllowedTimeInMillis) {
     try {
+      long deadline = System.nanoTime();
+      if (maxAllowedTimeInMillis > 0) {
+        deadline += TimeUnit.MILLISECONDS.toNanos(maxAllowedTimeInMillis);
+      } else {
+        deadline = Long.MAX_VALUE;
+      }
+
       while (pending.get() > 0) {
-        ShardResponse rsp = responses.take();
+        long waitTime = deadline - System.nanoTime();
+        ShardResponse rsp = responses.poll(waitTime, TimeUnit.NANOSECONDS);
+        if (rsp == null) return null;
         responseCancellableMap.remove(rsp);
 
         pending.decrementAndGet();
@@ -403,4 +417,14 @@ private boolean canShortCircuit(
   public ShardHandlerFactory getShardHandlerFactory() {
     return httpShardHandlerFactory;
   }
+
+  // test helper function
+  void setPendingRequest(int val) {
+    this.pending.set(val);
+  }
+
+  // test helper function
+  void setResponse(ShardResponse shardResponse) {
+    this.responses.add(shardResponse);
+  }
 }

solr/core/src/java/org/apache/solr/handler/component/SearchHandler.java

@@ -529,6 +529,7 @@ public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp) throw
       }
     } else {
       // a distributed request
+      long maxTimeAllowed = req.getParams().getLong(CommonParams.TIME_ALLOWED, -1);
 
       if (rb.outgoing == null) {
         rb.outgoing = new ArrayList<>();
@@ -593,7 +594,7 @@ public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp) throw
           while (rb.outgoing.size() == 0) {
             ShardResponse srsp =
                 tolerant
-                    ? shardHandler1.takeCompletedIncludingErrors()
+                    ? shardHandler1.takeCompletedIncludingErrorsWithTimeout(maxTimeAllowed)
                     : shardHandler1.takeCompletedOrError();
             if (srsp == null) break; // no more requests to wait for
 

solr/core/src/java/org/apache/solr/handler/component/ShardHandler.java

@@ -25,6 +25,9 @@ public abstract class ShardHandler {
 
   public abstract ShardResponse takeCompletedIncludingErrors();
 
+  public abstract ShardResponse takeCompletedIncludingErrorsWithTimeout(
+      long maxAllowedTimeInMillis);
+
   public abstract ShardResponse takeCompletedOrError();
 
   public abstract void cancelAll();

solr/core/src/test/org/apache/solr/core/MockShardHandlerFactory.java

@@ -48,6 +48,11 @@ public ShardResponse takeCompletedIncludingErrors() {
         return null;
       }
 
+      @Override
+      public ShardResponse takeCompletedIncludingErrorsWithTimeout(long maxAllowedTime) {
+        return null;
+      }
+
       @Override
       public ShardResponse takeCompletedOrError() {
         return null;

solr/core/src/test/org/apache/solr/handler/component/TestHttpShardHandlerFactory.java

@@ -27,6 +27,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import org.apache.solr.SolrTestCaseJ4;
 import org.apache.solr.client.solrj.impl.LBSolrClient;
 import org.apache.solr.client.solrj.request.QueryRequest;
@@ -155,4 +156,57 @@ public void testLiveNodesToHostUrl() {
     MatcherAssert.assertThat(hostSet, hasItem("1.2.3.4:9000"));
     MatcherAssert.assertThat(hostSet, hasItem("1.2.3.4:9001"));
   }
+
+  @Test
+  public void testHttpShardHandlerTimeout() {
+    HttpShardHandlerFactory httpShardHandlerFactory = new HttpShardHandlerFactory();
+    HttpShardHandler shardHandler = (HttpShardHandler) httpShardHandlerFactory.getShardHandler();
+
+    long startTime = System.nanoTime();
+    long timeAllowedInMillis = 120;
+    shardHandler.setPendingRequest(1);
+    ShardResponse shardResponse =
+        shardHandler.takeCompletedIncludingErrorsWithTimeout(timeAllowedInMillis);
+
+    assertEquals(null, shardResponse);
+
+    long timeTakenInMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime);
+    assertTrue(
+        "Should have taken more than 100 milli seconds " + timeTakenInMillis,
+        timeTakenInMillis > 100);
+    assertTrue(
+        "Should have timeout in 120 milli seconds " + timeTakenInMillis, timeTakenInMillis < 130);
+  }
+
+  @Test
+  public void testHttpShardHandlerWithResponse() {
+    HttpShardHandlerFactory httpShardHandlerFactory = new HttpShardHandlerFactory();
+    HttpShardHandler shardHandler = (HttpShardHandler) httpShardHandlerFactory.getShardHandler();
+
+    long startTime = System.nanoTime();
+    long timeAllowedInMillis = -1;
+    shardHandler.setPendingRequest(1);
+
+    ShardResponse shardResponse = new ShardResponse();
+    shardResponse.setShard("shard_1");
+    ShardRequest shardRequest = new ShardRequest();
+    shardRequest.actualShards = new String[] {"shard_1"};
+    shardResponse.setShardRequest(shardRequest);
+
+    Thread thread =
+        new Thread(
+            () -> {
+              shardHandler.setResponse(shardResponse);
+            });
+    thread.start();
+    ShardResponse gotResponse =
+        shardHandler.takeCompletedIncludingErrorsWithTimeout(timeAllowedInMillis);
+
+    assertEquals(shardResponse, gotResponse);
+
+    long timeTakenInMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTime);
+    assertTrue(
+        "Should have taken less than 100 milli seconds " + timeTakenInMillis,
+        timeTakenInMillis < 100);
+  }
 }