# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

from collections import defaultdict
import json
import logging
import socket
import threading
import traceback
import time
import urllib2
import uuid

from Types.ttypes import TNetworkAddress
from thrift.protocol import TBinaryProtocol
from thrift.server.TServer import TServer
from thrift.transport import TSocket
from thrift.transport import TTransport

import StatestoreService.StatestoreSubscriber as Subscriber
import StatestoreService.StatestoreService as Statestore
from StatestoreService.StatestoreSubscriber import TUpdateStateResponse
from StatestoreService.StatestoreSubscriber import TTopicRegistration
from ErrorCodes.ttypes import TErrorCode
from Status.ttypes import TStatus

from tests.common.environ import specific_build_type_timeout

LOG = logging.getLogger('test_statestore')

# Tests for the statestore. The StatestoreSubscriber class is a skeleton implementation of
# a Python-based statestore subscriber with additional hooks to allow testing. Each
# StatestoreSubscriber runs its own server so that the statestore may contact it.
#
# All tests in this file may be run in parallel. They assume that a statestore instance is
# already running, and is configured with out-of-the-box defaults (as is the case in our
# usual test environment) which govern failure-detector timeouts etc.
#
# These tests do not yet provide sufficient coverage.
#    If no topic entries, do the first and second subscribers always get a callback?
#    Adding topic entries to non-existant topic
#    Test for from_version and to_version behavior
#    Test with many concurrent subscribers
#    Test that only the subscribed-to topics are sent
#    Test that topic deletions take effect correctly.

def get_statestore_subscribers(host='localhost', port=25010):
  response = urllib2.urlopen("http://{0}:{1}/subscribers?json".format(host, port))
  page = response.read()
  return json.loads(page)

STATUS_OK = TStatus(TErrorCode.OK)
DEFAULT_UPDATE_STATE_RESPONSE = TUpdateStateResponse(status=STATUS_OK, topic_updates=[],
                                                     skipped=False)

# IMPALA-3501: the timeout needs to be higher in code coverage builds
WAIT_FOR_FAILURE_TIMEOUT = specific_build_type_timeout(40, code_coverage_build_timeout=60)

class WildcardServerSocket(TSocket.TSocketBase, TTransport.TServerTransportBase):
  """Specialised server socket that binds to a random port at construction"""
  def __init__(self, host=None, port=0):
    self.host = host
    self.handle = None
    self.handle = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    self.handle.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    self.handle.bind(('localhost', 0))
    _, self.port = self.handle.getsockname()

  def listen(self):
    self.handle.listen(128)

  def accept(self):
    client, addr = self.handle.accept()
    result = TSocket.TSocket()
    result.setHandle(client)
    return result


class KillableThreadedServer(TServer):
  """Based on TServer.TThreadedServer, this server may be shutdown (by calling
  shutdown()), after which no new connections may be made. Most of the implementation is
  directly copied from Thrift."""
  def __init__(self, *args, **kwargs):
    TServer.__init__(self, *args)
    self.daemon = kwargs.get("daemon", False)
    self.is_shutdown = False
    self.port = self.serverTransport.port

  def shutdown(self):
    self.is_shutdown = True
    self.serverTransport.close()
    self.wait_until_down()

  def wait_until_up(self, num_tries=10):
    for i in xrange(num_tries):
      cnxn = TSocket.TSocket('localhost', self.port)
      try:
        cnxn.open()
        return
      except Exception, e:
        if i == num_tries - 1: raise
      time.sleep(0.1)

  def wait_until_down(self, num_tries=10):
    for i in xrange(num_tries):
      cnxn = TSocket.TSocket('localhost', self.port)
      try:
        cnxn.open()
        time.sleep(0.1)
      except Exception, e:
        return
    raise Exception("Server did not stop")

  def serve(self):
    self.serverTransport.listen()
    while not self.is_shutdown:
      client = self.serverTransport.accept()
      # Since accept() can take a while, check again if the server is shutdown to avoid
      # starting an unnecessary thread.
      if self.is_shutdown: return
      t = threading.Thread(target=self.handle, args=(client,))
      t.setDaemon(self.daemon)
      t.start()

  def handle(self, client):
    itrans = self.inputTransportFactory.getTransport(client)
    otrans = self.outputTransportFactory.getTransport(client)
    iprot = self.inputProtocolFactory.getProtocol(itrans)
    oprot = self.outputProtocolFactory.getProtocol(otrans)
    try:
      while not self.is_shutdown:
        self.processor.process(iprot, oprot)
    except TTransport.TTransportException, tx:
      pass
    except Exception, x:
      print x

    itrans.close()
    otrans.close()


