Search in sources :

Example 1 with SelectionContext

use of io.trino.spi.resourcegroups.SelectionContext in project trino by trinodb.

the class TestDbResourceGroupConfigurationManager method testEnvironments.

@Test
public void testEnvironments() {
    H2DaoProvider daoProvider = setup("test_configuration");
    H2ResourceGroupsDao dao = daoProvider.get();
    dao.createResourceGroupsGlobalPropertiesTable();
    dao.createResourceGroupsTable();
    dao.createSelectorsTable();
    String prodEnvironment = "prod";
    String devEnvironment = "dev";
    dao.insertResourceGroupsGlobalProperties("cpu_quota_period", "1h");
    // two resource groups are the same except the group for the prod environment has a larger softMemoryLimit
    dao.insertResourceGroup(1, "prod_global", "10MB", 1000, 100, 100, "weighted", null, true, "1h", "1d", null, prodEnvironment);
    dao.insertResourceGroup(2, "dev_global", "1MB", 1000, 100, 100, "weighted", null, true, "1h", "1d", null, devEnvironment);
    dao.insertSelector(1, 1, ".*prod_user.*", null, null, null, null, null);
    dao.insertSelector(2, 2, ".*dev_user.*", null, null, null, null, null);
    // check the prod configuration
    DbResourceGroupConfigurationManager manager = new DbResourceGroupConfigurationManager(listener -> {
    }, new DbResourceGroupConfig(), daoProvider.get(), prodEnvironment);
    List<ResourceGroupSpec> groups = manager.getRootGroups();
    assertEquals(groups.size(), 1);
    InternalResourceGroup prodGlobal = new InternalResourceGroup("prod_global", (group, export) -> {
    }, directExecutor());
    manager.configure(prodGlobal, new SelectionContext<>(prodGlobal.getId(), new ResourceGroupIdTemplate("prod_global")));
    assertEqualsResourceGroup(prodGlobal, "10MB", 1000, 100, 100, WEIGHTED, DEFAULT_WEIGHT, true, Duration.ofHours(1), Duration.ofDays(1));
    assertEquals(manager.getSelectors().size(), 1);
    ResourceGroupSelector prodSelector = manager.getSelectors().get(0);
    ResourceGroupId prodResourceGroupId = prodSelector.match(new SelectionCriteria(true, "prod_user", ImmutableSet.of(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty())).get().getResourceGroupId();
    assertEquals(prodResourceGroupId.toString(), "prod_global");
    // check the dev configuration
    manager = new DbResourceGroupConfigurationManager(listener -> {
    }, new DbResourceGroupConfig(), daoProvider.get(), devEnvironment);
    assertEquals(groups.size(), 1);
    InternalResourceGroup devGlobal = new InternalResourceGroup("dev_global", (group, export) -> {
    }, directExecutor());
    manager.configure(devGlobal, new SelectionContext<>(prodGlobal.getId(), new ResourceGroupIdTemplate("dev_global")));
    assertEqualsResourceGroup(devGlobal, "1MB", 1000, 100, 100, WEIGHTED, DEFAULT_WEIGHT, true, Duration.ofHours(1), Duration.ofDays(1));
    assertEquals(manager.getSelectors().size(), 1);
    ResourceGroupSelector devSelector = manager.getSelectors().get(0);
    ResourceGroupId devResourceGroupId = devSelector.match(new SelectionCriteria(true, "dev_user", ImmutableSet.of(), Optional.empty(), ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES, Optional.empty())).get().getResourceGroupId();
    assertEquals(devResourceGroupId.toString(), "dev_global");
}
Also used : SelectionContext(io.trino.spi.resourcegroups.SelectionContext) Assertions.assertThat(org.assertj.core.api.Assertions.assertThat) SchedulingPolicy(io.trino.spi.resourcegroups.SchedulingPolicy) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) Assert.assertEquals(org.testng.Assert.assertEquals) Test(org.testng.annotations.Test) TrinoExceptionAssert.assertTrinoExceptionThrownBy(io.trino.testing.assertions.TrinoExceptionAssert.assertTrinoExceptionThrownBy) ArrayList(java.util.ArrayList) UnableToExecuteStatementException(org.jdbi.v3.core.statement.UnableToExecuteStatementException) Assertions.assertThatThrownBy(org.assertj.core.api.Assertions.assertThatThrownBy) ResourceEstimates(io.trino.spi.session.ResourceEstimates) Duration(java.time.Duration) ThreadLocalRandom(java.util.concurrent.ThreadLocalRandom) InternalResourceGroup(io.trino.execution.resourcegroups.InternalResourceGroup) WEIGHTED(io.trino.spi.resourcegroups.SchedulingPolicy.WEIGHTED) Assert.assertFalse(org.testng.Assert.assertFalse) ResourceGroupId(io.trino.spi.resourcegroups.ResourceGroupId) ImmutableSet(com.google.common.collect.ImmutableSet) ResourceGroupSelector(io.trino.plugin.resourcegroups.ResourceGroupSelector) TrinoException(io.trino.spi.TrinoException) MILLISECONDS(java.util.concurrent.TimeUnit.MILLISECONDS) ResourceGroupIdTemplate(io.trino.plugin.resourcegroups.ResourceGroupIdTemplate) DEFAULT_WEIGHT(io.trino.execution.resourcegroups.InternalResourceGroup.DEFAULT_WEIGHT) MoreExecutors.directExecutor(com.google.common.util.concurrent.MoreExecutors.directExecutor) FAIR(io.trino.spi.resourcegroups.SchedulingPolicy.FAIR) DataSize(io.airlift.units.DataSize) List(java.util.List) ResourceGroupSpec(io.trino.plugin.resourcegroups.ResourceGroupSpec) SelectionCriteria(io.trino.spi.resourcegroups.SelectionCriteria) Optional(java.util.Optional) Assert.assertTrue(org.testng.Assert.assertTrue) StaticSelector(io.trino.plugin.resourcegroups.StaticSelector) Pattern(java.util.regex.Pattern) Comparator(java.util.Comparator) ResourceGroupSelector(io.trino.plugin.resourcegroups.ResourceGroupSelector) ResourceGroupId(io.trino.spi.resourcegroups.ResourceGroupId) ResourceGroupSpec(io.trino.plugin.resourcegroups.ResourceGroupSpec) ResourceGroupIdTemplate(io.trino.plugin.resourcegroups.ResourceGroupIdTemplate) InternalResourceGroup(io.trino.execution.resourcegroups.InternalResourceGroup) SelectionCriteria(io.trino.spi.resourcegroups.SelectionCriteria) Test(org.testng.annotations.Test)

