/**
* 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 org.apache.aurora.scheduler.filter;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Function;
import com.google.common.base.Supplier;
import com.google.common.base.Suppliers;
import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableMultiset;
import com.google.common.collect.Multiset;
import org.apache.aurora.common.collections.Pair;
import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.storage.AttributeStore;
import org.apache.aurora.scheduler.storage.Storage.StoreProvider;
import org.apache.aurora.scheduler.storage.entities.IAttribute;
import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import static java.util.Objects.requireNonNull;
/**
* A temporary view of a job's state. Once constructed, instances of this class should be discarded
* once the job state may change (e.g. after exiting a write transaction). This is intended to
* capture job state once and avoid redundant queries.
* <p>
* TODO(wfarner): Consider preserving this as only a helper class to compute the Multiset
* representing the aggregate, since this class is now a thin wrapper over a Multiset.
*/
public final class AttributeAggregate {
/**
* A mapping from attribute name and value to the count of tasks with that name/value combination.
* See doc for {@link #getNumTasksWithAttribute(String, String)} for further details.
*/
private Supplier<Multiset<Pair<String, String>>> aggregate;
private boolean isInitialized = false;
private AttributeAggregate(Supplier<Multiset<Pair<String, String>>> aggregate) {
this.aggregate = Suppliers.memoize(
() -> {
initialize();
return aggregate.get();
}
);
}
private void initialize() {
isInitialized = true; // inlining this assignment yields a PMD false positive
}
/**
* Initializes an {@link AttributeAggregate} instance from data store.
*
* @param storeProvider Store provider to get data from.
* @param jobKey Job key.
* @return An {@link AttributeAggregate} instance.
*/
public static AttributeAggregate getJobActiveState(
final StoreProvider storeProvider,
final IJobKey jobKey) {
return create(
() -> storeProvider.getTaskStore()
.fetchTasks(Query.jobScoped(jobKey).byStatus(Tasks.SLAVE_ASSIGNED_STATES)),
storeProvider.getAttributeStore());
}
@VisibleForTesting
static AttributeAggregate create(
final Supplier<Iterable<IScheduledTask>> taskSupplier,
final AttributeStore attributeStore) {
final Function<String, Iterable<IAttribute>> getHostAttributes =
host -> {
// Note: this assumes we have access to attributes for hosts where all active tasks
// reside.
requireNonNull(host);
return attributeStore.getHostAttributes(host).get().getAttributes();
};
return create(Suppliers.compose(
tasks -> FluentIterable.from(tasks)
.transform(Tasks::scheduledToSlaveHost)
.transformAndConcat(getHostAttributes),
taskSupplier));
}
@VisibleForTesting
static AttributeAggregate create(Supplier<Iterable<IAttribute>> attributes) {
Supplier<Multiset<Pair<String, String>>> aggregator = Suppliers.compose(
attributes1 -> addAttributes(ImmutableMultiset.builder(), attributes1).build(),
attributes);
return new AttributeAggregate(aggregator);
}
private static ImmutableMultiset.Builder<Pair<String, String>> addAttributes(
ImmutableMultiset.Builder<Pair<String, String>> builder,
Iterable<IAttribute> attributes) {
for (IAttribute attribute : attributes) {
for (String value : attribute.getValues()) {
builder.add(Pair.of(attribute.getName(), value));
}
}
return builder;
}
public void updateAttributeAggregate(IHostAttributes attributes) {
// If the aggregate supplier has not been populated there is no need to update it here.
// All tasks attributes will be picked up by the wrapped task query if executed at a
// later point in time.
if (isInitialized) {
final Supplier<Multiset<Pair<String, String>>> previous = aggregate;
aggregate = Suppliers.memoize(
() -> {
ImmutableMultiset.Builder<Pair<String, String>> builder
= new ImmutableMultiset.Builder<>();
builder.addAll(previous.get());
addAttributes(builder, attributes.getAttributes());
return builder.build();
}
);
}
}
@VisibleForTesting
public static AttributeAggregate empty() {
return new AttributeAggregate(Suppliers.ofInstance(ImmutableMultiset.of()));
}
/**
* Gets the total number of tasks with a given attribute name and value combination.
* <p>
* For example, the counts for a group of tasks that are host-diverse but not rack-diverse, you
* may have counts like this:
* <pre>
* {
* ("host", "hostA"): 1
* ("host", "hostB"): 1
* ("rack", "rackA"): 2
* }
* </pre>
*
* @param name Name of the attribute.
* @param value Value of the attribute.
* @return Number of tasks in the job whose hosts have the provided attribute name and value.
*/
public long getNumTasksWithAttribute(String name, String value) {
return aggregate.get().count(Pair.of(name, value));
}
@VisibleForTesting
Multiset<Pair<String, String>> getAggregates() {
return aggregate.get();
}
@Override
public boolean equals(Object o) {
if (!(o instanceof AttributeAggregate)) {
return false;
}
AttributeAggregate other = (AttributeAggregate) o;
return getAggregates().equals(other.getAggregates());
}
@Override
public int hashCode() {
return getAggregates().hashCode();
}
}