class StatestoreSubscriber(object):
  """A bare-bones subscriber skeleton. Tests should create a new StatestoreSubscriber(),
  call start() and then register(). The subscriber will run a Thrift server on an unused
  port, and after registration the statestore will call Heartbeat() and UpdateState() via
  RPC. Tests can provide callbacks to the constructor that will be called during those
  RPCs, and this is the easiest way to check that the statestore protocol is being
  correctly followed. Tests should use wait_for_* methods to confirm that some event (like
  an RPC call) has happened asynchronously.

  Since RPC callbacks will execute on a different thread from the main one, any assertions
  there will not trigger a test failure without extra plumbing. What we do is simple: any
  exceptions during an RPC are caught and stored, and the check_thread_exceptions() method
  will re-raise them.

  The methods that may be called by a test deliberately return 'self' to allow for
  chaining, see test_failure_detected() for an example of how this makes the test flow
  more readable."""
  def __init__(self, heartbeat_cb=None, update_cb=None):
    self.heartbeat_event, self.heartbeat_count = threading.Condition(), 0
    # Track the number of updates received per topic.
    self.update_counts = defaultdict(lambda : 0)
    # Variables to notify for updates on each topic.
    self.update_event = threading.Condition()
    self.heartbeat_cb, self.update_cb = heartbeat_cb, update_cb
    self.subscriber_id = "python-test-client-%s" % uuid.uuid1()
    self.exception = None

  def Heartbeat(self, args):
    """Heartbeat RPC handler. Calls heartbeat callback if one exists."""
    self.heartbeat_event.acquire()
    try:
      self.heartbeat_count += 1
      response = Subscriber.THeartbeatResponse()
      if self.heartbeat_cb is not None and self.exception is None:
        try:
          response = self.heartbeat_cb(self, args)
        except Exception, e:
          self.exception = e
      self.heartbeat_event.notify()
    finally:
      self.heartbeat_event.release()
    return response

  def UpdateState(self, args):
    """UpdateState RPC handler. Calls update callback if one exists."""
    self.update_event.acquire()
    try:
      for topic_name in args.topic_deltas: self.update_counts[topic_name] += 1
      response = DEFAULT_UPDATE_STATE_RESPONSE
      if self.update_cb is not None and self.exception is None:
        try:
          response = self.update_cb(self, args)
        except Exception, e:
          # Print the original backtrace so it doesn't get lost.
          traceback.print_exc()
          self.exception = e
      self.update_event.notify()
    finally:
      self.update_event.release()
    return response

  def __init_server(self):
    processor = Subscriber.Processor(self)
    transport = WildcardServerSocket()
    tfactory = TTransport.TBufferedTransportFactory()
    pfactory = TBinaryProtocol.TBinaryProtocolFactory()
    self.server = KillableThreadedServer(processor, transport, tfactory, pfactory,
                                         daemon=True)
    self.server_thread = threading.Thread(target=self.server.serve)
    self.server_thread.setDaemon(True)
    self.server_thread.start()
    self.server.wait_until_up()
    self.port = self.server.port

  def __init_client(self):
    self.client_transport = \
        TTransport.TBufferedTransport(TSocket.TSocket('localhost', 24000))
    self.protocol = TBinaryProtocol.TBinaryProtocol(self.client_transport)
    self.client = Statestore.Client(self.protocol)
    self.client_transport.open()

  def check_thread_exceptions(self):
    """Checks if an exception was raised and stored in a callback thread"""
    if self.exception is not None: raise self.exception

  def kill(self):
    """Closes both the server and client sockets, and waits for the server to become
    unavailable"""
    self.client_transport.close()
    self.server.shutdown()
    return self

  def start(self):
    """Starts a subscriber server, and opens a client to the statestore. Returns only when
    the server is running."""
    self.__init_server()
    self.__init_client()
    return self

  def register(self, topics=None):
    """Call the Register() RPC"""
    if topics is None: topics = []
    request = Subscriber.TRegisterSubscriberRequest(
      topic_registrations=topics,
      subscriber_location=TNetworkAddress("localhost", self.port),
      subscriber_id=self.subscriber_id)
    response = self.client.RegisterSubscriber(request)
    if response.status.status_code == TErrorCode.OK:
      self.registration_id = response.registration_id
    else:
      raise Exception("Registration failed: %s, %s" %
                      (response.status.status_code,
                       '\n'.join(response.status.error_msgs)))
    return self

  def wait_for_heartbeat(self, count=None):
    """Waits for some number of heartbeats. If 'count' is provided, waits until the number
    of heartbeats seen by this subscriber exceeds count, otherwise waits for one further
    heartbeat."""
    self.heartbeat_event.acquire()
    try:
      if count is not None and self.heartbeat_count >= count: return self
      if count is None: count = self.heartbeat_count + 1
      while count > self.heartbeat_count:
        self.check_thread_exceptions()
        last_count = self.heartbeat_count
        self.heartbeat_event.wait(10)
        if last_count == self.heartbeat_count:
          raise Exception("Heartbeat not received within 10s (heartbeat count: %s)" %
                          self.heartbeat_count)
      self.check_thread_exceptions()
      return self
    finally:
      self.heartbeat_event.release()

  def wait_for_update(self, topic_name, count=None):
    """Waits for some number of updates of 'topic_name'. If 'count' is provided, waits
    until the number updates seen by this subscriber exceeds count, otherwise waits
    for one further update."""
    self.update_event.acquire()
    start_time = time.time()
    try:
      if count is not None and self.update_counts[topic_name] >= count: return self
      if count is None: count = self.update_counts[topic_name] + 1
      while count > self.update_counts[topic_name]:
        self.check_thread_exceptions()
        last_count = self.update_counts[topic_name]
        self.update_event.wait(10)
        if (time.time() > start_time + 10 and
            last_count == self.update_counts[topic_name]):
          raise Exception("Update not received for %s within 10s (update count: %s)" %
                          (topic_name, last_count))
      self.check_thread_exceptions()
      return self
    finally:
      self.update_event.release()


  def wait_for_failure(self, timeout=WAIT_FOR_FAILURE_TIMEOUT):
    """Waits until this subscriber no longer appears in the statestore's subscriber
    list. If 'timeout' seconds pass, throws an exception."""
    start = time.time()
    while time.time() - start < timeout:
      subs = [s["id"] for s in get_statestore_subscribers()["subscribers"]]
      if self.subscriber_id not in subs: return self
      time.sleep(0.2)
    raise Exception("Subscriber %s did not fail in %ss" % (self.subscriber_id, timeout))

