# Copyright 2009 Google Inc. 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.

"""Simple DNS server comparison benchmarking tool.

Designed to assist system administrators in selection and prioritization.
"""

__author__ = 'tstromberg@google.com (Thomas Stromberg)'

import Queue
import random
import threading
import time


class BenchmarkThreads(threading.Thread):
  """Benchmark multiple nameservers in parallel."""

  def __init__(self, input_queue, results_queue):
    threading.Thread.__init__(self)
    self.input = input_queue
    self.results = results_queue

  def run(self):
    """Iterate over the queue, processing each item."""
    while not self.input.empty():
      try:
        (ns, request_type, hostname) = self.input.get_nowait()
        # We've moved this here so that it's after all of the random selection goes through.
        if '__RANDOM__' in hostname:
          hostname = hostname.replace('__RANDOM__', str(random.random() * random.randint(0, 99999)))

        (response, duration, error_msg) = ns.TimedRequest(request_type, hostname)
        self.results.put((ns, request_type, hostname, response, duration, error_msg))
      except Queue.Empty:
        return


class Benchmark(object):
  """The main benchmarking class."""

  def __init__(self, nameservers, run_count=2, query_count=30, thread_count=1,
               status_callback=None):
    """Constructor.

    Args:
      nameservers: a list of NameServerData objects
      run_count: How many test-runs to perform on each nameserver (int)
      query_count: How many DNS lookups to test in each test-run (int)
      thread_count: How many benchmark threads to use (int)
      status_callback: Where to send msg() updates to.
    """
    self.query_count = query_count
    self.run_count = run_count
    self.thread_count = thread_count
    self.nameservers = nameservers
    self.results = {}
    self.status_callback = status_callback

  def msg(self, msg, **kwargs):
    if self.status_callback:
      self.status_callback(msg, **kwargs)

  def _CheckForIndexHostsInResults(self, test_records):
    """Check if we have already tested index hosts.

    Args:
      test_records: List of tuples of test records (type, record)

    Returns:
      A list of results that have already been tested
      A list of records that still need to be tested.
    """
    needs_test = []
    index_results = {}
    for test in test_records:
      matched = False
      for ns in self.results:
        for result in self.results[ns][0]:
          hostname, request_type = result[0:2]
          if (request_type, hostname) == test:
            matched = True
            index_results.setdefault(ns, []).append(result)
            # So that we don't include the second results if duplicates exist.
            break
      if not matched:
        needs_test.append(test)
    return (index_results, needs_test)

  def RunIndex(self, test_records):
    """Run index tests using the same mechanism as a standard benchmark."""
    if not test_records:
      print 'No records to test.'
      return None

    index_results, pending_tests = self._CheckForIndexHostsInResults(test_records)
    run_results = self._SingleTestRun(pending_tests)
    for ns in run_results:
      index_results.setdefault(ns, []).extend(run_results[ns])
    return index_results

  def Run(self, test_records=None):
    """Run all test runs for all nameservers."""

    # We don't want to keep stats on how many queries timed out from previous runs.
    for ns in self.nameservers.enabled_servers:
      ns.ResetErrorCounts()

    for _ in range(self.run_count):
      run_results = self._SingleTestRun(test_records)
      for ns in run_results:
        self.results.setdefault(ns, []).append(run_results[ns])
    return self.results

  def _SingleTestRun(self, test_records):
    """Manage and execute a single test-run on all nameservers.

    We used to run all tests for a nameserver, but the results proved to be
    unfair if the bandwidth was suddenly constrained. We now run a test on
    each server before moving on to the next.

    Args:
      test_records: a list of tuples in the form of (request_type, hostname)

    Returns:
      results: A dictionary of tuples, keyed by nameserver.
    """
    input_queue = Queue.Queue()
    shuffled_records = {}
    results = {}
    # Pre-compute the shuffled test records per-nameserver to avoid thread
    # contention.
    for ns in self.nameservers.enabled_servers:
      random.shuffle(test_records)
      shuffled_records[ns.ip] = list(test_records)

    # Feed the pre-computed records into the input queue.
    for i in range(len(test_records)):
      for ns in self.nameservers.enabled_servers:
        (request_type, hostname) = shuffled_records[ns.ip][i]
        input_queue.put((ns, request_type, hostname))

    results_queue = self._LaunchBenchmarkThreads(input_queue)
    errors = []
    while results_queue.qsize():
      (ns, request_type, hostname, response, duration, error_msg) = results_queue.get()
      if error_msg:
        duration = ns.timeout * 1000
        errors.append((ns, error_msg))
      results.setdefault(ns, []).append((hostname, request_type, duration, response, error_msg))

    for (ns, error_msg) in errors:
      self.msg('Error querying %s: %s' % (ns, error_msg))
    return results

  def _LaunchBenchmarkThreads(self, input_queue):
    """Launch and manage the benchmark threads."""
    results_queue = Queue.Queue()
    expected_total = input_queue.qsize()
    threads = []
    for unused_thread_num in range(0, self.thread_count):
      thread = BenchmarkThreads(input_queue, results_queue)
      thread.start()
      threads.append(thread)

    query_count = expected_total / len(self.nameservers.enabled_servers)
    status_message = ('Sending %s queries to %s servers' %
                      (query_count, len(self.nameservers.enabled_servers)))
    while results_queue.qsize() != expected_total:
      self.msg(status_message, count=results_queue.qsize(), total=expected_total)
      time.sleep(0.5)

    self.msg(status_message, count=results_queue.qsize(), total=expected_total)
    for thread in threads:
      thread.join()
    return results_queue