Example 2 with SelectionContext

use of io.trino.spi.resourcegroups.SelectionContext in project trino by trinodb.

the class TestDbResourceGroupConfigurationManager method testMissing.

@Test
public void testMissing() {
    H2DaoProvider daoProvider = setup("test_missing");
    H2ResourceGroupsDao dao = daoProvider.get();
    dao.createResourceGroupsGlobalPropertiesTable();
    dao.createResourceGroupsTable();
    dao.createSelectorsTable();
    dao.insertResourceGroup(1, "global", "1MB", 1000, 100, 100, "weighted", null, true, "1h", "1d", null, ENVIRONMENT);
    dao.insertResourceGroup(2, "sub", "2MB", 4, 3, 3, null, 5, null, null, null, 1L, ENVIRONMENT);
    dao.insertResourceGroupsGlobalProperties("cpu_quota_period", "1h");
    dao.insertSelector(2, 1, null, null, null, null, null, null);
    DbResourceGroupConfigurationManager manager = new DbResourceGroupConfigurationManager(listener -> {
    }, new DbResourceGroupConfig(), daoProvider.get(), ENVIRONMENT);
    InternalResourceGroup missing = new InternalResourceGroup("missing", (group, export) -> {
    }, directExecutor());
    assertThatThrownBy(() -> manager.configure(missing, new SelectionContext<>(missing.getId(), new ResourceGroupIdTemplate("missing")))).isInstanceOf(IllegalStateException.class).hasMessage("No matching configuration found for [missing] using [missing]");
}
Also used : InternalResourceGroup(io.trino.execution.resourcegroups.InternalResourceGroup) SelectionContext(io.trino.spi.resourcegroups.SelectionContext) ResourceGroupIdTemplate(io.trino.plugin.resourcegroups.ResourceGroupIdTemplate) Test(org.testng.annotations.Test)

Example 3 with SelectionContext

use of io.trino.spi.resourcegroups.SelectionContext in project trino by trinodb.

the class AbstractResourceConfigurationManager method parentGroupContext.

@Override
public SelectionContext<ResourceGroupIdTemplate> parentGroupContext(SelectionContext<ResourceGroupIdTemplate> context) {
    ResourceGroupId parentGroupId = context.getResourceGroupId().getParent().orElseThrow(() -> new IllegalArgumentException("Group has no parent group: " + context.getResourceGroupId()));
    List<ResourceGroupNameTemplate> parentGroupIdTemplate = new ArrayList<>(context.getContext().getSegments());
    parentGroupIdTemplate.remove(parentGroupIdTemplate.size() - 1);
    return new SelectionContext<>(parentGroupId, ResourceGroupIdTemplate.fromSegments(parentGroupIdTemplate));
}
Also used : ResourceGroupId(io.trino.spi.resourcegroups.ResourceGroupId) ArrayList(java.util.ArrayList) SelectionContext(io.trino.spi.resourcegroups.SelectionContext)

