<commit_before>// Copyright (c) 1999-2004 <NAME> (<EMAIL>)

package org.xbill.DNS;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

/**
 * An implementation of {@link Resolver} that can send queries to multiple servers, sending the
 * queries multiple times if necessary.
 *
 * @see Resolver
 * @author <NAME>
 */
@Slf4j
public class ExtendedResolver implements Resolver {
  private static class Resolution {
    private final Message query;
    private final int[] attempts;
    private final int retriesPerResolver;
    private List<ResolverEntry> resolvers;
    private int currentResolver;
    private long endTime;

    Resolution(ExtendedResolver eres, Message query) {
      resolvers = new ArrayList<>(eres.resolvers);
      endTime = System.nanoTime() + eres.timeout.toNanos();
      if (eres.loadBalance) {
        int start = eres.lbStart.updateAndGet(i -> i++ % resolvers.size());
        if (start > 0) {
          List<ResolverEntry> shuffle = new ArrayList<>(resolvers.size());
          for (int i = 0; i < resolvers.size(); i++) {
            int pos = (i + start) % resolvers.size();
            shuffle.add(resolvers.get(pos));
          }

          resolvers = shuffle;
        }
      } else {
        Collections.shuffle(resolvers);
        resolvers =
            resolvers.stream()
                .sorted(Comparator.comparingInt(re -> re.failures.get()))
                .collect(Collectors.toList());
      }

      attempts = new int[resolvers.size()];
      retriesPerResolver = eres.retries;
      this.query = query;
    }

    /* Asynchronously sends a message. */
    private CompletableFuture<Message> send() {
      ResolverEntry r = resolvers.get(currentResolver);
      log.debug(
          "Sending {}/{}, id={} to resolver {} ({}), attempt {} of {}",
          query.getQuestion().getName(),
          Type.string(query.getQuestion().getType()),
          query.getHeader().getID(),
          currentResolver,
          r.resolver,
          attempts[currentResolver] + 1,
          retriesPerResolver);
      attempts[currentResolver]++;
      return r.resolver.sendAsync(query).toCompletableFuture();
    }

    /* Start an asynchronous resolution */
    private CompletionStage<Message> startAsync() {
      CompletableFuture<Message> f = new CompletableFuture<>();
      send().handleAsync((result, ex) -> handle(result, ex, f));
      return f;
    }

    private Void handle(Message result, Throwable ex, CompletableFuture<Message> f) {
      AtomicInteger failureCounter = resolvers.get(currentResolver).failures;
      if (ex != null) {
        log.debug(
            "Failed to resolve {}/{}, id={} with resolver {} ({}) on attempt {} of {}, reason={}",
            query.getQuestion().getName(),
            Type.string(query.getQuestion().getType()),
            query.getHeader().getID(),
            currentResolver,
            resolvers.get(currentResolver).resolver,
            attempts[currentResolver],
            retriesPerResolver,
            ex.getMessage());

        failureCounter.incrementAndGet();

        if (endTime - System.nanoTime() < 0) {
          f.completeExceptionally(
              new IOException(
                  "Timed out while trying to resolve "
                      + query.getQuestion().getName()
                      + "/"
                      + Type.string(query.getQuestion().type)
                      + ", id="
                      + query.getHeader().getID()));
        } else {
          // go to next resolver, until retries on all resolvers are exhausted
          currentResolver = (currentResolver + 1) % resolvers.size();
          if (attempts[currentResolver] < retriesPerResolver) {
            send().handleAsync((r, t) -> handle(r, t, f));
            return null;
          }

          f.completeExceptionally(ex);
        }
      } else {
        failureCounter.updateAndGet(i -> i > 0 ? (int) Math.log(i) : 0);
        f.complete(result);
      }

      return null;
    }
  }

  @RequiredArgsConstructor
  private static class ResolverEntry {
    private final Resolver resolver;
    private final AtomicInteger failures;

    ResolverEntry(Resolver r) {
      this(r, new AtomicInteger(0));
    }

    @Override
    public String toString() {
      return resolver.toString();
    }
  }

  /**
   * Default timeout until resolving is aborted.
   *
   * @since 3.2
   */
  public static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(10);

  /**
   * Default timeout until resolving with one of the used resolvers fails.
   *
   * @since 3.2
   */
  public static final Duration DEFAULT_RESOLVER_TIMEOUT = Duration.ofSeconds(5);

  private final List<ResolverEntry> resolvers = new CopyOnWriteArrayList<>();
  private final AtomicInteger lbStart = new AtomicInteger();
  private boolean loadBalance;
  private int retries = 3;
  private Duration timeout = DEFAULT_TIMEOUT;

  /**
   * Creates a new Extended Resolver. The default {@link ResolverConfig} is used to determine the
   * servers for which {@link SimpleResolver}s are initialized. The timeout for each server is
   * initialized with {@link #DEFAULT_RESOLVER_TIMEOUT}.
   */
  public ExtendedResolver() {
    List<InetSocketAddress> servers = ResolverConfig.getCurrentConfig().servers();
    resolvers.addAll(
        servers.stream()
            .map(
                server -> {
                  Resolver r = new SimpleResolver(server);
                  r.setTimeout(DEFAULT_RESOLVER_TIMEOUT);
                  return new ResolverEntry(r);
                })
            .collect(Collectors.toSet()));
  }

