apache · janhoy · Sep 6, 2023 · Sep 6, 2023
diff --git a/solr/CHANGES.txt b/solr/CHANGES.txt
@@ -72,6 +72,8 @@ New Features
 ---------------------
 * SOLR-16654: Add support for node-level caches (Michael Gibney)
 
+* SOLR-16954: Make Circuit Breakers available for Update Requests (janhoy, Christine Poerschke, Pierre Salagnac)
+
 Improvements
 ---------------------
 * SOLR-16490: `/admin/cores?action=backupcore` now has a v2 equivalent, available at

diff --git a/solr/core/src/java/org/apache/solr/handler/ContentStreamHandlerBase.java b/solr/core/src/java/org/apache/solr/handler/ContentStreamHandlerBase.java
@@ -16,6 +16,11 @@
  */
 package org.apache.solr.handler;
 
+import static org.apache.solr.common.params.CommonParams.FAILURE;
+import static org.apache.solr.common.params.CommonParams.STATUS;
+
+import java.util.List;
+import org.apache.solr.client.solrj.SolrRequest.SolrRequestType;
 import org.apache.solr.common.SolrException;
 import org.apache.solr.common.params.SolrParams;
 import org.apache.solr.common.util.ContentStream;
@@ -26,6 +31,8 @@
 import org.apache.solr.update.SolrCoreState;
 import org.apache.solr.update.processor.UpdateRequestProcessor;
 import org.apache.solr.update.processor.UpdateRequestProcessorChain;
+import org.apache.solr.util.circuitbreaker.CircuitBreaker;
+import org.apache.solr.util.circuitbreaker.CircuitBreakerRegistry;
 
 /**
  * Shares common code between various handlers that manipulate {@link
@@ -49,6 +56,10 @@ public void init(NamedList<?> args) {
 
   @Override
   public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp) throws Exception {
+    if (checkCircuitBreakers(req, rsp)) {
+      return; // Circuit breaker tripped, return immediately
+    }
+
     /*
        We track update requests so that we can preserve consistency by waiting for them to complete
        on a node shutdown and then immediately trigger a leader election without waiting for the core to close.
@@ -101,6 +112,30 @@ public void handleRequestBody(SolrQueryRequest req, SolrQueryResponse rsp) throw
     }
   }
 
+  /**
+   * Check if {@link SolrRequestType#UPDATE} circuit breakers are tripped. Override this method in
+   * sub classes that do not want to check circuit breakers.
+   *
+   * @return true if circuit breakers are tripped, false otherwise.
+   */
+  protected boolean checkCircuitBreakers(SolrQueryRequest req, SolrQueryResponse rsp) {
+    CircuitBreakerRegistry circuitBreakerRegistry = req.getCore().getCircuitBreakerRegistry();
+    if (circuitBreakerRegistry.isEnabled(SolrRequestType.UPDATE)) {
+      List<CircuitBreaker> trippedCircuitBreakers =
+          circuitBreakerRegistry.checkTripped(SolrRequestType.UPDATE);
+      if (trippedCircuitBreakers != null) {
+        String errorMessage = CircuitBreakerRegistry.toErrorMessage(trippedCircuitBreakers);
+        rsp.add(STATUS, FAILURE);
+        rsp.setException(
+            new SolrException(
+                CircuitBreaker.getErrorCode(trippedCircuitBreakers),
+                "Circuit Breakers tripped " + errorMessage));
+        return true;
+      }
+    }
+    return false;
+  }
+
   protected abstract ContentStreamLoader newLoader(
       SolrQueryRequest req, UpdateRequestProcessor processor);
 }
diff --git a/solr/core/src/java/org/apache/solr/handler/component/SearchHandler.java b/solr/core/src/java/org/apache/solr/handler/component/SearchHandler.java
@@ -35,6 +35,7 @@
 import java.util.concurrent.atomic.AtomicLong;
 import org.apache.lucene.index.ExitableDirectoryReader;
 import org.apache.lucene.search.TotalHits;
+import org.apache.solr.client.solrj.SolrRequest.SolrRequestType;
 import org.apache.solr.client.solrj.SolrServerException;
 import org.apache.solr.cloud.ZkController;
 import org.apache.solr.common.SolrDocumentList;