class TestStatestore():
  def make_topic_update(self, topic_name, key_template="foo", value_template="bar",
                        num_updates=1, clear_topic_entries=False):
    topic_entries = [
      Subscriber.TTopicItem(key=key_template + str(x), value=value_template + str(x))
      for x in xrange(num_updates)]
    return Subscriber.TTopicDelta(topic_name=topic_name,
                                  topic_entries=topic_entries,
                                  is_delta=False,
                                  clear_topic_entries=clear_topic_entries)

  def test_registration_ids_different(self):
    """Test that if a subscriber with the same id registers twice, the registration ID is
    different"""
    sub = StatestoreSubscriber()
    sub.start().register()
    old_reg_id = sub.registration_id
    sub.register()
    assert old_reg_id != sub.registration_id

  def test_receive_heartbeats(self):
    """Smoke test to confirm that heartbeats get sent to a correctly registered
    subscriber"""
    sub = StatestoreSubscriber()
    sub.start().register().wait_for_heartbeat(5)

  def test_receive_updates(self):
    """Test that updates are correctly received when a subscriber alters a topic"""
    topic_name = "topic_delta_%s" % uuid.uuid1()

    def topic_update_correct(sub, args):
      delta = self.make_topic_update(topic_name)
      update_count = sub.update_counts[topic_name]
      if topic_name not in args.topic_deltas:
        # The update doesn't contain our topic.
        pass
      elif update_count == 1:
        return TUpdateStateResponse(status=STATUS_OK, topic_updates=[delta],
                                    skipped=False)
      elif update_count == 2:
        assert len(args.topic_deltas) == 1, args.topic_deltas
        assert args.topic_deltas[topic_name].topic_entries == delta.topic_entries
        assert args.topic_deltas[topic_name].topic_name == delta.topic_name
      elif update_count == 3:
        # After the content-bearing update was processed, the next delta should be empty
        assert len(args.topic_deltas[topic_name].topic_entries) == 0

      return DEFAULT_UPDATE_STATE_RESPONSE

    sub = StatestoreSubscriber(update_cb=topic_update_correct)
    reg = TTopicRegistration(topic_name=topic_name, is_transient=False)
    (
      sub.start()
         .register(topics=[reg])
         .wait_for_update(topic_name, 3)
    )

  def test_update_is_delta(self):
    """Test that the 'is_delta' flag is correctly set. The first update for a topic should
    always not be a delta, and so should all subsequent updates until the subscriber says
    it has not skipped the update."""
    topic_name = "test_update_is_delta_%s" % uuid.uuid1()

    def check_delta(sub, args):
      update_count = sub.update_counts[topic_name]
      if topic_name not in args.topic_deltas:
        # The update doesn't contain our topic.
        pass
      elif update_count == 1:
        assert args.topic_deltas[topic_name].is_delta == False
        delta = self.make_topic_update(topic_name)
        return TUpdateStateResponse(status=STATUS_OK, topic_updates=[delta],
                                    skipped=False)
      elif update_count == 2:
        assert args.topic_deltas[topic_name].is_delta == False
      elif update_count == 3:
        assert args.topic_deltas[topic_name].is_delta == True
        assert len(args.topic_deltas[topic_name].topic_entries) == 0
        assert args.topic_deltas[topic_name].to_version == 1

      return DEFAULT_UPDATE_STATE_RESPONSE

    sub = StatestoreSubscriber(update_cb=check_delta)
    reg = TTopicRegistration(topic_name=topic_name, is_transient=False)
    (
      sub.start()
         .register(topics=[reg])
         .wait_for_update(topic_name, 3)
    )

  def test_skipped(self):
    """Test that skipping an update causes it to be resent"""
    topic_name = "test_skipped_%s" % uuid.uuid1()

    def check_skipped(sub, args):
      # Ignore responses that don't contain our topic.
      if topic_name not in args.topic_deltas: return DEFAULT_UPDATE_STATE_RESPONSE
      update_count = sub.update_counts[topic_name]
      if update_count == 1:
        update = self.make_topic_update(topic_name)
        return TUpdateStateResponse(status=STATUS_OK, topic_updates=[update],
                                    skipped=False)
      # All subsequent updates: set skipped=True and expected the full topic to be resent
      # every time
      assert args.topic_deltas[topic_name].is_delta == False
      assert len(args.topic_deltas[topic_name].topic_entries) == 1
      return TUpdateStateResponse(status=STATUS_OK, skipped=True)

    sub = StatestoreSubscriber(update_cb=check_skipped)
    reg = TTopicRegistration(topic_name=topic_name, is_transient=False)
    (
      sub.start()
         .register(topics=[reg])
         .wait_for_update(topic_name, 3)
    )

  def test_failure_detected(self):
    sub = StatestoreSubscriber()
    topic_name = "test_failure_detected"
    reg = TTopicRegistration(topic_name=topic_name, is_transient=True)
    (
      sub.start()
         .register(topics=[reg])
         .wait_for_update(topic_name, 1)
         .kill()
         .wait_for_failure()
    )

  def test_hung_heartbeat(self):
    """Test for IMPALA-1712: If heartbeats hang (which we simulate by sleeping for five
    minutes) the statestore should time them out every 3s and then eventually fail after
    40s (10 times (3 + 1), where the 1 is the inter-heartbeat delay)"""
    sub = StatestoreSubscriber(heartbeat_cb=lambda sub, args: time.sleep(300))
    topic_name = "test_hung_heartbeat"
    reg = TTopicRegistration(topic_name=topic_name, is_transient=True)
    (
      sub.start()
         .register(topics=[reg])
         .wait_for_update(topic_name, 1)
         .wait_for_failure(timeout=60)
    )

  def test_topic_persistence(self):
    """Test that persistent topic entries survive subscriber failure, but transent topic
    entries are erased when the associated subscriber fails"""
    topic_id = str(uuid.uuid1())
    persistent_topic_name = "test_topic_persistence_persistent_%s" % topic_id
    transient_topic_name = "test_topic_persistence_transient_%s" % topic_id

    def add_entries(sub, args):
      # None of, one or both of the topics may be in the update.
      updates = []
      if (persistent_topic_name in args.topic_deltas and
          sub.update_counts[persistent_topic_name] == 1):
        updates.append(self.make_topic_update(persistent_topic_name))

      if (transient_topic_name in args.topic_deltas and
          sub.update_counts[transient_topic_name] == 1):
        updates.append(self.make_topic_update(transient_topic_name))

      if len(updates) > 0:
        return TUpdateStateResponse(status=STATUS_OK, topic_updates=updates,
                                    skipped=False)
      return DEFAULT_UPDATE_STATE_RESPONSE

    def check_entries(sub, args):
      # None of, one or both of the topics may be in the update.
      if (persistent_topic_name in args.topic_deltas and
          sub.update_counts[persistent_topic_name] == 1):
        assert len(args.topic_deltas[persistent_topic_name].topic_entries) == 1
        # Statestore should not send deletions when the update is not a delta, see
        # IMPALA-1891
        assert args.topic_deltas[persistent_topic_name].topic_entries[0].deleted == False
      if (transient_topic_name in args.topic_deltas and
          sub.update_counts[persistent_topic_name] == 1):
        assert len(args.topic_deltas[transient_topic_name].topic_entries) == 0
      return DEFAULT_UPDATE_STATE_RESPONSE

    reg = [TTopicRegistration(topic_name=persistent_topic_name, is_transient=False),
           TTopicRegistration(topic_name=transient_topic_name, is_transient=True)]

    sub = StatestoreSubscriber(update_cb=add_entries)
    (
      sub.start()
         .register(topics=reg)
         .wait_for_update(persistent_topic_name, 2)
         .wait_for_update(transient_topic_name, 2)
         .kill()
         .wait_for_failure()
    )

    sub2 = StatestoreSubscriber(update_cb=check_entries)
    (
      sub2.start()
          .register(topics=reg)
          .wait_for_update(persistent_topic_name, 1)
          .wait_for_update(transient_topic_name, 1)
    )

  def test_update_with_clear_entries_flag(self):
    """Test that the statestore clears all topic entries when a subscriber
    sets the clear_topic_entries flag in a topic update message (IMPALA-6948)."""
    topic_name = "test_topic_%s" % str(uuid.uuid1())

    def add_entries(sub, args):
      updates = []
      if (topic_name in args.topic_deltas and sub.update_counts[topic_name] == 1):
        updates.append(self.make_topic_update(topic_name, num_updates=2,
            key_template="old"))

      if (topic_name in args.topic_deltas and sub.update_counts[topic_name] == 2):
        updates.append(self.make_topic_update(topic_name, num_updates=1,
            key_template="new", clear_topic_entries=True))

      if len(updates) > 0:
        return TUpdateStateResponse(status=STATUS_OK, topic_updates=updates,
            skipped=False)

      return DEFAULT_UPDATE_STATE_RESPONSE

    def check_entries(sub, args):
      if (topic_name in args.topic_deltas and sub.update_counts[topic_name] == 1):
        assert len(args.topic_deltas[topic_name].topic_entries) == 1
        assert args.topic_deltas[topic_name].topic_entries[0].key == "new0"

      return DEFAULT_UPDATE_STATE_RESPONSE

    reg = [TTopicRegistration(topic_name=topic_name, is_transient=False)]
    sub1 = StatestoreSubscriber(update_cb=add_entries)
    (
      sub1.start()
        .register(topics=reg)
        .wait_for_update(topic_name, 1)
        .kill()
        .wait_for_failure()
        .start()
        .register(topics=reg)
        .wait_for_update(topic_name, 1)
    )

    sub2 = StatestoreSubscriber(update_cb=check_entries)
    (
      sub2.start()
        .register(topics=reg)
        .wait_for_update(topic_name, 2)
    )

  def test_heartbeat_failure_reset(self):
    """Regression test for IMPALA-6785: the heartbeat failure count for the subscriber ID
    should be reset when it resubscribes, not after the first successful heartbeat. Delay
    the heartbeat to force the topic update to finish first."""

    sub = StatestoreSubscriber(heartbeat_cb=lambda sub, args: time.sleep(0.5))
    topic_name = "test_heartbeat_failure_reset"
    reg = TTopicRegistration(topic_name=topic_name, is_transient=True)
    sub.start()
    sub.register(topics=[reg])
    LOG.info("Registered with id {0}".format(sub.subscriber_id))
    sub.wait_for_heartbeat(1)
    sub.kill()
    LOG.info("Killed, waiting for statestore to detect failure via heartbeats")
    sub.wait_for_failure()
    # IMPALA-6785 caused only one topic update to be send. Wait for multiple updates to
    # be received to confirm that the subsequent updates are being scheduled repeatedly.
    target_updates = sub.update_counts[topic_name] + 5
    sub.start()
    sub.register(topics=[reg])
    LOG.info("Re-registered with id {0}, waiting for update".format(sub.subscriber_id))
    sub.wait_for_update(topic_name, target_updates)