  /**
   * Creates a new instance with {@link SimpleResolver}s. The timeout for each server is initialized
   * with {@link #DEFAULT_RESOLVER_TIMEOUT}.
   *<commit_msg>Keep order of servers when not load-balancing<commit_after>// Copyright (c) 1999-2004 <NAME> (<EMAIL>)

package org.xbill.DNS;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

/**
 * An implementation of {@link Resolver} that can send queries to multiple servers, sending the
 * queries multiple times if necessary.
 *
 * @see Resolver
 * @author <NAME>
 */
@Slf4j
public class ExtendedResolver implements Resolver {
  private static class Resolution {
    private final Message query;
    private final int[] attempts;
    private final int retriesPerResolver;
    private List<ResolverEntry> resolvers;
    private int currentResolver;
    private long endTime;

    Resolution(ExtendedResolver eres, Message query) {
      resolvers = new ArrayList<>(eres.resolvers);
      endTime = System.nanoTime() + eres.timeout.toNanos();
      if (eres.loadBalance) {
        int start = eres.lbStart.updateAndGet(i -> i++ % resolvers.size());
        if (start > 0) {
          List<ResolverEntry> shuffle = new ArrayList<>(resolvers.size());
          for (int i = 0; i < resolvers.size(); i++) {
            int pos = (i + start) % resolvers.size();
            shuffle.add(resolvers.get(pos));
          }

          resolvers = shuffle;
        }
      } else {
        resolvers =
            resolvers.stream()
                .sorted(Comparator.comparingInt(re -> re.failures.get()))
                .collect(Collectors.toList());
      }

      attempts = new int[resolvers.size()];
      retriesPerResolver = eres.retries;
      this.query = query;
    }

    /* Asynchronously sends a message. */
    private CompletableFuture<Message> send() {
      ResolverEntry r = resolvers.get(currentResolver);
      log.debug(
          "Sending {}/{}, id={} to resolver {} ({}), attempt {} of {}",
          query.getQuestion().getName(),
          Type.string(query.getQuestion().getType()),
          query.getHeader().getID(),
          currentResolver,
          r.resolver,
          attempts[currentResolver] + 1,
          retriesPerResolver);
      attempts[currentResolver]++;
      return r.resolver.sendAsync(query).toCompletableFuture();
    }

    /* Start an asynchronous resolution */
    private CompletionStage<Message> startAsync() {
      CompletableFuture<Message> f = new CompletableFuture<>();
      send().handleAsync((result, ex) -> handle(result, ex, f));
      return f;
    }

    private Void handle(Message result, Throwable ex, CompletableFuture<Message> f) {
      AtomicInteger failureCounter = resolvers.get(currentResolver).failures;
      if (ex != null) {
        log.debug(
            "Failed to resolve {}/{}, id={} with resolver {} ({}) on attempt {} of {}, reason={}",
            query.getQuestion().getName(),
            Type.string(query.getQuestion().getType()),
            query.getHeader().getID(),
            currentResolver,
            resolvers.get(currentResolver).resolver,
            attempts[currentResolver],
            retriesPerResolver,
            ex.getMessage());

        failureCounter.incrementAndGet();

        if (endTime - System.nanoTime() < 0) {
          f.completeExceptionally(
              new IOException(
                  "Timed out while trying to resolve "
                      + query.getQuestion().getName()
                      + "/"
                      + Type.string(query.getQuestion().type)
                      + ", id="
                      + query.getHeader().getID()));
        } else {
          // go to next resolver, until retries on all resolvers are exhausted
          currentResolver = (currentResolver + 1) % resolvers.size();
          if (attempts[currentResolver] < retriesPerResolver) {
            send().handleAsync((r, t) -> handle(r, t, f));
            return null;
          }

          f.completeExceptionally(ex);
        }
      } else {
        failureCounter.updateAndGet(i -> i > 0 ? (int) Math.log(i) : 0);
        f.complete(result);
      }

      return null;
    }
  }

  @RequiredArgsConstructor
  private static class ResolverEntry {
    private final Resolver resolver;
    private final AtomicInteger failures;

    ResolverEntry(Resolver r) {
      this(r, new AtomicInteger(0));
    }

    @Override
    public String toString() {
      return resolver.toString();
    }
  }

  /**
   * Default timeout until resolving is aborted.
   *
   * @since 3.2
   */
  public static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(10);

  /**
   * Default timeout until resolving with one of the used resolvers fails.
   *
   * @since 3.2
   */
  public static final Duration DEFAULT_RESOLVER_TIMEOUT = Duration.ofSeconds(5);

  private final List<ResolverEntry> resolvers = new CopyOnWriteArrayList<>();
  private final AtomicInteger lbStart = new AtomicInteger();
  private boolean loadBalance;
  private int retries = 3;
  private Duration timeout = DEFAULT_TIMEOUT;

  /**
   * Creates a new Extended Resolver. The default {@link ResolverConfig} is used to determine the
   * servers for which {@link SimpleResolver}s are initialized. The timeout for each server is
   * initialized with {@link #DEFAULT_RESOLVER_TIMEOUT}.
   */
  public ExtendedResolver() {
    List<InetSocketAddress> servers = ResolverConfig.getCurrentConfig().servers();
    resolvers.addAll(
        servers.stream()
            .map(
                server -> {
                  Resolver r = new SimpleResolver(server);
                  r.setTimeout(DEFAULT_RESOLVER_TIMEOUT);
                  return new ResolverEntry(r);
                })
            .collect(Collectors.toList()));
  }

  /**
   * Creates a new instance with {@link SimpleResolver}s. The timeout for each server is initialized
   * with {@link #DEFAULT_RESOLVER_TIMEOUT}.
   *
{'fraction_non_alphanumeric': 0.11283027850511783, 'fraction_numerical': 0.0042846941204475126, 'mean_word_length': 2.5128205128205128, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6882256', 'n_tokens_mistral': 3541, 'n_tokens_neox': 3422, 'n_words': 861}