make rate limiting return an error

confluentinc · Feb 25, 2022 · 67b1011 · 67b1011
1 parent fba9fe4
commit 67b1011
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/computation/CommandRunner.java b/ksqldb-rest-app/src/main/java/io/confluent/ksql/rest/server/computation/CommandRunner.java
@@ -21,11 +21,13 @@
 import io.confluent.ksql.rest.Errors;
 import io.confluent.ksql.rest.entity.ClusterTerminateRequest;
 import io.confluent.ksql.rest.server.resources.IncompatibleKsqlCommandVersionException;
+import io.confluent.ksql.rest.server.resources.KsqlRestException;
 import io.confluent.ksql.rest.server.state.ServerState;
 import io.confluent.ksql.rest.util.ClusterTerminator;
 import io.confluent.ksql.rest.util.PersistentQueryCleanupImpl;
 import io.confluent.ksql.rest.util.TerminateCluster;
 import io.confluent.ksql.services.KafkaTopicClient;
+import io.confluent.ksql.util.KsqlException;
 import io.confluent.ksql.util.Pair;
 import io.confluent.ksql.util.PersistentQueryMetadata;
 import io.confluent.ksql.util.QueryMetadata;
@@ -352,7 +354,11 @@ private void executeStatement(final QueuedCommand queuedCommand) {
       if (closed) {
         LOG.info("Execution aborted as system is closing down");
       } else {
-        rateLimiter.acquire();
+        if (!rateLimiter.tryAcquire()) {
+          throw new KsqlRestException(
+            Errors.tooManyRequests("Too many requests to the command topic within a 1 second timeframe")
+          );
+        }
         statementExecutor.handleStatement(queuedCommand);
         LOG.info("Executed statement: " + commandStatement);
       }

diff --git a/ksqldb-rest-model/src/main/java/io/confluent/ksql/rest/Errors.java b/ksqldb-rest-model/src/main/java/io/confluent/ksql/rest/Errors.java
@@ -23,6 +23,7 @@
 import static io.netty.handler.codec.http.HttpResponseStatus.NOT_FOUND;
 import static io.netty.handler.codec.http.HttpResponseStatus.PRECONDITION_REQUIRED;
 import static io.netty.handler.codec.http.HttpResponseStatus.SERVICE_UNAVAILABLE;
+import static io.netty.handler.codec.http.HttpResponseStatus.TOO_MANY_REQUESTS;
 import static io.netty.handler.codec.http.HttpResponseStatus.UNAUTHORIZED;
 
 import io.confluent.ksql.rest.entity.KsqlEntityList;
@@ -72,6 +73,8 @@ public final class Errors {
   public static final int ERROR_CODE_SERVER_ERROR =
       toErrorCode(INTERNAL_SERVER_ERROR.code());
 
+  public static final int ERROR_CODE_TOO_MANY_REQUESTS = toErrorCode(TOO_MANY_REQUESTS.code());
+
   private final ErrorMessages errorMessages;
 
   public static int toStatusCode(final int errorCode) {
@@ -208,6 +211,13 @@ public static EndpointResponse serverNotReady(final KsqlErrorMessage error) {
         .build();
   }
 
+  public static EndpointResponse tooManyRequests(final String msg) {
+    return EndpointResponse.create()
+      .status(TOO_MANY_REQUESTS.code())
+      .entity(new KsqlErrorMessage(ERROR_CODE_TOO_MANY_REQUESTS, msg))
+      .build();
+  }
+
   public Errors(final ErrorMessages errorMessages) {
     this.errorMessages = Objects.requireNonNull(errorMessages, "errorMessages");
   }