Examples with ResourceEstimates io.trino.spi.session.ResourceEstimates used on opensource projects

Example 1 with ResourceEstimates

use of io.trino.spi.session.ResourceEstimates in project trino by trinodb.

the class HttpRequestSessionContextFactory method createSessionContext.

public SessionContext createSessionContext(MultivaluedMap<String, String> headers, Optional<String> alternateHeaderName, Optional<String> remoteAddress, Optional<Identity> authenticatedIdentity) throws WebApplicationException {
    ProtocolHeaders protocolHeaders;
    try {
        protocolHeaders = detectProtocol(alternateHeaderName, headers.keySet());
    } catch (ProtocolDetectionException e) {
        throw badRequest(e.getMessage());
    }
    Optional<String> catalog = Optional.ofNullable(trimEmptyToNull(headers.getFirst(protocolHeaders.requestCatalog())));
    Optional<String> schema = Optional.ofNullable(trimEmptyToNull(headers.getFirst(protocolHeaders.requestSchema())));
    Optional<String> path = Optional.ofNullable(trimEmptyToNull(headers.getFirst(protocolHeaders.requestPath())));
    assertRequest((catalog.isPresent()) || (schema.isEmpty()), "Schema is set but catalog is not");
    requireNonNull(authenticatedIdentity, "authenticatedIdentity is null");
    Identity identity = buildSessionIdentity(authenticatedIdentity, protocolHeaders, headers);
    SelectedRole selectedRole = parseSystemRoleHeaders(protocolHeaders, headers);
    Optional<String> source = Optional.ofNullable(headers.getFirst(protocolHeaders.requestSource()));
    Optional<String> traceToken = Optional.ofNullable(trimEmptyToNull(headers.getFirst(protocolHeaders.requestTraceToken())));
    Optional<String> userAgent = Optional.ofNullable(headers.getFirst(USER_AGENT));
    Optional<String> remoteUserAddress = requireNonNull(remoteAddress, "remoteAddress is null");
    Optional<String> timeZoneId = Optional.ofNullable(headers.getFirst(protocolHeaders.requestTimeZone()));
    Optional<String> language = Optional.ofNullable(headers.getFirst(protocolHeaders.requestLanguage()));
    Optional<String> clientInfo = Optional.ofNullable(headers.getFirst(protocolHeaders.requestClientInfo()));
    Set<String> clientTags = parseClientTags(protocolHeaders, headers);
    Set<String> clientCapabilities = parseClientCapabilities(protocolHeaders, headers);
    ResourceEstimates resourceEstimates = parseResourceEstimate(protocolHeaders, headers);
    // parse session properties
    ImmutableMap.Builder<String, String> systemProperties = ImmutableMap.builder();
    Map<String, Map<String, String>> catalogSessionProperties = new HashMap<>();
    for (Entry<String, String> entry : parseSessionHeaders(protocolHeaders, headers).entrySet()) {
        String fullPropertyName = entry.getKey();
        String propertyValue = entry.getValue();
        List<String> nameParts = DOT_SPLITTER.splitToList(fullPropertyName);
        if (nameParts.size() == 1) {
            String propertyName = nameParts.get(0);
            assertRequest(!propertyName.isEmpty(), "Invalid %s header", protocolHeaders.requestSession());
            // catalog session properties cannot be validated until the transaction has stated, so we delay system property validation also
            systemProperties.put(propertyName, propertyValue);
        } else if (nameParts.size() == 2) {
            String catalogName = nameParts.get(0);
            String propertyName = nameParts.get(1);
            assertRequest(!catalogName.isEmpty(), "Invalid %s header", protocolHeaders.requestSession());
            assertRequest(!propertyName.isEmpty(), "Invalid %s header", protocolHeaders.requestSession());
            // catalog session properties cannot be validated until the transaction has stated
            catalogSessionProperties.computeIfAbsent(catalogName, id -> new HashMap<>()).put(propertyName, propertyValue);
        } else {
            throw badRequest(format("Invalid %s header", protocolHeaders.requestSession()));
        }
    }
    requireNonNull(catalogSessionProperties, "catalogSessionProperties is null");
    catalogSessionProperties = catalogSessionProperties.entrySet().stream().collect(toImmutableMap(Entry::getKey, entry -> ImmutableMap.copyOf(entry.getValue())));
    Map<String, String> preparedStatements = parsePreparedStatementsHeaders(protocolHeaders, headers);
    String transactionIdHeader = headers.getFirst(protocolHeaders.requestTransactionId());
    boolean clientTransactionSupport = transactionIdHeader != null;
    Optional<TransactionId> transactionId = parseTransactionId(transactionIdHeader);
    return new SessionContext(protocolHeaders, catalog, schema, path, authenticatedIdentity, identity, selectedRole, source, traceToken, userAgent, remoteUserAddress, timeZoneId, language, clientTags, clientCapabilities, resourceEstimates, systemProperties.buildOrThrow(), catalogSessionProperties, preparedStatements, transactionId, clientTransactionSupport, clientInfo);
}

