Rate limiter (apache#278)

polaris-server.yml

@@ -166,3 +166,7 @@ logging:
 
 # Limits the size of request bodies sent to Polaris. -1 means no limit.
 maxRequestBodyBytes: -1
+
+# Optional, not specifying a "rateLimiter" section also means no rate limiter
+rateLimiter:
+ type: no-op

polaris-service/src/main/java/org/apache/polaris/service/PolarisApplication.java

@@ -101,6 +101,7 @@
 import org.apache.polaris.service.context.PolarisCallContextCatalogFactory;
 import org.apache.polaris.service.context.RealmContextResolver;
 import org.apache.polaris.service.persistence.InMemoryPolarisMetaStoreManagerFactory;
+import org.apache.polaris.service.ratelimiter.RateLimiterFilter;
 import org.apache.polaris.service.storage.PolarisStorageIntegrationProviderImpl;
 import org.apache.polaris.service.task.ManifestFileCleanupTaskHandler;
 import org.apache.polaris.service.task.TableCleanupTaskHandler;
@@ -252,6 +253,14 @@ public void run(PolarisApplicationConfig configuration, Environment environment)
  .servlets()
  .addFilter("tracing", new TracingFilter(openTelemetry))
  .addMappingForUrlPatterns(EnumSet.of(DispatcherType.REQUEST), true, "/*");
+
+ if (configuration.getRateLimiter() != null) {
+ environment
+ .servlets()
+ .addFilter("ratelimiter", new RateLimiterFilter(configuration.getRateLimiter()))
+ .addMappingForUrlPatterns(EnumSet.of(DispatcherType.REQUEST), true, "/*");
+ }
+
  DiscoverableAuthenticator<String, AuthenticatedPolarisPrincipal> authenticator =
  configuration.getPolarisAuthenticator();
  authenticator.setEntityManagerFactory(entityManagerFactory);

polaris-service/src/main/java/org/apache/polaris/service/config/PolarisApplicationConfig.java

@@ -24,6 +24,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import javax.annotation.Nullable;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.polaris.core.PolarisConfigurationStore;
 import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal;
@@ -32,6 +33,7 @@
 import org.apache.polaris.service.catalog.FileIOFactory;
 import org.apache.polaris.service.context.CallContextResolver;
 import org.apache.polaris.service.context.RealmContextResolver;
+import org.apache.polaris.service.ratelimiter.RateLimiter;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
@@ -55,6 +57,7 @@ public class PolarisApplicationConfig extends Configuration {
  private String awsAccessKey;
  private String awsSecretKey;
  private FileIOFactory fileIOFactory;
+ private RateLimiter rateLimiter;
 
  public static final long REQUEST_BODY_BYTES_NO_LIMIT = -1;
  private long maxRequestBodyBytes = REQUEST_BODY_BYTES_NO_LIMIT;
@@ -137,6 +140,16 @@ public void setCorsConfiguration(CorsConfiguration corsConfiguration) {
  this.corsConfiguration = corsConfiguration;
  }
 
+ @JsonProperty("rateLimiter")
+ public RateLimiter getRateLimiter() {
+ return rateLimiter;
+ }
+
+ @JsonProperty("rateLimiter")
+ public void setRateLimiter(@Nullable RateLimiter rateLimiter) {
+ this.rateLimiter = rateLimiter;
+ }
+
  public void setTaskHandler(TaskHandlerConfiguration taskHandler) {
  this.taskHandler = taskHandler;
  }

polaris-service/src/main/java/org/apache/polaris/service/ratelimiter/NoOpRateLimiter.java

@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.polaris.service.ratelimiter;
+
+import com.fasterxml.jackson.annotation.JsonTypeName;
+
+/** Rate limiter that always allows the request */
+@JsonTypeName("no-op")
+public class NoOpRateLimiter implements RateLimiter {
+ @Override
+ public boolean tryAcquire() {
+ return true;
+ }
+}

polaris-service/src/main/java/org/apache/polaris/service/ratelimiter/RateLimiter.java

@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.polaris.service.ratelimiter;
+
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import io.dropwizard.jackson.Discoverable;
+
+/** Interface for rate limiting requests */
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, include = JsonTypeInfo.As.PROPERTY, property = "type")
+public interface RateLimiter extends Discoverable {
+ /**
+ * This signifies that a request is being made. That is, the rate limiter should count the request
+ * at this point.
+ *
+ * @return Whether the request is allowed to proceed by the rate limiter
+ */
+ boolean tryAcquire();
+}

polaris-service/src/main/java/org/apache/polaris/service/ratelimiter/RateLimiterFilter.java

@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.polaris.service.ratelimiter;
+
+import jakarta.annotation.Priority;
+import jakarta.servlet.Filter;
+import jakarta.servlet.FilterChain;
+import jakarta.servlet.ServletException;
+import jakarta.servlet.ServletRequest;
+import jakarta.servlet.ServletResponse;
+import jakarta.servlet.http.HttpServletResponse;
+import jakarta.ws.rs.Priorities;
+import jakarta.ws.rs.core.Response;
+import java.io.IOException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Request filter that returns a 429 Too Many Requests if the rate limiter says so */
+@Priority(Priorities.AUTHORIZATION + 1)
+public class RateLimiterFilter implements Filter {
+ private static final Logger LOGGER = LoggerFactory.getLogger(RateLimiterFilter.class);
+
+ private final RateLimiter rateLimiter;
+
+ public RateLimiterFilter(RateLimiter rateLimiter) {
+ this.rateLimiter = rateLimiter;
+ }
+
+ /** Returns a 429 if the rate limiter says so. Otherwise, forwards the request along. */
+ @Override
+ public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
+ throws IOException, ServletException {
+ if (response instanceof HttpServletResponse servletResponse && !rateLimiter.tryAcquire()) {
+ servletResponse.setStatus(Response.Status.TOO_MANY_REQUESTS.getStatusCode());
+ LOGGER.atDebug().log("Rate limiting request");
+ return;
+ }
+ chain.doFilter(request, response);
+ }
+}

polaris-service/src/main/java/org/apache/polaris/service/ratelimiter/RealmTokenBucketRateLimiter.java

@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.polaris.service.ratelimiter;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonTypeName;
+import java.time.Clock;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import org.apache.polaris.core.context.CallContext;
+import org.apache.polaris.core.context.RealmContext;
+import org.jetbrains.annotations.VisibleForTesting;
+
+/**
+ * Rate limiter that maps the request's realm identifier to its own TokenBucketRateLimiter, with its
+ * own capacity.
+ */
+@JsonTypeName("realm-token-bucket")
+public class RealmTokenBucketRateLimiter implements RateLimiter {
+ private final long requestsPerSecond;
+ private final long windowSeconds;
+ private final Map<String, RateLimiter> perRealmLimiters;
+
+ @VisibleForTesting
+ @JsonCreator
+ public RealmTokenBucketRateLimiter(
+ @JsonProperty("requestsPerSecond") final long requestsPerSecond,
+ @JsonProperty("windowSeconds") final long windowSeconds) {
+ this.requestsPerSecond = requestsPerSecond;
+ this.windowSeconds = windowSeconds;
+ this.perRealmLimiters = new ConcurrentHashMap<>();
+ }
+
+ /**
+ * This signifies that a request is being made. That is, the rate limiter should count the request
+ * at this point.
+ *
+ * @return Whether the request is allowed to proceed by the rate limiter
+ */
+ @Override
+ public boolean tryAcquire() {
+ String key =
+ Optional.ofNullable(CallContext.getCurrentContext())
+ .map(CallContext::getRealmContext)
+ .map(RealmContext::getRealmIdentifier)
+ .orElse("");
+
+ return perRealmLimiters
+ .computeIfAbsent(
+ key,
+ (k) ->
+ new TokenBucketRateLimiter(
+ requestsPerSecond,
+ Math.multiplyExact(requestsPerSecond, windowSeconds),
+ getClock()))
+ .tryAcquire();
+ }
+
+ @VisibleForTesting
+ protected Clock getClock() {
+ return Clock.systemUTC();
+ }
+}

polaris-service/src/main/java/org/apache/polaris/service/ratelimiter/TokenBucketRateLimiter.java

@@ -0,0 +1,64 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.polaris.service.ratelimiter;
+
+import java.time.InstantSource;
+
+/**
+ * Token bucket implementation of a Polaris RateLimiter. Acquires tokens at a fixed rate and has a
+ * maximum amount of tokens. Each successful "tryAcquire" costs 1 token.
+ */
+public class TokenBucketRateLimiter implements RateLimiter {
+ private final double tokensPerMilli;
+ private final long maxTokens;
+ private final InstantSource instantSource;
+
+ private double tokens;
+ private long lastTokenGenerationMillis;
+
+ public TokenBucketRateLimiter(long tokensPerSecond, long maxTokens, InstantSource instantSource) {
+ this.tokensPerMilli = tokensPerSecond / 1000D;
+ this.maxTokens = maxTokens;
+ this.instantSource = instantSource;
+
+ tokens = maxTokens;
+ lastTokenGenerationMillis = instantSource.millis();
+ }
+
+ /**
+ * Tries to acquire and spend 1 token. Doesn't block if a token isn't available.
+ *
+ * @return whether a token was successfully acquired & spent
+ */
+ @Override
+ public synchronized boolean tryAcquire() {
+ // Grant tokens for the time that has passed since our last tryAcquire()
+ long t = instantSource.millis();
+ long millisPassed = Math.subtractExact(t, lastTokenGenerationMillis);
+ lastTokenGenerationMillis = t;
+ tokens = Math.min(maxTokens, tokens + (millisPassed * tokensPerMilli));
+
+ // Take a token if they have one available
+ if (tokens >= 1) {
+ tokens--;
+ return true;
+ }
+ return false;
+ }
+}

polaris-service/src/main/resources/META-INF/services/io.dropwizard.jackson.Discoverable

@@ -24,3 +24,4 @@ org.apache.polaris.service.context.RealmContextResolver
 org.apache.polaris.service.context.CallContextResolver
 org.apache.polaris.service.auth.TokenBrokerFactory
 org.apache.polaris.service.catalog.FileIOFactory
+org.apache.polaris.service.ratelimiter.RateLimiter

polaris-service/src/main/resources/META-INF/services/org.apache.polaris.service.ratelimiter.RateLimiter

@@ -0,0 +1,21 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+org.apache.polaris.service.ratelimiter.RealmTokenBucketRateLimiter
+org.apache.polaris.service.ratelimiter.NoOpRateLimiter