content: /**
   * 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.durability;
  
  import java.util.Collection;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  
  import com.google.common.base.Function;
  import com.google.common.base.Preconditions;
  import com.google.common.collect.ImmutableSet;
  
  import org.apache.aurora.gen.storage.Op;
  import org.apache.aurora.gen.storage.RemoveHostMaintenanceRequest;
  import org.apache.aurora.gen.storage.RemoveJob;
  import org.apache.aurora.gen.storage.RemoveQuota;
  import org.apache.aurora.gen.storage.RemoveTasks;
  import org.apache.aurora.gen.storage.SaveCronJob;
  import org.apache.aurora.gen.storage.SaveFrameworkId;
  import org.apache.aurora.gen.storage.SaveHostAttributes;
  import org.apache.aurora.gen.storage.SaveHostMaintenanceRequest;
  import org.apache.aurora.gen.storage.SaveJobInstanceUpdateEvent;
  import org.apache.aurora.gen.storage.SaveJobUpdate;
  import org.apache.aurora.gen.storage.SaveJobUpdateEvent;
  import org.apache.aurora.gen.storage.SaveQuota;
  import org.apache.aurora.gen.storage.SaveTasks;
  import org.apache.aurora.scheduler.base.Query;
  import org.apache.aurora.scheduler.events.EventSink;
  import org.apache.aurora.scheduler.events.PubsubEvent;
  import org.apache.aurora.scheduler.storage.AttributeStore;
  import org.apache.aurora.scheduler.storage.CronJobStore;
  import org.apache.aurora.scheduler.storage.HostMaintenanceStore;
  import org.apache.aurora.scheduler.storage.JobUpdateStore;
  import org.apache.aurora.scheduler.storage.QuotaStore;
  import org.apache.aurora.scheduler.storage.SchedulerStore;
  import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
  import org.apache.aurora.scheduler.storage.TaskStore;
  import org.apache.aurora.scheduler.storage.durability.DurableStorage.TransactionManager;
  import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
  import org.apache.aurora.scheduler.storage.entities.IHostMaintenanceRequest;
  import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
  import org.apache.aurora.scheduler.storage.entities.IJobInstanceUpdateEvent;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  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.IJobUpdateKey;
  import org.apache.aurora.scheduler.storage.entities.IJobUpdateQuery;
  import org.apache.aurora.scheduler.storage.entities.IResourceAggregate;
  import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
  import org.slf4j.Logger;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Mutable stores implementation that translates all operations to {@link Op}s (which are passed
   * to a provided {@link TransactionManager}) before forwarding the operations to delegate mutable
   * stores.
   */
  public class WriteRecorder implements
      MutableStoreProvider,
      SchedulerStore.Mutable,
      CronJobStore.Mutable,
      TaskStore.Mutable,
      QuotaStore.Mutable,
      AttributeStore.Mutable,
      JobUpdateStore.Mutable,
      HostMaintenanceStore.Mutable {
  
    private final TransactionManager transactionManager;
    private final SchedulerStore.Mutable schedulerStore;
    private final CronJobStore.Mutable jobStore;
    private final TaskStore.Mutable taskStore;
    private final QuotaStore.Mutable quotaStore;
    private final AttributeStore.Mutable attributeStore;
    private final JobUpdateStore.Mutable jobUpdateStore;
    private final HostMaintenanceStore.Mutable hostMaintenanceStore;
    private final Logger log;
    private final EventSink eventSink;
  
    /**
     * Creates a new write-ahead storage that delegates to the providing default stores.
     *
     * @param transactionManager External controller for transaction operations.
     * @param schedulerStore Delegate.
     * @param jobStore       Delegate.
     * @param taskStore      Delegate.
     * @param quotaStore     Delegate.
     * @param attributeStore Delegate.
     * @param jobUpdateStore Delegate.
     */
    public WriteRecorder(
        TransactionManager transactionManager,
        SchedulerStore.Mutable schedulerStore,
        CronJobStore.Mutable jobStore,
        TaskStore.Mutable taskStore,
        QuotaStore.Mutable quotaStore,
        AttributeStore.Mutable attributeStore,
        JobUpdateStore.Mutable jobUpdateStore,
        HostMaintenanceStore.Mutable hostMaintenanceStore,
        Logger log,
        EventSink eventSink) {
  
      this.transactionManager = requireNonNull(transactionManager);
      this.schedulerStore = requireNonNull(schedulerStore);
      this.jobStore = requireNonNull(jobStore);
      this.taskStore = requireNonNull(taskStore);
      this.quotaStore = requireNonNull(quotaStore);
      this.attributeStore = requireNonNull(attributeStore);
      this.jobUpdateStore = requireNonNull(jobUpdateStore);
      this.hostMaintenanceStore = requireNonNull(hostMaintenanceStore);
      this.log = requireNonNull(log);
      this.eventSink = requireNonNull(eventSink);
    }
  
    private void write(Op op) {
      Preconditions.checkState(
          transactionManager.hasActiveTransaction(),
          "Mutating operations must be within a transaction.");
      transactionManager.log(op);
    }
  
    @Override
    public void saveFrameworkId(final String frameworkId) {
      requireNonNull(frameworkId);
  
      write(Op.saveFrameworkId(new SaveFrameworkId(frameworkId)));
      schedulerStore.saveFrameworkId(frameworkId);
    }
  
    @Override
    public void deleteTasks(final Set<String> taskIds) {
      requireNonNull(taskIds);
  
      write(Op.removeTasks(new RemoveTasks(taskIds)));
      taskStore.deleteTasks(taskIds);
    }
  
    @Override
    public void saveTasks(final Set<IScheduledTask> newTasks) {
      requireNonNull(newTasks);
  
      write(Op.saveTasks(new SaveTasks(IScheduledTask.toBuildersSet(newTasks))));
      taskStore.saveTasks(newTasks);
    }
  
    @Override
    public Optional<IScheduledTask> mutateTask(
        String taskId,
        Function<IScheduledTask, IScheduledTask> mutator) {
  
      Optional<IScheduledTask> mutated = taskStore.mutateTask(taskId, mutator);
      log.debug("Storing updated task to log: {}={}", taskId, mutated.get().getStatus());
      write(Op.saveTasks(new SaveTasks(ImmutableSet.of(mutated.get().newBuilder()))));
  
      return mutated;
    }
  
    @Override
    public void saveQuota(final String role, final IResourceAggregate quota) {
      requireNonNull(role);
      requireNonNull(quota);
  
      write(Op.saveQuota(new SaveQuota(role, quota.newBuilder())));
      quotaStore.saveQuota(role, quota);
    }
  
    @Override
    public boolean saveHostAttributes(final IHostAttributes attrs) {
      requireNonNull(attrs);
  
      boolean changed = attributeStore.saveHostAttributes(attrs);
      if (changed) {
        write(Op.saveHostAttributes(new SaveHostAttributes(attrs.newBuilder())));
        eventSink.post(new PubsubEvent.HostAttributesChanged(attrs));
      }
      return changed;
    }
  
    @Override
    public void removeJob(final IJobKey jobKey) {
      requireNonNull(jobKey);
  
      write(Op.removeJob(new RemoveJob().setJobKey(jobKey.newBuilder())));
      jobStore.removeJob(jobKey);
    }
  
    @Override
    public void saveAcceptedJob(final IJobConfiguration jobConfig) {
      requireNonNull(jobConfig);
  
      write(Op.saveCronJob(new SaveCronJob(jobConfig.newBuilder())));
      jobStore.saveAcceptedJob(jobConfig);
    }
  
    @Override
    public void removeQuota(final String role) {
      requireNonNull(role);
  
      write(Op.removeQuota(new RemoveQuota(role)));
      quotaStore.removeQuota(role);
    }
  
    @Override
    public void saveJobUpdate(IJobUpdate update) {
      requireNonNull(update);
  
      write(Op.saveJobUpdate(new SaveJobUpdate().setJobUpdate(update.newBuilder())));
      jobUpdateStore.saveJobUpdate(update);
    }
  
    @Override
    public void saveJobUpdateEvent(IJobUpdateKey key, IJobUpdateEvent event) {
      requireNonNull(key);
      requireNonNull(event);
  
      write(Op.saveJobUpdateEvent(new SaveJobUpdateEvent(event.newBuilder(), key.newBuilder())));
      jobUpdateStore.saveJobUpdateEvent(key, event);
    }
  
    @Override
    public void saveJobInstanceUpdateEvent(IJobUpdateKey key, IJobInstanceUpdateEvent event) {
      requireNonNull(key);
      requireNonNull(event);
  
      write(Op.saveJobInstanceUpdateEvent(
          new SaveJobInstanceUpdateEvent(event.newBuilder(), key.newBuilder())));
      jobUpdateStore.saveJobInstanceUpdateEvent(key, event);
    }
  
    @Override
    public void removeJobUpdates(Set<IJobUpdateKey> keys) {
      requireNonNull(keys);
  
      // Compatibility mode - RemoveJobUpdates is not yet written since older versions cannot
      // read it.  JobUpdates are only removed implicitly when a snapshot is taken.
      jobUpdateStore.removeJobUpdates(keys);
    }
  
    @Override
    public void saveHostMaintenanceRequest(IHostMaintenanceRequest hostMaintenanceRequest) {
      requireNonNull(hostMaintenanceRequest);
  
      write(Op.saveHostMaintenanceRequest(
          new SaveHostMaintenanceRequest(hostMaintenanceRequest.newBuilder())));
      this.hostMaintenanceStore.saveHostMaintenanceRequest(hostMaintenanceRequest);
    }
  
    @Override
    public void removeHostMaintenanceRequest(String host) {
      requireNonNull(host);
  
      write(Op.removeHostMaintenanceRequest(
          new RemoveHostMaintenanceRequest(host)));
      this.hostMaintenanceStore.removeHostMaintenanceRequest(host);
    }
  
    @Override
    public void deleteAllTasks() {
      throw new UnsupportedOperationException(
          "Unsupported since casual storage users should never be doing this.");
    }
  
    @Override
    public void deleteHostAttributes() {
      throw new UnsupportedOperationException(
          "Unsupported since casual storage users should never be doing this.");
    }
  
    @Override
    public void deleteHostMaintenanceRequests() {
      throw new UnsupportedOperationException(
          "Unsupported since casual storage users should never be doing this.");
    }
  
    @Override
    public void deleteJobs() {
      throw new UnsupportedOperationException(
          "Unsupported since casual storage users should never be doing this.");
    }
  
    @Override
    public void deleteQuotas() {
      throw new UnsupportedOperationException(
          "Unsupported since casual storage users should never be doing this.");
    }
  
    @Override
    public void deleteAllUpdates() {
      throw new UnsupportedOperationException(
          "Unsupported since casual storage users should never be doing this.");
    }
  
    @Override
    public SchedulerStore.Mutable getSchedulerStore() {
      return this;
    }
  
    @Override
    public CronJobStore.Mutable getCronJobStore() {
      return this;
    }
  
    @Override
    public TaskStore.Mutable getUnsafeTaskStore() {
      return this;
    }
  
    @Override
    public QuotaStore.Mutable getQuotaStore() {
      return this;
    }
  
    @Override
    public AttributeStore.Mutable getAttributeStore() {
      return this;
    }
  
    @Override
    public TaskStore getTaskStore() {
      return this;
    }
  
    @Override
    public JobUpdateStore.Mutable getJobUpdateStore() {
      return this;
    }
  
    @Override
    public HostMaintenanceStore.Mutable getHostMaintenanceStore() {
      return this;
    }
  
    @Override
    public Optional<String> fetchFrameworkId() {
      return this.schedulerStore.fetchFrameworkId();
    }
  
    @Override
    public Iterable<IJobConfiguration> fetchJobs() {
      return this.jobStore.fetchJobs();
    }
  
    @Override
    public Optional<IJobConfiguration> fetchJob(IJobKey jobKey) {
      return this.jobStore.fetchJob(jobKey);
    }
  
    @Override
    public Optional<IScheduledTask> fetchTask(String taskId) {
      return this.taskStore.fetchTask(taskId);
    }
  
    @Override
    public Collection<IScheduledTask> fetchTasks(Query.Builder query) {
      return this.taskStore.fetchTasks(query);
    }
  
    @Override
    public Set<IJobKey> getJobKeys() {
      return this.taskStore.getJobKeys();
    }
  
    @Override
    public Optional<IResourceAggregate> fetchQuota(String role) {
      return this.quotaStore.fetchQuota(role);
    }
  
    @Override
    public Map<String, IResourceAggregate> fetchQuotas() {
      return this.quotaStore.fetchQuotas();
    }
  
    @Override
    public Optional<IHostAttributes> getHostAttributes(String host) {
      return this.attributeStore.getHostAttributes(host);
    }
  
    @Override
    public Set<IHostAttributes> getHostAttributes() {
      return this.attributeStore.getHostAttributes();
    }
  
    @Override
    public List<IJobUpdateDetails> fetchJobUpdates(IJobUpdateQuery query) {
      return this.jobUpdateStore.fetchJobUpdates(query);
    }
  
    @Override
    public Optional<IJobUpdateDetails> fetchJobUpdate(IJobUpdateKey key) {
      return this.jobUpdateStore.fetchJobUpdate(key);
    }
  
    @Override
    public Optional<IHostMaintenanceRequest> getHostMaintenanceRequest(String host) {
      return this.hostMaintenanceStore.getHostMaintenanceRequest(host);
    }
  
    @Override
    public Set<IHostMaintenanceRequest> getHostMaintenanceRequests() {
      return this.hostMaintenanceStore.getHostMaintenanceRequests();
    }
  }
  