Example 2 with ResourceEstimates

use of io.trino.spi.session.ResourceEstimates in project trino by trinodb.

the class TestQueues method testSelectorResourceEstimateBasedSelection.

@Test(timeOut = 240_000)
public void testSelectorResourceEstimateBasedSelection() throws Exception {
    try (DistributedQueryRunner queryRunner = createQueryRunner()) {
        queryRunner.installPlugin(new ResourceGroupManagerPlugin());
        queryRunner.getCoordinator().getResourceGroupManager().get().setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_resource_estimate_based_config.json")));
        assertResourceGroup(queryRunner, newSessionWithResourceEstimates(new ResourceEstimates(Optional.of(java.time.Duration.ofMinutes(4)), Optional.empty(), Optional.of(DataSize.of(400, MEGABYTE).toBytes()))), LONG_LASTING_QUERY, createResourceGroupId("global", "small"));
        assertResourceGroup(queryRunner, newSessionWithResourceEstimates(new ResourceEstimates(Optional.of(java.time.Duration.ofMinutes(4)), Optional.empty(), Optional.of(DataSize.of(600, MEGABYTE).toBytes()))), LONG_LASTING_QUERY, createResourceGroupId("global", "other"));
        assertResourceGroup(queryRunner, newSessionWithResourceEstimates(new ResourceEstimates(Optional.of(java.time.Duration.ofMinutes(4)), Optional.empty(), Optional.empty())), LONG_LASTING_QUERY, createResourceGroupId("global", "other"));
        assertResourceGroup(queryRunner, newSessionWithResourceEstimates(new ResourceEstimates(Optional.of(java.time.Duration.ofSeconds(1)), Optional.of(java.time.Duration.ofSeconds(1)), Optional.of(DataSize.of(6, TERABYTE).toBytes()))), LONG_LASTING_QUERY, createResourceGroupId("global", "huge_memory"));
        assertResourceGroup(queryRunner, newSessionWithResourceEstimates(new ResourceEstimates(Optional.of(java.time.Duration.ofHours(100)), Optional.empty(), Optional.of(DataSize.of(4, TERABYTE).toBytes()))), LONG_LASTING_QUERY, createResourceGroupId("global", "other"));
    }
}

Example 3 with ResourceEstimates

use of io.trino.spi.session.ResourceEstimates in project trino by trinodb.

the class AbstractTestingTrinoClient method toClientSession.

