Java tutorial
/* * Licensed 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 com.facebook.presto.resourceGroups.db; import com.facebook.presto.resourceGroups.AbstractResourceConfigurationManager; import com.facebook.presto.resourceGroups.ManagerSpec; import com.facebook.presto.resourceGroups.ResourceGroupIdTemplate; import com.facebook.presto.resourceGroups.ResourceGroupSpec; import com.facebook.presto.resourceGroups.SelectorSpec; import com.facebook.presto.spi.memory.ClusterMemoryPoolManager; import com.facebook.presto.spi.resourceGroups.ResourceGroup; import com.facebook.presto.spi.resourceGroups.ResourceGroupId; import com.facebook.presto.spi.resourceGroups.ResourceGroupSelector; import com.facebook.presto.spi.resourceGroups.SelectionContext; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Sets; import io.airlift.units.Duration; import javax.annotation.PostConstruct; import javax.annotation.PreDestroy; import javax.annotation.concurrent.GuardedBy; import javax.inject.Inject; import java.util.AbstractMap; import java.util.HashMap; import java.util.HashSet; import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import static com.google.common.base.Preconditions.checkState; import static io.airlift.concurrent.Threads.daemonThreadsNamed; import static java.util.Objects.requireNonNull; import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; public class DbResourceGroupConfigurationManager extends AbstractResourceConfigurationManager { private final ResourceGroupsDao dao; private final ConcurrentMap<ResourceGroupId, ResourceGroup> groups = new ConcurrentHashMap<>(); @GuardedBy("this") private Map<ResourceGroupIdTemplate, ResourceGroupSpec> resourceGroupSpecs = new HashMap<>(); private final ConcurrentMap<ResourceGroupIdTemplate, List<ResourceGroupId>> configuredGroups = new ConcurrentHashMap<>(); private final AtomicReference<List<ResourceGroupSpec>> rootGroups = new AtomicReference<>(ImmutableList.of()); private final AtomicReference<List<ResourceGroupSelector>> selectors = new AtomicReference<>(); private final AtomicReference<Optional<Duration>> cpuQuotaPeriod = new AtomicReference<>(Optional.empty()); private final ScheduledExecutorService configExecutor = newSingleThreadScheduledExecutor( daemonThreadsNamed("DbResourceGroupConfigurationManager")); private final AtomicBoolean started = new AtomicBoolean(); @Inject public DbResourceGroupConfigurationManager(ClusterMemoryPoolManager memoryPoolManager, ResourceGroupsDao dao) { super(memoryPoolManager); requireNonNull(memoryPoolManager, "memoryPoolManager is null"); requireNonNull(dao, "daoProvider is null"); this.dao = dao; this.dao.createResourceGroupsGlobalPropertiesTable(); this.dao.createResourceGroupsTable(); this.dao.createSelectorsTable(); load(); } @Override protected Optional<Duration> getCpuQuotaPeriod() { return cpuQuotaPeriod.get(); } @Override protected List<ResourceGroupSpec> getRootGroups() { return rootGroups.get(); } @PreDestroy public void destroy() { configExecutor.shutdownNow(); } @PostConstruct public void start() { if (started.compareAndSet(false, true)) { configExecutor.scheduleWithFixedDelay(this::load, 1, 1, TimeUnit.SECONDS); } } @Override public void configure(ResourceGroup group, SelectionContext context) { Map.Entry<ResourceGroupIdTemplate, ResourceGroupSpec> entry = getMatchingSpec(group, context); if (groups.putIfAbsent(group.getId(), group) == null) { // If a new spec replaces the spec returned from getMatchingSpec the group will be reconfigured on the next run of load(). configuredGroups.computeIfAbsent(entry.getKey(), v -> new LinkedList<>()).add(group.getId()); } synchronized (getRootGroup(group.getId())) { configureGroup(group, entry.getValue()); } } @Override public List<ResourceGroupSelector> getSelectors() { return this.selectors.get(); } private synchronized Optional<Duration> getCpuQuotaPeriodFromDb() { List<ResourceGroupGlobalProperties> globalProperties = dao.getResourceGroupGlobalProperties(); checkState(globalProperties.size() <= 1, "There is more than one cpu_quota_period"); return (!globalProperties.isEmpty()) ? globalProperties.get(0).getCpuQuotaPeriod() : Optional.empty(); } @VisibleForTesting public synchronized void load() { Map.Entry<ManagerSpec, Map<ResourceGroupIdTemplate, ResourceGroupSpec>> specsFromDb = buildSpecsFromDb(); ManagerSpec managerSpec = specsFromDb.getKey(); Map<ResourceGroupIdTemplate, ResourceGroupSpec> resourceGroupSpecs = specsFromDb.getValue(); Set<ResourceGroupIdTemplate> changedSpecs = new HashSet<>(); Set<ResourceGroupIdTemplate> deletedSpecs = Sets.difference(this.resourceGroupSpecs.keySet(), resourceGroupSpecs.keySet()); for (Map.Entry<ResourceGroupIdTemplate, ResourceGroupSpec> entry : resourceGroupSpecs.entrySet()) { if (!entry.getValue().sameConfig(this.resourceGroupSpecs.get(entry.getKey()))) { changedSpecs.add(entry.getKey()); } } this.resourceGroupSpecs = resourceGroupSpecs; this.cpuQuotaPeriod.set(managerSpec.getCpuQuotaPeriod()); this.rootGroups.set(managerSpec.getRootGroups()); this.selectors.set(buildSelectors(managerSpec)); configureChangedGroups(changedSpecs); disableDeletedGroups(deletedSpecs); } // Populate temporary data structures to build resource group specs and selectors from db private synchronized void populateFromDbHelper(Map<Long, ResourceGroupSpecBuilder> recordMap, Set<Long> rootGroupIds, Map<Long, ResourceGroupIdTemplate> resourceGroupIdTemplateMap, Map<Long, Set<Long>> subGroupIdsToBuild) { List<ResourceGroupSpecBuilder> records = dao.getResourceGroups(); for (ResourceGroupSpecBuilder record : records) { recordMap.put(record.getId(), record); if (!record.getParentId().isPresent()) { rootGroupIds.add(record.getId()); resourceGroupIdTemplateMap.put(record.getId(), new ResourceGroupIdTemplate(record.getNameTemplate().toString())); } else { subGroupIdsToBuild.computeIfAbsent(record.getParentId().get(), k -> new HashSet<>()) .add(record.getId()); } } } private synchronized Map.Entry<ManagerSpec, Map<ResourceGroupIdTemplate, ResourceGroupSpec>> buildSpecsFromDb() { // New resource group spec map Map<ResourceGroupIdTemplate, ResourceGroupSpec> resourceGroupSpecs = new HashMap<>(); // Set of root group db ids Set<Long> rootGroupIds = new HashSet<>(); // Map of id from db to resource group spec Map<Long, ResourceGroupSpec> resourceGroupSpecMap = new HashMap<>(); // Map of id from db to resource group template id Map<Long, ResourceGroupIdTemplate> resourceGroupIdTemplateMap = new HashMap<>(); // Map of id from db to resource group spec builder Map<Long, ResourceGroupSpecBuilder> recordMap = new HashMap<>(); // Map of subgroup id's not yet built Map<Long, Set<Long>> subGroupIdsToBuild = new HashMap<>(); populateFromDbHelper(recordMap, rootGroupIds, resourceGroupIdTemplateMap, subGroupIdsToBuild); // Build up resource group specs from leaf to root for (LinkedList<Long> queue = new LinkedList<>(rootGroupIds); !queue.isEmpty();) { Long id = queue.pollFirst(); resourceGroupIdTemplateMap.computeIfAbsent(id, k -> { ResourceGroupSpecBuilder builder = recordMap.get(id); return ResourceGroupIdTemplate.forSubGroupNamed( resourceGroupIdTemplateMap.get(builder.getParentId().get()), builder.getNameTemplate().toString()); }); Set<Long> childrenToBuild = subGroupIdsToBuild.getOrDefault(id, ImmutableSet.of()); // Add to resource group specs if no more child resource groups are left to build if (childrenToBuild.isEmpty()) { ResourceGroupSpecBuilder builder = recordMap.get(id); ResourceGroupSpec resourceGroupSpec = builder.build(); resourceGroupSpecMap.put(id, resourceGroupSpec); // Add newly built spec to spec map resourceGroupSpecs.put(resourceGroupIdTemplateMap.get(id), resourceGroupSpec); // Add this resource group spec to parent subgroups and remove id from subgroup ids to build builder.getParentId().ifPresent(parentId -> { recordMap.get(parentId).addSubGroup(resourceGroupSpec); subGroupIdsToBuild.get(parentId).remove(id); }); } else { // Add this group back to queue since it still has subgroups to build queue.addFirst(id); // Add this group's subgroups to the queue so that when this id is dequeued again childrenToBuild will be empty queue.addAll(0, childrenToBuild); } } // Specs are built from db records, validate and return manager spec List<ResourceGroupSpec> rootGroups = rootGroupIds.stream().map(resourceGroupSpecMap::get) .collect(Collectors.toList()); List<SelectorSpec> selectors = dao.getSelectors().stream() .map(selectorRecord -> new SelectorSpec(selectorRecord.getUserRegex(), selectorRecord.getSourceRegex(), Optional.empty(), resourceGroupIdTemplateMap.get(selectorRecord.getResourceGroupId()))) .collect(Collectors.toList()); ManagerSpec managerSpec = new ManagerSpec(rootGroups, selectors, getCpuQuotaPeriodFromDb()); validateRootGroups(managerSpec); return new AbstractMap.SimpleImmutableEntry<>(managerSpec, resourceGroupSpecs); } private synchronized void configureChangedGroups(Set<ResourceGroupIdTemplate> changedSpecs) { for (ResourceGroupIdTemplate resourceGroupIdTemplate : changedSpecs) { for (ResourceGroupId resourceGroupId : configuredGroups.getOrDefault(resourceGroupIdTemplate, ImmutableList.of())) { synchronized (getRootGroup(resourceGroupId)) { configureGroup(groups.get(resourceGroupId), resourceGroupSpecs.get(resourceGroupIdTemplate)); } } } } private synchronized void disableDeletedGroups(Set<ResourceGroupIdTemplate> deletedSpecs) { for (ResourceGroupIdTemplate resourceGroupIdTemplate : deletedSpecs) { for (ResourceGroupId resourceGroupId : configuredGroups.getOrDefault(resourceGroupIdTemplate, ImmutableList.of())) { disableGroup(groups.get(resourceGroupId)); } } } private synchronized void disableGroup(ResourceGroup group) { // Disable groups that are removed from the db group.setMaxRunningQueries(0); group.setMaxQueuedQueries(0); } private ResourceGroup getRootGroup(ResourceGroupId groupId) { Optional<ResourceGroupId> parent = groupId.getParent(); while (parent.isPresent()) { groupId = parent.get(); parent = groupId.getParent(); } // GroupId is guaranteed to be in groups: it is added before the first call to this method in configure() return groups.get(groupId); } }