google-nomulus/java/google/registry/dns/ReadDnsQueueAction.java

// Copyright 2017 The Nomulus Authors. All Rights Reserved.
//
// 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 google.registry.dns;

import static com.google.appengine.api.taskqueue.TaskOptions.Builder.withUrl;
import static com.google.common.collect.Sets.difference;
import static google.registry.dns.DnsConstants.DNS_PUBLISH_PUSH_QUEUE_NAME;
import static google.registry.dns.DnsConstants.DNS_TARGET_NAME_PARAM;
import static google.registry.dns.DnsConstants.DNS_TARGET_TYPE_PARAM;
import static google.registry.model.registry.Registries.getTlds;
import static java.util.concurrent.TimeUnit.SECONDS;

import com.google.appengine.api.taskqueue.Queue;
import com.google.appengine.api.taskqueue.TaskHandle;
import com.google.appengine.api.taskqueue.TaskOptions;
import com.google.auto.value.AutoValue;
import com.google.common.collect.ComparisonChain;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.SortedSetMultimap;
import com.google.common.collect.TreeMultimap;
import google.registry.config.RegistryConfig.Config;
import google.registry.dns.DnsConstants.TargetType;
import google.registry.model.registry.Registry;
import google.registry.request.Action;
import google.registry.request.Parameter;
import google.registry.request.RequestParameters;
import google.registry.request.auth.Auth;
import google.registry.util.FormattingLogger;
import google.registry.util.TaskEnqueuer;
import java.io.UnsupportedEncodingException;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.Set;
import javax.inject.Inject;
import javax.inject.Named;
import org.joda.time.Duration;

/**
 * Action for fanning out DNS refresh tasks by TLD, using data taken from the DNS pull queue.
 *
 * <h3>Parameters Reference</h3>
 *
 * <ul>
 * <li>{@code jitterSeconds} Randomly delay each task by up to this many seconds.
 * </ul>
 */
@Action(
  path = "/_dr/cron/readDnsQueue",
  automaticallyPrintOk = true,
  auth = Auth.AUTH_INTERNAL_ONLY
)
public final class ReadDnsQueueAction implements Runnable {

  private static final String PARAM_JITTER_SECONDS = "jitterSeconds";
  private static final Random random = new Random();
  private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();

  @Inject @Config("dnsTldUpdateBatchSize") int tldUpdateBatchSize;
  @Inject @Config("dnsWriteLockTimeout") Duration writeLockTimeout;
  @Inject @Named(DNS_PUBLISH_PUSH_QUEUE_NAME) Queue dnsPublishPushQueue;
  @Inject @Parameter(PARAM_JITTER_SECONDS) Optional<Integer> jitterSeconds;
  @Inject DnsQueue dnsQueue;
  @Inject TaskEnqueuer taskEnqueuer;
  @Inject ReadDnsQueueAction() {}

  /** Container for items we pull out of the DNS pull queue and process for fanout. */
  @AutoValue
  abstract static class RefreshItem implements Comparable<RefreshItem> {
    static RefreshItem create(TargetType type, String name) {
      return new AutoValue_ReadDnsQueueAction_RefreshItem(type, name);
    }

    abstract TargetType type();

    abstract String name();

    @Override
    public int compareTo(RefreshItem other) {
      return ComparisonChain.start()
          .compare(this.type(), other.type())
          .compare(this.name(), other.name())
          .result();
    }
  }

