org.apache.aurora.scheduler.storage.mem.MemJobUpdateStore.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.aurora.scheduler.storage.mem.MemJobUpdateStore.java

Source

/**
 * 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.storage.mem;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Maps;
import com.google.common.collect.Ordering;
import com.google.common.primitives.Longs;
import com.google.inject.Inject;

import org.apache.aurora.common.base.MorePreconditions;
import org.apache.aurora.common.inject.TimedInterceptor.Timed;
import org.apache.aurora.common.stats.StatsProvider;
import org.apache.aurora.gen.JobInstanceUpdateEvent;
import org.apache.aurora.gen.JobUpdateDetails;
import org.apache.aurora.gen.JobUpdateEvent;
import org.apache.aurora.gen.JobUpdateState;
import org.apache.aurora.scheduler.storage.JobUpdateStore;
import org.apache.aurora.scheduler.storage.Storage.StorageException;
import org.apache.aurora.scheduler.storage.entities.IJobInstanceUpdateEvent;
import org.apache.aurora.scheduler.storage.entities.IJobUpdate;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateDetails;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateEvent;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateInstructions;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateQuery;

import static java.util.Objects.requireNonNull;

public class MemJobUpdateStore implements JobUpdateStore.Mutable {
    @VisibleForTesting
    static final String UPDATE_STORE_SIZE = "mem_storage_update_size";

    private static final Ordering<IJobUpdateDetails> REVERSE_LAST_MODIFIED_ORDER = Ordering.natural().reverse()
            .onResultOf(u -> u.getUpdate().getSummary().getState().getLastModifiedTimestampMs());

    private final Map<IJobUpdateKey, IJobUpdateDetails> updates = Maps.newConcurrentMap();

    @Inject
    MemJobUpdateStore(StatsProvider statsProvider) {
        statsProvider.makeGauge(UPDATE_STORE_SIZE, updates::size);
    }

    @Timed("job_update_store_fetch_details_query")
    @Override
    public synchronized List<IJobUpdateDetails> fetchJobUpdates(IJobUpdateQuery query) {
        return performQuery(query).collect(Collectors.toList());
    }

    @Timed("job_update_store_fetch_details")
    @Override
    public synchronized Optional<IJobUpdateDetails> fetchJobUpdate(IJobUpdateKey key) {
        return Optional.ofNullable(updates.get(key));
    }

    private static void validateInstructions(IJobUpdateInstructions instructions) {
        if (!instructions.isSetDesiredState() && instructions.getInitialState().isEmpty()) {
            throw new IllegalArgumentException(
                    "Missing both initial and desired states. At least one is required.");
        }

        if (!instructions.getInitialState().isEmpty()) {
            if (instructions.getInitialState().stream().anyMatch(t -> t.getTask() == null)) {
                throw new NullPointerException("Invalid initial instance state.");
            }
            Preconditions.checkArgument(
                    instructions.getInitialState().stream().noneMatch(t -> t.getInstances().isEmpty()),
                    "Invalid intial instance state ranges.");
        }

        if (instructions.getDesiredState() != null) {
            MorePreconditions.checkNotBlank(instructions.getDesiredState().getInstances());
            Preconditions.checkNotNull(instructions.getDesiredState().getTask());
        }
    }

    @Timed("job_update_store_save_update")
    @Override
    public synchronized void saveJobUpdate(IJobUpdate update) {
        requireNonNull(update);
        validateInstructions(update.getInstructions());

        JobUpdateDetails mutable = new JobUpdateDetails().setUpdate(update.newBuilder())
                .setUpdateEvents(ImmutableList.of()).setInstanceEvents(ImmutableList.of());
        mutable.getUpdate().getSummary().setState(synthesizeUpdateState(mutable));

        updates.put(update.getSummary().getKey(), IJobUpdateDetails.build(mutable));
    }

    private static final Ordering<JobUpdateEvent> EVENT_ORDERING = Ordering.natural()
            .onResultOf(JobUpdateEvent::getTimestampMs);

    @Timed("job_update_store_save_event")
    @Override
    public synchronized void saveJobUpdateEvent(IJobUpdateKey key, IJobUpdateEvent event) {
        IJobUpdateDetails update = updates.get(key);
        if (update == null) {
            throw new StorageException("Update not found: " + key);
        }

        JobUpdateDetails mutable = update.newBuilder();
        mutable.addToUpdateEvents(event.newBuilder());
        mutable.setUpdateEvents(EVENT_ORDERING.sortedCopy(mutable.getUpdateEvents()));
        mutable.getUpdate().getSummary().setState(synthesizeUpdateState(mutable));
        updates.put(key, IJobUpdateDetails.build(mutable));
    }

    private static final Ordering<JobInstanceUpdateEvent> INSTANCE_EVENT_ORDERING = Ordering.natural()
            .onResultOf(JobInstanceUpdateEvent::getTimestampMs);

    @Timed("job_update_store_save_instance_event")
    @Override
    public synchronized void saveJobInstanceUpdateEvent(IJobUpdateKey key, IJobInstanceUpdateEvent event) {

        IJobUpdateDetails update = updates.get(key);
        if (update == null) {
            throw new StorageException("Update not found: " + key);
        }

        JobUpdateDetails mutable = update.newBuilder();
        mutable.addToInstanceEvents(event.newBuilder());
        mutable.setInstanceEvents(INSTANCE_EVENT_ORDERING.sortedCopy(mutable.getInstanceEvents()));
        mutable.getUpdate().getSummary().setState(synthesizeUpdateState(mutable));
        updates.put(key, IJobUpdateDetails.build(mutable));
    }

    @Timed("job_update_store_delete_updates")
    @Override
    public synchronized void removeJobUpdates(Set<IJobUpdateKey> key) {
        requireNonNull(key);
        updates.keySet().removeAll(key);
    }

    @Timed("job_update_store_delete_all")
    @Override
    public synchronized void deleteAllUpdates() {
        updates.clear();
    }

    private static JobUpdateState synthesizeUpdateState(JobUpdateDetails update) {
        JobUpdateState state = new JobUpdateState();

        JobUpdateEvent firstEvent = Iterables.getFirst(update.getUpdateEvents(), null);
        if (firstEvent != null) {
            state.setCreatedTimestampMs(firstEvent.getTimestampMs());
        }

        JobUpdateEvent lastEvent = Iterables.getLast(update.getUpdateEvents(), null);
        if (lastEvent != null) {
            state.setStatus(lastEvent.getStatus());
            state.setLastModifiedTimestampMs(lastEvent.getTimestampMs());
        }

        JobInstanceUpdateEvent lastInstanceEvent = Iterables.getLast(update.getInstanceEvents(), null);
        if (lastInstanceEvent != null) {
            state.setLastModifiedTimestampMs(
                    Longs.max(state.getLastModifiedTimestampMs(), lastInstanceEvent.getTimestampMs()));
        }

        return state;
    }

    private Stream<IJobUpdateDetails> performQuery(IJobUpdateQuery query) {
        Predicate<IJobUpdateDetails> filter = u -> true;
        if (query.getRole() != null) {
            filter = filter
                    .and(u -> u.getUpdate().getSummary().getKey().getJob().getRole().equals(query.getRole()));
        }
        if (query.getKey() != null) {
            filter = filter.and(u -> u.getUpdate().getSummary().getKey().equals(query.getKey()));
        }
        if (query.getJobKey() != null) {
            filter = filter.and(u -> u.getUpdate().getSummary().getKey().getJob().equals(query.getJobKey()));
        }
        if (query.getUser() != null) {
            filter = filter.and(u -> u.getUpdate().getSummary().getUser().equals(query.getUser()));
        }
        if (query.getUpdateStatuses() != null && !query.getUpdateStatuses().isEmpty()) {
            filter = filter.and(
                    u -> query.getUpdateStatuses().contains(u.getUpdate().getSummary().getState().getStatus()));
        }

        // TODO(wfarner): Modification time is not a stable ordering for pagination, but we use it as
        // such here.  The behavior is carried over from DbJobupdateStore; determine if it is desired.
        Stream<IJobUpdateDetails> matches = updates.values().stream().filter(filter)
                .sorted(REVERSE_LAST_MODIFIED_ORDER).skip(query.getOffset());

        if (query.getLimit() > 0) {
            matches = matches.limit(query.getLimit());
        }

        return matches;
    }
}