private static ClientSession toClientSession(Session session, URI server, Duration clientRequestTimeout) {
    ImmutableMap.Builder<String, String> properties = ImmutableMap.builder();
    properties.putAll(session.getSystemProperties());
    for (Entry<String, Map<String, String>> catalogAndConnectorProperties : session.getCatalogProperties().entrySet()) {
        for (Entry<String, String> connectorProperties : catalogAndConnectorProperties.getValue().entrySet()) {
            String catalogName = catalogAndConnectorProperties.getKey();
            properties.put(catalogName + "." + connectorProperties.getKey(), connectorProperties.getValue());
        }
    }
    ImmutableMap.Builder<String, String> resourceEstimates = ImmutableMap.builder();
    ResourceEstimates estimates = session.getResourceEstimates();
    estimates.getExecutionTime().ifPresent(e -> resourceEstimates.put(EXECUTION_TIME, e.toString()));
    estimates.getCpuTime().ifPresent(e -> resourceEstimates.put(CPU_TIME, e.toString()));
    estimates.getPeakMemoryBytes().ifPresent(e -> resourceEstimates.put(PEAK_MEMORY, e.toString()));
    return new ClientSession(server, session.getIdentity().getUser(), Optional.empty(), session.getSource().orElse(null), session.getTraceToken(), session.getClientTags(), session.getClientInfo().orElse(null), session.getCatalog().orElse(null), session.getSchema().orElse(null), session.getPath().toString(), ZoneId.of(session.getTimeZoneKey().getId()), session.getLocale(), resourceEstimates.buildOrThrow(), properties.buildOrThrow(), session.getPreparedStatements(), getRoles(session), session.getIdentity().getExtraCredentials(), session.getTransactionId().map(Object::toString).orElse(null), clientRequestTimeout, true);
}

Example 4 with ResourceEstimates

use of io.trino.spi.session.ResourceEstimates in project trino by trinodb.

the class TestStaticSelector method testSelectorResourceEstimate.

@Test
public void testSelectorResourceEstimate() {
    ResourceGroupId resourceGroupId = new ResourceGroupId(new ResourceGroupId("global"), "foo");
    StaticSelector smallQuerySelector = new StaticSelector(Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.of(new SelectorResourceEstimate(Optional.of(new Range<>(Optional.empty(), Optional.of(new Duration(5, MINUTES)))), Optional.empty(), Optional.of(new Range<>(Optional.empty(), Optional.of(DataSize.valueOf("500MB")))))), Optional.empty(), new ResourceGroupIdTemplate("global.foo"));
    assertEquals(smallQuerySelector.match(newSelectionCriteria("userA", null, ImmutableSet.of("tag1", "tag2"), new ResourceEstimates(Optional.of(java.time.Duration.ofMinutes(4)), Optional.empty(), Optional.of(DataSize.of(400, MEGABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
    assertEquals(smallQuerySelector.match(newSelectionCriteria("A.user", "a source b", ImmutableSet.of("tag1"), new ResourceEstimates(Optional.of(java.time.Duration.ofMinutes(4)), Optional.empty(), Optional.of(DataSize.of(600, MEGABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.empty());
    assertEquals(smallQuerySelector.match(newSelectionCriteria("userB", "source", ImmutableSet.of(), new ResourceEstimates(Optional.of(java.time.Duration.ofMinutes(4)), Optional.empty(), Optional.empty()))).map(SelectionContext::getResourceGroupId), Optional.empty());
    StaticSelector largeQuerySelector = new StaticSelector(Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.of(new SelectorResourceEstimate(Optional.empty(), Optional.empty(), Optional.of(new Range<>(Optional.of(DataSize.valueOf("5TB")), Optional.empty())))), Optional.empty(), new ResourceGroupIdTemplate("global.foo"));
    assertEquals(largeQuerySelector.match(newSelectionCriteria("userA", null, ImmutableSet.of("tag1", "tag2"), new ResourceEstimates(Optional.of(java.time.Duration.ofHours(100)), Optional.empty(), Optional.of(DataSize.of(4, TERABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.empty());
    assertEquals(largeQuerySelector.match(newSelectionCriteria("A.user", "a source b", ImmutableSet.of("tag1"), new ResourceEstimates(Optional.empty(), Optional.empty(), Optional.of(DataSize.of(6, TERABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
    assertEquals(largeQuerySelector.match(newSelectionCriteria("userB", "source", ImmutableSet.of(), new ResourceEstimates(Optional.of(java.time.Duration.ofSeconds(1)), Optional.of(java.time.Duration.ofSeconds(1)), Optional.of(DataSize.of(6, TERABYTE).toBytes())))).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId));
}