  /** Leases all tasks from the pull queue and creates per-tld update actions for them. */
  @Override
  public void run() {
    Set<String> tldsOfInterest = getTlds();

    List<TaskHandle> tasks = dnsQueue.leaseTasks(writeLockTimeout);
    if (tasks.isEmpty()) {
      return;
    }
    logger.infofmt("Leased %d DNS update tasks.", tasks.size());
    // Normally, all tasks will be deleted from the pull queue. But some might have to remain if
    // we are not interested in the associated TLD, or if the TLD is paused. Remember which these
    // are.
    Set<TaskHandle> tasksToKeep = new HashSet<>();
    // The paused TLDs for which we found at least one refresh request.
    Set<String> pausedTlds = new HashSet<>();
    // Create a sorted multimap into which we will insert the refresh items, so that the items for
    // each TLD will be grouped together, and domains and hosts will be grouped within a TLD. The
    // grouping and ordering of domains and hosts is not technically necessary, but a predictable
    // ordering makes it possible to write detailed tests.
    SortedSetMultimap<String, RefreshItem> refreshItemMultimap = TreeMultimap.create();
    // Read all tasks on the DNS pull queue and load them into the refresh item multimap.
    for (TaskHandle task : tasks) {
      try {
        Map<String, String> params = ImmutableMap.copyOf(task.extractParams());
        String tld = params.get(RequestParameters.PARAM_TLD);
        if (tld == null) {
          logger.severe("Discarding invalid DNS refresh request; no TLD specified.");
        } else if (!tldsOfInterest.contains(tld)) {
          tasksToKeep.add(task);
        } else if (Registry.get(tld).getDnsPaused()) {
          tasksToKeep.add(task);
          pausedTlds.add(tld);
        } else {
          String typeString = params.get(DNS_TARGET_TYPE_PARAM);
          String name = params.get(DNS_TARGET_NAME_PARAM);
          TargetType type = TargetType.valueOf(typeString);
          switch (type) {
            case DOMAIN:
            case HOST:
              refreshItemMultimap.put(tld, RefreshItem.create(type, name));
              break;
            default:
              logger.severefmt("Discarding DNS refresh request of type %s.", typeString);
              break;
          }
        }
      } catch (RuntimeException | UnsupportedEncodingException e) {
        logger.severefmt(e, "Discarding invalid DNS refresh request (task %s).", task);
      }
    }
    if (!pausedTlds.isEmpty()) {
      logger.infofmt("The dns-pull queue is paused for TLDs: %s.", pausedTlds);
    }
    // Loop through the multimap by TLD and generate refresh tasks for the hosts and domains for
    // each configured DNS writer.
    for (Map.Entry<String, Collection<RefreshItem>> tldRefreshItemsEntry
        : refreshItemMultimap.asMap().entrySet()) {
      String tld = tldRefreshItemsEntry.getKey();
      for (List<RefreshItem> chunk :
          Iterables.partition(tldRefreshItemsEntry.getValue(), tldUpdateBatchSize)) {
        for (String dnsWriter : Registry.get(tld).getDnsWriters()) {
          TaskOptions options = withUrl(PublishDnsUpdatesAction.PATH)
              .countdownMillis(jitterSeconds.isPresent()
                  ? random.nextInt((int) SECONDS.toMillis(jitterSeconds.get()))
                  : 0)
              .param(RequestParameters.PARAM_TLD, tld)
              .param(PublishDnsUpdatesAction.PARAM_DNS_WRITER, dnsWriter);
          for (RefreshItem refreshItem : chunk) {
            options.param(
                (refreshItem.type() == TargetType.HOST)
                    ? PublishDnsUpdatesAction.PARAM_HOSTS
                    : PublishDnsUpdatesAction.PARAM_DOMAINS,
                refreshItem.name());
          }
          taskEnqueuer.enqueue(dnsPublishPushQueue, options);
        }
      }
    }
    Set<TaskHandle> tasksToDelete = difference(ImmutableSet.copyOf(tasks), tasksToKeep);
    // Either delete or drop the lease of each task.
    logger.infofmt("Deleting %d DNS update tasks.", tasksToDelete.size());
    dnsQueue.deleteTasks(ImmutableList.copyOf(tasksToDelete));
    logger.infofmt("Dropping %d DNS update tasks.", tasksToKeep.size());
    for (TaskHandle task : tasksToKeep) {
      dnsQueue.dropTaskLease(task);
    }
    logger.infofmt("Done processing DNS tasks.");
  }
}