diff --git a/docs/src/main/sphinx/security/opa-access-control.md b/docs/src/main/sphinx/security/opa-access-control.md index 6e284227ac87..cb19bf085ebe 100644 --- a/docs/src/main/sphinx/security/opa-access-control.md +++ b/docs/src/main/sphinx/security/opa-access-control.md @@ -61,6 +61,9 @@ The following table lists the configuration properties for the OPA access contro * - `opa.allow-permission-management-operations` - Configure if permission management operations are allowed. Find more details in [](opa-permission-management). Defaults to `false`. +* - `opa.max-outstanding-requests` + - Maximum number of requests awaiting responses from OPA across all queries. + Defaults to `512`. * - `opa.http-client.*` - Optional HTTP client configurations for the connection from Trino to OPA, for example `opa.http-client.http-proxy` for configuring the HTTP proxy. diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaConfig.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaConfig.java index 219265a18985..9f2c13abfbef 100644 --- a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaConfig.java +++ b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaConfig.java @@ -16,6 +16,7 @@ import io.airlift.configuration.Config; import io.airlift.configuration.ConfigDescription; import io.airlift.configuration.validation.FileExists; +import jakarta.validation.constraints.Min; import jakarta.validation.constraints.NotNull; import java.net.URI; @@ -34,6 +35,7 @@ public class OpaConfig private Optional opaColumnMaskingUri = Optional.empty(); private Optional opaBatchColumnMaskingUri = Optional.empty(); private Optional additionalContextFile = Optional.empty(); + private int maxOutstandingRequests = 512; @NotNull public URI getOpaUri() @@ -156,4 +158,18 @@ public OpaConfig setAdditionalContextFile(Path additionalContextFile) this.additionalContextFile = Optional.ofNullable(additionalContextFile); return this; } + + @Min(1) + public int getMaxOutstandingRequests() + { + return maxOutstandingRequests; + } + + @Config("opa.max-outstanding-requests") + @ConfigDescription("Maximum number of OPA requests awaiting responses") + public OpaConfig setMaxOutstandingRequests(int maxOutstandingRequests) + { + this.maxOutstandingRequests = maxOutstandingRequests; + return this; + } } diff --git a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaHttpClient.java b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaHttpClient.java index ac3880001047..1cdd99022093 100644 --- a/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaHttpClient.java +++ b/plugin/trino-opa/src/main/java/io/trino/plugin/opa/OpaHttpClient.java @@ -44,6 +44,7 @@ import java.util.Set; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.Semaphore; import java.util.function.BiFunction; import java.util.function.Function; @@ -51,6 +52,7 @@ import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.collect.ImmutableSet.toImmutableSet; import static com.google.common.net.MediaType.JSON_UTF_8; +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static io.airlift.http.client.FullJsonResponseHandler.createFullJsonResponseHandler; import static io.airlift.http.client.HeaderNames.CONTENT_TYPE; import static io.airlift.http.client.JsonBodyGenerator.jsonBodyGenerator; @@ -66,6 +68,7 @@ public class OpaHttpClient private final Executor executor; private final boolean logRequests; private final boolean logResponses; + private final Semaphore outstandingRequests; private static final Logger log = Logger.get(OpaHttpClient.class); @Inject @@ -80,6 +83,7 @@ public OpaHttpClient( this.executor = requireNonNull(executor, "executor is null"); this.logRequests = config.getLogRequests(); this.logResponses = config.getLogResponses(); + this.outstandingRequests = new Semaphore(config.getMaxOutstandingRequests(), true); } public FluentFuture submitOpaRequest(OpaQueryInput input, URI uri, JsonCodec deserializer) @@ -105,7 +109,24 @@ public FluentFuture submitOpaRequest(OpaQueryInput input, URI uri, JsonCo new String(requestBodyGenerator.getBody(), UTF_8), request.getHeaders()); } - return FluentFuture.from(httpClient.executeAsync(request, createFullJsonResponseHandler(deserializer))) + try { + outstandingRequests.acquire(); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new OpaQueryException.QueryFailed(e); + } + + ListenableFuture> responseFuture; + try { + responseFuture = httpClient.executeAsync(request, createFullJsonResponseHandler(deserializer)); + } + catch (RuntimeException e) { + outstandingRequests.release(); + throw e; + } + responseFuture.addListener(outstandingRequests::release, directExecutor()); + return FluentFuture.from(responseFuture) .transform(response -> parseOpaResponse(response, uri), executor); } diff --git a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/HttpClientUtils.java b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/HttpClientUtils.java index 5eee19a7efa5..3ea34df9ef56 100644 --- a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/HttpClientUtils.java +++ b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/HttpClientUtils.java @@ -30,6 +30,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.concurrent.ExecutorService; import java.util.function.Function; import static com.google.common.net.MediaType.JSON_UTF_8; @@ -103,12 +104,23 @@ public InstrumentedHttpClient(URI expectedURI, String expectedMethod, String exp this(new RecordingHttpProcessor(expectedURI, expectedMethod, expectedContentType, handler)); } + public InstrumentedHttpClient(URI expectedURI, String expectedMethod, String expectedContentType, Function handler, ExecutorService executor) + { + this(new RecordingHttpProcessor(expectedURI, expectedMethod, expectedContentType, handler), executor); + } + public InstrumentedHttpClient(RecordingHttpProcessor processor) { super(processor); this.httpProcessor = processor; } + public InstrumentedHttpClient(RecordingHttpProcessor processor, ExecutorService executor) + { + super(processor, executor); + this.httpProcessor = processor; + } + public List getRequests() { return httpProcessor.getRequests(); diff --git a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestHelpers.java b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestHelpers.java index 7de1b1649b7f..6fdf64b95be7 100644 --- a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestHelpers.java +++ b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestHelpers.java @@ -23,6 +23,7 @@ import java.net.URI; import java.util.Map; import java.util.Optional; +import java.util.concurrent.ExecutorService; import java.util.function.Consumer; import java.util.function.Function; @@ -45,6 +46,11 @@ public static InstrumentedHttpClient createMockHttpClient(URI expectedUri, Funct return new InstrumentedHttpClient(expectedUri, "POST", JSON_UTF_8.toString(), handler); } + public static InstrumentedHttpClient createMockHttpClient(URI expectedUri, Function handler, ExecutorService executor) + { + return new InstrumentedHttpClient(expectedUri, "POST", JSON_UTF_8.toString(), handler, executor); + } + public static OpaAccessControl createOpaAuthorizer(OpaConfig config, InstrumentedHttpClient mockHttpClient) { return (OpaAccessControl) OpaAccessControlFactory.create(opaConfigToDict(config), Optional.of(mockHttpClient), Optional.of(SYSTEM_ACCESS_CONTROL_CONTEXT)); @@ -87,7 +93,8 @@ public static Map opaConfigToDict(OpaConfig config) .put("opa.policy.uri", config.getOpaUri().toString()) .put("opa.log-requests", String.valueOf(config.getLogRequests())) .put("opa.log-responses", String.valueOf(config.getLogResponses())) - .put("opa.allow-permission-management-operations", String.valueOf(config.getAllowPermissionManagementOperations())); + .put("opa.allow-permission-management-operations", String.valueOf(config.getAllowPermissionManagementOperations())) + .put("opa.max-outstanding-requests", String.valueOf(config.getMaxOutstandingRequests())); config.getOpaBatchUri().ifPresent(batchUri -> configBuilder.put("opa.policy.batched-uri", batchUri.toString())); config.getOpaRowFiltersUri().ifPresent(rowFiltersUri -> configBuilder.put("opa.policy.row-filters-uri", rowFiltersUri.toString())); config.getOpaColumnMaskingUri().ifPresent(columnMaskingUri -> configBuilder.put("opa.policy.column-masking-uri", columnMaskingUri.toString())); diff --git a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaBatchAccessControlFiltering.java b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaBatchAccessControlFiltering.java index 6d29b173094f..b230dd8dd2f7 100644 --- a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaBatchAccessControlFiltering.java +++ b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaBatchAccessControlFiltering.java @@ -29,11 +29,16 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import java.util.function.Function; import java.util.stream.Stream; import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.util.concurrent.Uninterruptibles.awaitUninterruptibly; import static io.trino.plugin.opa.RequestTestUtilities.assertStringRequestsEqual; import static io.trino.plugin.opa.RequestTestUtilities.buildValidatingRequestHandler; import static io.trino.plugin.opa.TestConstants.OK_RESPONSE; @@ -45,6 +50,10 @@ import static io.trino.plugin.opa.TestHelpers.assertAccessControlMethodThrowsForResponse; import static io.trino.plugin.opa.TestHelpers.createMockHttpClient; import static io.trino.plugin.opa.TestHelpers.createOpaAuthorizer; +import static java.util.concurrent.Executors.newFixedThreadPool; +import static java.util.concurrent.Executors.newSingleThreadExecutor; +import static java.util.concurrent.TimeUnit.MILLISECONDS; +import static java.util.concurrent.TimeUnit.SECONDS; import static org.assertj.core.api.Assertions.assertThat; final class TestOpaBatchAccessControlFiltering @@ -255,6 +264,63 @@ void testFilterColumns() .buildOrThrow()); } + @Test + void testFilterColumnsLimitsOutstandingRequests() + throws Exception + { + SchemaTableName tableOne = SchemaTableName.schemaTableName("my_schema", "table_one"); + SchemaTableName tableTwo = SchemaTableName.schemaTableName("my_schema", "table_two"); + SchemaTableName tableThree = SchemaTableName.schemaTableName("my_schema", "table_three"); + Map> requestedColumns = ImmutableMap.of( + tableOne, ImmutableSet.of("column"), + tableTwo, ImmutableSet.of("column"), + tableThree, ImmutableSet.of("column")); + CountDownLatch twoRequestsStarted = new CountDownLatch(2); + CountDownLatch threeRequestsStarted = new CountDownLatch(3); + CountDownLatch releaseResponses = new CountDownLatch(1); + AtomicInteger activeRequests = new AtomicInteger(); + AtomicInteger maximumActiveRequests = new AtomicInteger(); + ExecutorService httpExecutor = newFixedThreadPool(3); + InstrumentedHttpClient mockClient = createMockHttpClient( + OPA_SERVER_BATCH_URI, + _ -> { + int currentActiveRequests = activeRequests.incrementAndGet(); + maximumActiveRequests.accumulateAndGet(currentActiveRequests, Math::max); + twoRequestsStarted.countDown(); + threeRequestsStarted.countDown(); + try { + assertThat(awaitUninterruptibly(releaseResponses, 10, SECONDS)).isTrue(); + return new MockResponse("{\"result\": [0]}", 200); + } + finally { + activeRequests.decrementAndGet(); + } + }, + httpExecutor); + OpaAccessControl authorizer = createOpaAuthorizer( + batchFilteringOpaConfig().setMaxOutstandingRequests(2), + mockClient); + ExecutorService executor = newSingleThreadExecutor(); + try { + Future>> result = executor.submit(() -> authorizer.filterColumns( + TEST_SECURITY_CONTEXT, + "my_catalog", + requestedColumns)); + + assertThat(twoRequestsStarted.await(10, SECONDS)).isTrue(); + assertThat(threeRequestsStarted.await(200, MILLISECONDS)).isFalse(); + releaseResponses.countDown(); + + assertThat(result.get(10, SECONDS)).containsExactlyInAnyOrderEntriesOf(requestedColumns); + assertThat(maximumActiveRequests).hasValue(2); + } + finally { + releaseResponses.countDown(); + executor.shutdownNow(); + httpExecutor.shutdownNow(); + } + } + @Test void testEmptyFilterColumns() { diff --git a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaConfig.java b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaConfig.java index 8f823fab89a7..890a3168e3c5 100644 --- a/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaConfig.java +++ b/plugin/trino-opa/src/test/java/io/trino/plugin/opa/TestOpaConfig.java @@ -38,7 +38,8 @@ void testDefaults() .setLogRequests(false) .setLogResponses(false) .setAllowPermissionManagementOperations(false) - .setAdditionalContextFile(null)); + .setAdditionalContextFile(null) + .setMaxOutstandingRequests(512)); } @Test @@ -54,6 +55,7 @@ void testExplicitPropertyMappings() .put("opa.log-responses", "true") .put("opa.allow-permission-management-operations", "true") .put("opa.context-file", "src/test/resources/additional-context.properties") + .put("opa.max-outstanding-requests", "37") .buildOrThrow(); OpaConfig expected = new OpaConfig() @@ -65,7 +67,8 @@ void testExplicitPropertyMappings() .setLogRequests(true) .setLogResponses(true) .setAllowPermissionManagementOperations(true) - .setAdditionalContextFile(Path.of("src/test/resources/additional-context.properties")); + .setAdditionalContextFile(Path.of("src/test/resources/additional-context.properties")) + .setMaxOutstandingRequests(37); assertFullMapping(properties, expected); }