Example 4 with SelectionContext

use of io.trino.spi.resourcegroups.SelectionContext in project trino by trinodb.

the class StaticSelector method match.

@Override
public Optional<SelectionContext<ResourceGroupIdTemplate>> match(SelectionCriteria criteria) {
    Map<String, String> variables = new HashMap<>();
    if (userRegex.isPresent()) {
        Matcher userMatcher = userRegex.get().matcher(criteria.getUser());
        if (!userMatcher.matches()) {
            return Optional.empty();
        }
        addVariableValues(userRegex.get(), criteria.getUser(), variables);
    }
    if (userGroupRegex.isPresent() && criteria.getUserGroups().stream().noneMatch(group -> userGroupRegex.get().matcher(group).matches())) {
        return Optional.empty();
    }
    if (sourceRegex.isPresent()) {
        String source = criteria.getSource().orElse("");
        if (!sourceRegex.get().matcher(source).matches()) {
            return Optional.empty();
        }
        addVariableValues(sourceRegex.get(), source, variables);
    }
    if (!clientTags.isEmpty() && !criteria.getTags().containsAll(clientTags)) {
        return Optional.empty();
    }
    if (selectorResourceEstimate.isPresent() && !selectorResourceEstimate.get().match(criteria.getResourceEstimates())) {
        return Optional.empty();
    }
    if (queryType.isPresent()) {
        String contextQueryType = criteria.getQueryType().orElse("");
        if (!queryType.get().equalsIgnoreCase(contextQueryType)) {
            return Optional.empty();
        }
    }
    variables.putIfAbsent(USER_VARIABLE, criteria.getUser());
    // Special handling for source, which is an optional field that is part of the standard variables
    variables.putIfAbsent(SOURCE_VARIABLE, criteria.getSource().orElse(""));
    ResourceGroupId id = group.expandTemplate(new VariableMap(variables));
    return Optional.of(new SelectionContext<>(id, group));
}
Also used : ResourceGroupId(io.trino.spi.resourcegroups.ResourceGroupId) ImmutableSet(com.google.common.collect.ImmutableSet) SelectionContext(io.trino.spi.resourcegroups.SelectionContext) Set(java.util.Set) HashMap(java.util.HashMap) Sets(com.google.common.collect.Sets) HashSet(java.util.HashSet) List(java.util.List) Preconditions.checkArgument(com.google.common.base.Preconditions.checkArgument) Matcher(java.util.regex.Matcher) ImmutableList(com.google.common.collect.ImmutableList) SelectionCriteria(io.trino.spi.resourcegroups.SelectionCriteria) Map(java.util.Map) Objects.requireNonNull(java.util.Objects.requireNonNull) Optional(java.util.Optional) VisibleForTesting(com.google.common.annotations.VisibleForTesting) Pattern(java.util.regex.Pattern) ResourceGroupId(io.trino.spi.resourcegroups.ResourceGroupId) HashMap(java.util.HashMap) Matcher(java.util.regex.Matcher)

Aggregations

SelectionContext (io.trino.spi.resourcegroups.SelectionContext)4 ResourceGroupId (io.trino.spi.resourcegroups.ResourceGroupId)3 ImmutableSet (com.google.common.collect.ImmutableSet)2 InternalResourceGroup (io.trino.execution.resourcegroups.InternalResourceGroup)2 ResourceGroupIdTemplate (io.trino.plugin.resourcegroups.ResourceGroupIdTemplate)2 SelectionCriteria (io.trino.spi.resourcegroups.SelectionCriteria)2 ArrayList (java.util.ArrayList)2 List (java.util.List)2 Optional (java.util.Optional)2 Pattern (java.util.regex.Pattern)2 Test (org.testng.annotations.Test)2 VisibleForTesting (com.google.common.annotations.VisibleForTesting)1 Preconditions.checkArgument (com.google.common.base.Preconditions.checkArgument)1 ImmutableList (com.google.common.collect.ImmutableList)1 Sets (com.google.common.collect.Sets)1 MoreExecutors.directExecutor (com.google.common.util.concurrent.MoreExecutors.directExecutor)1 DataSize (io.airlift.units.DataSize)1 DEFAULT_WEIGHT (io.trino.execution.resourcegroups.InternalResourceGroup.DEFAULT_WEIGHT)1 ResourceGroupSelector (io.trino.plugin.resourcegroups.ResourceGroupSelector)1 ResourceGroupSpec (io.trino.plugin.resourcegroups.ResourceGroupSpec)1