com.yahoo.vespa.hosted.controller.api.integration.resource.ResourceDatabaseClientMock Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of controller-api Show documentation
There is a newer version: 8.253.3
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package com.yahoo.vespa.hosted.controller.api.integration.resource;

import com.yahoo.config.provision.ApplicationId;
import com.yahoo.config.provision.ApplicationName;
import com.yahoo.config.provision.TenantName;
import com.yahoo.vespa.hosted.controller.api.identifiers.ClusterId;
import com.yahoo.vespa.hosted.controller.api.identifiers.DeploymentId;
import com.yahoo.vespa.hosted.controller.api.integration.billing.Plan;
import com.yahoo.vespa.hosted.controller.api.integration.billing.PlanRegistry;
import com.yahoo.vespa.hosted.controller.api.integration.configserver.Cluster;

import java.math.BigDecimal;
import java.time.Duration;
import java.time.Instant;
import java.time.YearMonth;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * @author olaa
 */
public class ResourceDatabaseClientMock implements ResourceDatabaseClient {

    PlanRegistry planRegistry;
    Map planMap = new HashMap<>();
    List resourceSnapshots = new ArrayList<>();
    Map> scalingEvents = new HashMap<>();
    private boolean hasRefreshedMaterializedView = false;

    public ResourceDatabaseClientMock(PlanRegistry planRegistry) {
        this.planRegistry = planRegistry;
    }

    @Override
    public void writeResourceSnapshots(Collection items) {
        this.resourceSnapshots.addAll(items);
    }

    @Override
    public Set getMonthsWithSnapshotsForTenant(TenantName tenantName) {
        return Collections.emptySet();
    }

    @Override
    public List getRawSnapshotHistoryForTenant(TenantName tenantName, YearMonth yearMonth) {
        return resourceSnapshots;
    }

    @Override
    public Set getTenants() {
        return resourceSnapshots.stream()
                .map(snapshot -> snapshot.getApplicationId().tenant())
                .collect(Collectors.toSet());
    }

    private List resourceUsageFromSnapshots(Plan plan, List snapshots) {
        snapshots.sort(Comparator.comparing(ResourceSnapshot::getTimestamp));

        return IntStream.range(0, snapshots.size())
                .mapToObj(idx -> {
                    var a = snapshots.get(idx);
                    var b = (idx + 1) < snapshots.size() ? snapshots.get(idx + 1) : null;
                    var start = a.getTimestamp();
                    var end   = Optional.ofNullable(b).map(ResourceSnapshot::getTimestamp).orElse(start.plusSeconds(120));
                    var d = BigDecimal.valueOf(Duration.between(start, end).toMillis());
                    return new ResourceUsage(
                            a.getApplicationId(),
                            a.getZoneId(),
                            plan,
                            a.resources().architecture(),
                            a.getMajorVersion(),
                            a.getAccount(),
                            BigDecimal.valueOf(a.resources().vcpu()).multiply(d),
                            BigDecimal.valueOf(a.resources().memoryGb()).multiply(d),
                            BigDecimal.valueOf(a.resources().diskGb()).multiply(d),
                            BigDecimal.valueOf(a.resources().gpuResources().count() * a.resources().gpuResources().memoryGb()).multiply(d));
                })
                .toList();
    }

    private ResourceUsage resourceUsageAdd(ResourceUsage a, ResourceUsage b) {
        assert a.getApplicationId().equals(b.getApplicationId());
        assert a.getZoneId().equals(b.getZoneId());
        assert a.getPlan().equals(b.getPlan());
        assert a.getArchitecture().equals(b.getArchitecture());
        assert a.getCloudAccount().equals(b.getCloudAccount());
        return new ResourceUsage(
                a.getApplicationId(),
                a.getZoneId(),
                a.getPlan(),
                a.getArchitecture(),
                a.getMajorVersion(),
                a.getCloudAccount(),
                a.getCpuMillis().add(b.getCpuMillis()),
                a.getMemoryMillis().add(b.getMemoryMillis()),
                a.getDiskMillis().add(b.getDiskMillis()),
                a.getGpuMillis().add(b.getGpuMillis()));
    }

    @Override
    public List getResourceSnapshotsForPeriod(TenantName tenantName, long start, long end) {
        var tenantPlan = planMap.getOrDefault(tenantName, planRegistry.defaultPlan());

        return resourceSnapshots.stream()
                .filter(snapshot -> snapshot.getTimestamp().isAfter(Instant.ofEpochMilli(start)))
                .filter(snapshot -> snapshot.getTimestamp().isBefore(Instant.ofEpochMilli(end)))
                .filter(snapshot -> snapshot.getApplicationId().tenant().equals(tenantName))
                .collect(Collectors.groupingBy(
                        usage -> Objects.hash(usage.getApplicationId(), usage.getZoneId(), tenantPlan.id().value(), usage.resources().architecture(), usage.getMajorVersion())
                ))
                .values().stream()
                .map(snapshots -> resourceUsageFromSnapshots(tenantPlan, snapshots))
                .map(usages -> usages.stream().reduce(this::resourceUsageAdd))
                .filter(Optional::isPresent)
                .map(Optional::get)
                .toList();
    }

    @Override
    public void refreshMaterializedView() {
        hasRefreshedMaterializedView = true;
    }

    @Override
    public Instant getOldestSnapshotTimestamp(Set deployments) {
        return Instant.ofEpochMilli(987654L);
    }

    @Override
    public void writeScalingEvents(ClusterId clusterId, Collection scalingEvents) {
        this.scalingEvents.put(clusterId, List.copyOf(scalingEvents));
    }

    @Override
    public Map> scalingEvents(Instant from, Instant to, DeploymentId deploymentId) {
        return Map.of();
    }

    public void setPlan(TenantName tenant, Plan plan) {
        planMap.put(tenant, plan);
    }

    public boolean hasRefreshedMaterializedView() {
        return hasRefreshedMaterializedView;
    }

    public List resourceSnapshots() {
        return resourceSnapshots;
    }

}