From 5278d99f9d05bd66cf8e3a07af0df4982c84976e Mon Sep 17 00:00:00 2001
From: Alexander Schwartz <aschwart@redhat.com>
Date: Tue, 25 Jul 2023 16:49:19 +0200
Subject: [PATCH] Provide an HTTP 503 response when the thread pool rejected
 the request

This allows for limiting the queue of the thread pool to avoid overloading Keycloak.

Closes #21962
---
 .../integration/web/QuarkusRequestFilter.java | 35 ++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/integration/web/QuarkusRequestFilter.java b/quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/integration/web/QuarkusRequestFilter.java
index b3eead824fce..e91af33ba5fa 100644
--- a/quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/integration/web/QuarkusRequestFilter.java
+++ b/quarkus/runtime/src/main/java/org/keycloak/quarkus/runtime/integration/web/QuarkusRequestFilter.java
@@ -18,7 +18,12 @@
 package org.keycloak.quarkus.runtime.integration.web;
 
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Predicate;
+
+import org.apache.http.HttpStatus;
+import org.jboss.logging.Logger;
 import org.keycloak.common.ClientConnection;
 import org.keycloak.common.util.Resteasy;
 import org.keycloak.models.KeycloakSession;
@@ -46,6 +51,8 @@
  */
 public class QuarkusRequestFilter implements Handler<RoutingContext>, TransactionalSessionHandler {
 
+    private final Logger logger = Logger.getLogger(QuarkusRequestFilter.class);
+
     private final ExecutorService executor;
 
     private Predicate<RoutingContext> contextFilter;
@@ -59,6 +66,9 @@ public QuarkusRequestFilter(Predicate<RoutingContext> contextFilter, ExecutorSer
         this.executor = executor;
     }
 
+    private final LongAdder rejectedRequests = new LongAdder();
+    private volatile boolean loadSheddingActive;
+
     @Override
     public void handle(RoutingContext context) {
         if (ignoreContext(context)) {
@@ -67,7 +77,30 @@ public void handle(RoutingContext context) {
         }
         // our code should always be run as blocking until we don't provide a better support for running non-blocking code
         // in the event loop
-        executor.execute(createBlockingHandler(context));
+        try {
+            executor.execute(createBlockingHandler(context));
+            if (loadSheddingActive) {
+                synchronized (rejectedRequests) {
+                    if (loadSheddingActive) {
+                        loadSheddingActive = false;
+                        // rejectedRequests.sumThenReset() is approximative when concurrent increments are active, still it should be accurate enough for this log message
+                        logger.warnf("Executor thread pool no longer exhausted, request processing continues after %s discarded request(s)", rejectedRequests.sumThenReset());
+                    }
+                }
+            }
+        } catch (RejectedExecutionException e) {
+            if (!loadSheddingActive) {
+                synchronized (rejectedRequests) {
+                    if (!loadSheddingActive) {
+                        loadSheddingActive = true;
+                        logger.warn("Executor thread pool exhausted, starting to reject requests");
+                    }
+                }
+            }
+            rejectedRequests.increment();
+            // if the thread pool has been configured with a maximum queue size, it might reject the request
+            context.fail(HttpStatus.SC_SERVICE_UNAVAILABLE);
+        }
     }
 
     private boolean ignoreContext(RoutingContext context) {