@@ -330,8 +331,8 @@ protected ResponseBuilder newResponseBuilder(
   }
 
   /**
-   * Check if circuit breakers are tripped. Override this method in sub classes that do not want to
-   * check circuit breakers.
+   * Check if {@link SolrRequestType#QUERY} circuit breakers are tripped. Override this method in
+   * sub classes that do not want to check circuit breakers.
    *
    * @return true if circuit breakers are tripped, false otherwise.
    */
@@ -340,26 +341,26 @@ protected boolean checkCircuitBreakers(
     final RTimerTree timer = rb.isDebug() ? req.getRequestTimer() : null;
 
     final CircuitBreakerRegistry circuitBreakerRegistry = req.getCore().getCircuitBreakerRegistry();
-    if (circuitBreakerRegistry.isEnabled()) {
+    if (circuitBreakerRegistry.isEnabled(SolrRequestType.QUERY)) {
       List<CircuitBreaker> trippedCircuitBreakers;
 
       if (timer != null) {
         RTimerTree subt = timer.sub("circuitbreaker");
         rb.setTimer(subt);
 
-        trippedCircuitBreakers = circuitBreakerRegistry.checkTripped();
+        trippedCircuitBreakers = circuitBreakerRegistry.checkTripped(SolrRequestType.QUERY);
 
         rb.getTimer().stop();
       } else {
-        trippedCircuitBreakers = circuitBreakerRegistry.checkTripped();
+        trippedCircuitBreakers = circuitBreakerRegistry.checkTripped(SolrRequestType.QUERY);
       }
 
       if (trippedCircuitBreakers != null) {
         String errorMessage = CircuitBreakerRegistry.toErrorMessage(trippedCircuitBreakers);
         rsp.add(STATUS, FAILURE);
         rsp.setException(
             new SolrException(
-                SolrException.ErrorCode.SERVICE_UNAVAILABLE,
+                CircuitBreaker.getErrorCode(trippedCircuitBreakers),
                 "Circuit Breakers tripped " + errorMessage));
         return true;
       }

diff --git a/solr/core/src/java/org/apache/solr/util/circuitbreaker/CircuitBreaker.java b/solr/core/src/java/org/apache/solr/util/circuitbreaker/CircuitBreaker.java
@@ -17,6 +17,12 @@
 
 package org.apache.solr.util.circuitbreaker;
 
+import java.util.List;
+import java.util.Locale;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.apache.solr.client.solrj.SolrRequest.SolrRequestType;
+import org.apache.solr.common.SolrException;
 import org.apache.solr.common.util.NamedList;
 import org.apache.solr.util.SolrPluginUtils;
 import org.apache.solr.util.plugin.NamedListInitializedPlugin;
@@ -36,6 +42,10 @@
  * @lucene.experimental
  */
 public abstract class CircuitBreaker implements NamedListInitializedPlugin {
+  // Only query requests are checked by default
+  private Set<SolrRequestType> requestTypes = Set.of(SolrRequestType.QUERY);
+  private final List<SolrRequestType> SUPPORTED_TYPES =
+      List.of(SolrRequestType.QUERY, SolrRequestType.UPDATE);
 
   @Override
   public void init(NamedList<?> args) {
@@ -52,4 +62,48 @@ public CircuitBreaker() {}
 
   /** Get error message when the circuit breaker triggers */
   public abstract String getErrorMessage();
+
+  /**
+   * Set the request types for which this circuit breaker should be checked. If not called, the
+   * circuit breaker will be checked for the {@link SolrRequestType#QUERY} request type only.
+   *
+   * @param requestTypes list of strings representing request types
+   * @throws IllegalArgumentException if the request type is not valid
+   */
+  public void setRequestTypes(List<String> requestTypes) {
+    this.requestTypes =
+        requestTypes.stream()
+            .map(t -> SolrRequestType.valueOf(t.toUpperCase(Locale.ROOT)))
+            .peek(
+                t -> {
+                  if (!SUPPORTED_TYPES.contains(t)) {
+                    throw new IllegalArgumentException(
+                        String.format(
+                            Locale.ROOT,
+                            "Request type %s is not supported for circuit breakers",
+                            t.name()));
+                  }
+                })
+            .collect(Collectors.toSet());
+  }
+
+  public Set<SolrRequestType> getRequestTypes() {
+    return requestTypes;
+  }
+
+  /**
+   * Return the proper error code to use in exception. For legacy use of {@link CircuitBreaker} we
+   * return 503 for backward compatibility, else return 429.
+   *
+   * @deprecated Remove in 10.0
+   */
+  @Deprecated(since = "9.4")
+  public static SolrException.ErrorCode getErrorCode(List<CircuitBreaker> trippedCircuitBreakers) {
+    if (trippedCircuitBreakers != null
+        && trippedCircuitBreakers.stream().anyMatch(cb -> cb instanceof CircuitBreakerManager)) {
+      return SolrException.ErrorCode.SERVICE_UNAVAILABLE;
+    } else {
+      return SolrException.ErrorCode.TOO_MANY_REQUESTS;
+    }
+  }
 }
diff --git a/solr/core/src/java/org/apache/solr/util/circuitbreaker/CircuitBreakerRegistry.java b/solr/core/src/java/org/apache/solr/util/circuitbreaker/CircuitBreakerRegistry.java
@@ -19,7 +19,11 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import org.apache.solr.client.solrj.SolrRequest.SolrRequestType;
 
 /**
  * Keeps track of all registered circuit breaker instances for various request types. Responsible
@@ -30,27 +34,37 @@
  */
 public class CircuitBreakerRegistry {
 
-  private final List<CircuitBreaker> circuitBreakerList = new ArrayList<>();
+  private final Map<SolrRequestType, List<CircuitBreaker>> circuitBreakerMap = new HashMap<>();
 
   public CircuitBreakerRegistry() {}
 
   public void register(CircuitBreaker circuitBreaker) {
-    circuitBreakerList.add(circuitBreaker);
+    circuitBreaker
+        .getRequestTypes()
+        .forEach(
+            r -> {
+              List<CircuitBreaker> list =
+                  circuitBreakerMap.computeIfAbsent(r, k -> new ArrayList<>());
+              list.add(circuitBreaker);
+            });
   }
 
   @VisibleForTesting
   public void deregisterAll() {
-    circuitBreakerList.clear();
+    circuitBreakerMap.clear();
   }
+
   /**
    * Check and return circuit breakers that have triggered
    *
+   * @param requestType {@link SolrRequestType} to check for.
    * @return CircuitBreakers which have triggered, null otherwise.
    */
-  public List<CircuitBreaker> checkTripped() {
+  public List<CircuitBreaker> checkTripped(SolrRequestType requestType) {
     List<CircuitBreaker> triggeredCircuitBreakers = null;
 
-    for (CircuitBreaker circuitBreaker : circuitBreakerList) {
+    for (CircuitBreaker circuitBreaker :
+        circuitBreakerMap.getOrDefault(requestType, Collections.emptyList())) {
       if (circuitBreaker.isTripped()) {
         if (triggeredCircuitBreakers == null) {
           triggeredCircuitBreakers = new ArrayList<>();
@@ -63,22 +77,6 @@ public List<CircuitBreaker> checkTripped() {
     return triggeredCircuitBreakers;
   }
 
-  /**
-   * Returns true if *any* circuit breaker has triggered, false if none have triggered.
-   *
-   * <p>NOTE: This method short circuits the checking of circuit breakers -- the method will return
-   * as soon as it finds a circuit breaker that has triggered.
-   */
-  public boolean checkAnyTripped() {
-    for (CircuitBreaker circuitBreaker : circuitBreakerList) {
-      if (circuitBreaker.isTripped()) {
-        return true;
-      }
-    }
-
-    return false;
-  }
-
   /**
    * Construct the final error message to be printed when circuit breakers trip.
    *
@@ -96,7 +94,7 @@ public static String toErrorMessage(List<CircuitBreaker> circuitBreakerList) {
     return sb.toString();
   }
 
-  public boolean isEnabled() {
-    return !circuitBreakerList.isEmpty();
+  public boolean isEnabled(SolrRequestType requestType) {
+    return circuitBreakerMap.containsKey(requestType);
   }
 }
diff --git a/solr/core/src/test-files/solr/collection1/conf/solrconfig-pluggable-circuitbreaker.xml b/solr/core/src/test-files/solr/collection1/conf/solrconfig-pluggable-circuitbreaker.xml
@@ -80,6 +80,18 @@
 
   <circuitBreaker class="solr.MemoryCircuitBreaker">
     <double  name="threshold">75</double>
+    <arr name="requestTypes">
+      <str>update</str>
+    </arr>
+  </circuitBreaker>
+
+  <circuitBreaker class="solr.MemoryCircuitBreaker">
+    <double  name="threshold">80</double>
+<!-- Default is query
+    <arr name="requestTypes">
+      <str>query</str>
+    </arr>
+-->
   </circuitBreaker>
 
   <circuitBreaker class="solr.CPUCircuitBreaker">

diff --git a/solr/core/src/test/org/apache/solr/util/BaseTestCircuitBreaker.java b/solr/core/src/test/org/apache/solr/util/BaseTestCircuitBreaker.java
@@ -32,6 +32,7 @@
 import org.apache.solr.common.util.SolrNamedThreadFactory;
 import org.apache.solr.util.circuitbreaker.CPUCircuitBreaker;
 import org.apache.solr.util.circuitbreaker.CircuitBreaker;
+import org.apache.solr.util.circuitbreaker.CircuitBreakerManager;
 import org.apache.solr.util.circuitbreaker.MemoryCircuitBreaker;
 import org.hamcrest.MatcherAssert;
 import org.junit.After;
@@ -79,21 +80,42 @@ public void testCBAlwaysTrips() {
         });
   }
 
-  public void testCBFakeMemoryPressure() {
+  public void testCBFakeMemoryPressure() throws Exception {
     removeAllExistingCircuitBreakers();
 
-    CircuitBreaker circuitBreaker = new FakeMemoryPressureCircuitBreaker();
-    MemoryCircuitBreaker memoryCircuitBreaker = (MemoryCircuitBreaker) circuitBreaker;
+    // Update and query will not trip
+    h.update(
+        "<add><doc><field name=\"id\">1</field><field name=\"name\">john smith</field></doc></add>");
+    h.query(req("name:\"john smith\""));
 
-    memoryCircuitBreaker.setThreshold(75);
+    MemoryCircuitBreaker searchBreaker = new FakeMemoryPressureCircuitBreaker();
+    searchBreaker.setThreshold(80);
+    // Default request type is "query"
+    // searchBreaker.setRequestTypes(List.of("query"));
+    h.getCore().getCircuitBreakerRegistry().register(searchBreaker);
 
-    h.getCore().getCircuitBreakerRegistry().register(circuitBreaker);
+    // Query will trip, but not update due to defaults
+    expectThrows(SolrException.class, () -> h.query(req("name:\"john smith\"")));
+    h.update(
+        "<add><doc><field name=\"id\">2</field><field name=\"name\">john smith</field></doc></add>");
+
+    MemoryCircuitBreaker updateBreaker = new FakeMemoryPressureCircuitBreaker();
+    updateBreaker.setThreshold(75);
+    updateBreaker.setRequestTypes(List.of("update"));
+    h.getCore().getCircuitBreakerRegistry().register(updateBreaker);
 
+    // Now also update will trip
     expectThrows(
         SolrException.class,
-        () -> {
-          h.query(req("name:\"john smith\""));
-        });
+        () ->
+            h.update(
+                "<add><doc><field name=\"id\">1</field><field name=\"name\">john smith</field></doc></add>"));
+  }
+
+  public void testBadRequestType() {
+    expectThrows(
+        IllegalArgumentException.class,
+        () -> new MemoryCircuitBreaker().setRequestTypes(List.of("badRequestType")));
   }
 
   public void testBuildingMemoryPressure() {
@@ -238,6 +260,15 @@ public void testResponseWithCBTiming() {
         "//lst[@name='process']/double[@name='time']");
   }
 
+  public void testErrorCode() {
+    assertEquals(
+        SolrException.ErrorCode.SERVICE_UNAVAILABLE,
+        CircuitBreaker.getErrorCode(List.of(new CircuitBreakerManager())));
+    assertEquals(
+        SolrException.ErrorCode.TOO_MANY_REQUESTS,
+        CircuitBreaker.getErrorCode(List.of(new MemoryCircuitBreaker())));
+  }
+
   private void removeAllExistingCircuitBreakers() {
     h.getCore().getCircuitBreakerRegistry().deregisterAll();
   }