# 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
from random import randint
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 build_flavor_timeout
from tests.common.skip import SkipIfDockerizedCluster

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 = build_flavor_timeout(40, code_coverage_build_timeout=60)
WAIT_FOR_HEARTBEAT_TIMEOUT = build_flavor_timeout(
    40, code_coverage_build_timeout=60)
WAIT_FOR_UPDATE_TIMEOUT = build_flavor_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()
    # The processor contains a reference to a StatestoreSubscriber. Clean up that
    # reference to avoid a circular reference that would prevent object deletion.
    self.processor = None

  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 __enter__(self):
    return self

  def __exit__(self, *args):
    self.kill()
    self.wait_for_failure()

  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"""
    if self.client_transport:
      self.client_transport.close()
    if self.server:
      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(WAIT_FOR_HEARTBEAT_TIMEOUT)
        if last_count == self.heartbeat_count:
          raise Exception(
              "Heartbeat not received within {0}s (heartbeat count: {1})".format(
                WAIT_FOR_HEARTBEAT_TIMEOUT, 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(WAIT_FOR_UPDATE_TIMEOUT)
        if (time.time() > start_time + WAIT_FOR_UPDATE_TIMEOUT and
              last_count == self.update_counts[topic_name]):
          raise Exception(
              "Update not received for {0} within {1} (update count: {2})".format(
                topic_name, WAIT_FOR_UPDATE_TIMEOUT, 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 True:
      subs = [s["id"] for s in get_statestore_subscribers()["subscribers"]]
      if self.subscriber_id not in subs: return self
      if time.time() - start > timeout:
        raise Exception("Subscriber {0} did not fail in {1}s".format(
            self.subscriber_id, timeout))
      time.sleep(0.2)


@SkipIfDockerizedCluster.statestore_not_exposed
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"""
    with StatestoreSubscriber() as sub:
      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"""
    with StatestoreSubscriber() as sub:
      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

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

  def test_filter_prefix(self):
    topic_name = "topic_delta_%s" % uuid.uuid1()

    def topic_update_correct(sub, args):
      foo_delta = self.make_topic_update(topic_name, num_updates=1)
      bar_delta = self.make_topic_update(topic_name, num_updates=2, key_template='bar')

      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:
        # Send some values with both prefixes.
        return TUpdateStateResponse(status=STATUS_OK,
                                    topic_updates=[foo_delta, bar_delta],
                                    skipped=False)
      elif update_count == 2:
        # We should only get the 'bar' entries back.
        assert len(args.topic_deltas) == 1, args.topic_deltas
        assert args.topic_deltas[topic_name].topic_entries == bar_delta.topic_entries
        assert args.topic_deltas[topic_name].topic_name == bar_delta.topic_name
      elif update_count == 3:
        # Send some more updates that only have 'foo' prefixes.
        return TUpdateStateResponse(status=STATUS_OK,
                                    topic_updates=[foo_delta],
                                    skipped=False)
      elif update_count == 4:
        # We shouldn't see any entries from the above update, but we should still see
        # the version number change due to the new entries in the topic.
        assert len(args.topic_deltas[topic_name].topic_entries) == 0
        assert args.topic_deltas[topic_name].from_version == 3
        assert args.topic_deltas[topic_name].to_version == 4
      elif update_count == 5:
        # After the content-bearing update was processed, the next delta should be empty
        assert len(args.topic_deltas[topic_name].topic_entries) == 0
        assert args.topic_deltas[topic_name].from_version == 4
        assert args.topic_deltas[topic_name].to_version == 4

      return DEFAULT_UPDATE_STATE_RESPONSE

    with StatestoreSubscriber(update_cb=topic_update_correct) as sub:
      reg = TTopicRegistration(topic_name=topic_name, is_transient=False,
                               filter_prefix="bar")
      (
        sub.start()
           .register(topics=[reg])
           .wait_for_update(topic_name, 5)
      )

  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

    with StatestoreSubscriber(update_cb=check_delta) as sub:
      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)

    with StatestoreSubscriber(update_cb=check_skipped) as sub:
      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):
    with StatestoreSubscriber() as sub:
      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)"""
    with StatestoreSubscriber(heartbeat_cb=lambda sub, args: time.sleep(300)) as sub:
      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_intermittent_hung_heartbeats(self):
    """Heartbeats that occasionally time out should not cause a failure to be detected."""
    heartbeat_count = [0]  # Use array to allow mutating from inside callback.

    def heartbeat_cb(sub, args):
      heartbeat_count[0] += 1
      # Delay every second heartbeat.
      if (heartbeat_count[0] % 2 == 1):
        time.sleep(4)
      return Subscriber.THeartbeatResponse()

    with StatestoreSubscriber(heartbeat_cb=heartbeat_cb) as sub:
      topic_name = "test_intermittent_hung_heartbeats"
      reg = TTopicRegistration(topic_name=topic_name, is_transient=True)
      (
        sub.start()
           .register(topics=[reg])
           .wait_for_update(topic_name, 30)
           .kill()
           .wait_for_failure()
       )

  def test_slow_subscriber(self):
    """Test for IMPALA-6644: This test kills a healthy subscriber and sleeps for multiple
    intervals of about 1 second each, this lets the heartbeats to the subscriber fail.
    It polls the subscribers page of the statestore to ensure that the
    'secs_since_heartbeat' field is updated with an acceptable value. This test only
    checks for a strictly increasing value since the actual value of time might depend
    on the system load. It stops polling the page once the subscriber is removed from
    the set of active subscribers. It also checks that a valid heartbeat record of the
    subscriber is found at least once."""
    sub = StatestoreSubscriber()
    sub.start().register().wait_for_heartbeat(1)
    sub.kill()
    # secs_since_heartbeat is initially unknown.
    secs_since_heartbeat = -1
    valid_heartbeat_record = False
    while secs_since_heartbeat != 0:
      sleep_start_time = time.time()
      while time.time() - sleep_start_time < 1:
        time.sleep(0.1)
      prev_secs_since_heartbeat = secs_since_heartbeat
      secs_since_heartbeat = 0
      subscribers = get_statestore_subscribers()["subscribers"]
      for s in subscribers:
        if str(s["id"]) == sub.subscriber_id:
          secs_since_heartbeat = float(s["secs_since_heartbeat"])
          assert (secs_since_heartbeat > prev_secs_since_heartbeat)
          valid_heartbeat_record = True
    assert valid_heartbeat_record

  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)]

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

    with StatestoreSubscriber(update_cb=check_entries) as sub2:
      (
         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)]
    with StatestoreSubscriber(update_cb=add_entries) as sub1:
      (
        sub1.start()
            .register(topics=reg)
            .wait_for_update(topic_name, 1)
            .kill()
            .wait_for_failure()
            .start()
            .register(topics=reg)
            .wait_for_update(topic_name, 2)
      )

    with StatestoreSubscriber(update_cb=check_entries) as sub2:
      (
        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."""

    with StatestoreSubscriber(heartbeat_cb=lambda sub, args: time.sleep(0.5)) as sub:
      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)

  def test_min_subscriber_topic_version(self):
    self._do_test_min_subscriber_topic_version(False)

  def test_min_subscriber_topic_version_with_straggler(self):
    self._do_test_min_subscriber_topic_version(True)

  def _do_test_min_subscriber_topic_version(self, simulate_straggler):
    """Implementation of test that the 'min_subscriber_topic_version' flag is correctly
    set when requested. This tests runs two subscribers concurrently and tracks the
    minimum version each has processed. If 'simulate_straggler' is true, one subscriber
    rejects updates so that its version is not advanced."""
    topic_name = "test_min_subscriber_topic_version_%s" % uuid.uuid1()

    # This lock is held while processing the update to protect last_to_versions.
    update_lock = threading.Lock()
    last_to_versions = {}
    TOTAL_SUBSCRIBERS = 2
    def callback(sub, args, is_producer, sub_name):
      """Callback for subscriber to verify min_subscriber_topic_version behaviour.
      If 'is_producer' is true, this acts as the producer, otherwise it acts as the
      consumer. 'sub_name' is a name used to index into last_to_versions."""
      if topic_name not in args.topic_deltas:
        # The update doesn't contain our topic.
        pass
      with update_lock:
        LOG.info("{0} got update {1}".format(sub_name,
            repr(args.topic_deltas[topic_name])))
        LOG.info("Versions: {0}".format(last_to_versions))
        to_version = args.topic_deltas[topic_name].to_version
        from_version = args.topic_deltas[topic_name].from_version
        min_subscriber_topic_version = \
            args.topic_deltas[topic_name].min_subscriber_topic_version

        if is_producer:
          assert min_subscriber_topic_version is not None
          assert (to_version == 0 and min_subscriber_topic_version == 0) or\
              min_subscriber_topic_version < to_version,\
              "'to_version' hasn't been created yet by this subscriber."
          # Only validate version once all subscribers have processed an update.
          if len(last_to_versions) == TOTAL_SUBSCRIBERS:
            min_to_version = min(last_to_versions.values())
            assert min_subscriber_topic_version <= min_to_version,\
                "The minimum subscriber topic version seen by the producer cannot get " +\
                "ahead of the minimum version seem by the consumer, by definition."
            assert min_subscriber_topic_version >= min_to_version - 2,\
                "The min topic version can be two behind the last version seen by " + \
                "this subscriber because the updates for both subscribers are " + \
                "prepared in parallel and because it's possible that the producer " + \
                "processes two updates in-between consumer updates. This is not " + \
                "absolute but depends on updates not being delayed a large amount."
        else:
          # Consumer did not request topic version.
          assert min_subscriber_topic_version is None

        # Check the 'to_version' and update 'last_to_versions'.
        last_to_version = last_to_versions.get(sub_name, 0)
        if to_version > 0:
          # Non-empty update.
          assert from_version == last_to_version
        # Stragglers should accept the first update then skip later ones.
        skip_update = simulate_straggler and not is_producer and last_to_version > 0
        if not skip_update: last_to_versions[sub_name] = to_version

        if is_producer:
          delta = self.make_topic_update(topic_name)
          return TUpdateStateResponse(status=STATUS_OK, topic_updates=[delta],
                                      skipped=False)
        elif skip_update:
          return TUpdateStateResponse(status=STATUS_OK, topic_updates=[], skipped=True)
        else:
          return DEFAULT_UPDATE_STATE_RESPONSE

    # Two concurrent subscribers, which pushes out updates and checks the minimum
    # version, the other which just consumes the updates.
    def producer_callback(sub, args): return callback(sub, args, True, "producer")
    def consumer_callback(sub, args): return callback(sub, args, False, "consumer")
    with StatestoreSubscriber(update_cb=consumer_callback) as consumer_sub:
      with StatestoreSubscriber(update_cb=producer_callback) as producer_sub:
        consumer_reg = TTopicRegistration(topic_name=topic_name, is_transient=True)
        producer_reg = TTopicRegistration(topic_name=topic_name, is_transient=True,
            populate_min_subscriber_topic_version=True)
        NUM_UPDATES = 6
        (
          consumer_sub.start()
                      .register(topics=[consumer_reg])
        )
        (
          producer_sub.start()
                      .register(topics=[producer_reg])
                      .wait_for_update(topic_name, NUM_UPDATES)
        )
        consumer_sub.wait_for_update(topic_name, NUM_UPDATES)

  def test_transient_entry_removal_race(self):
    """IMPALA-7306: transient entries were not deleted if the subscriber is unregistered
    while it is in the middle of a callback. This test exercises that case by blocking
    the update callback so that it is still running when the statestore unregisters the
    subscriber for failed heartbeats. It also confirms that non-transient entries are not
    removed."""
    transient_topic_name = "test_transient_entry_removal_race_transient"
    non_transient_topic_name = "test_transient_entry_removal_race_non_transient"
    topic_regs = [TTopicRegistration(topic_name=transient_topic_name, is_transient=True),
          TTopicRegistration(topic_name=non_transient_topic_name, is_transient=False)]
    # The heartbeat timeout is 3s, so sleep for long enough for it to expire
    HEARTBEAT_DELAY = 10

    def delayed_heartbeat(sub, args):
      LOG.info("Heartbeat callback called")
      time.sleep(HEARTBEAT_DELAY)
      LOG.debug("Heartbeat callback about to return")

    def add_transient_entries_after_hb_failure(sub, args):
      LOG.info("Update callback called")
      # Add an additional delay so that this returns after the heartbeat.
      time.sleep(WAIT_FOR_FAILURE_TIMEOUT)
      updates = [self.make_topic_update(transient_topic_name, "k", "v"),
                 self.make_topic_update(non_transient_topic_name, "k", "v")]
      LOG.debug("Update callback about to return")
      return TUpdateStateResponse(status=STATUS_OK, topic_updates=updates, skipped=False)

    # Subscriber with delay creates a transient entry, which should not be added since
    # the subscriber failed and was unregistered.
    with StatestoreSubscriber(heartbeat_cb=delayed_heartbeat,
                              update_cb=add_transient_entries_after_hb_failure) as sub:
      # Wait for the first update (which should happen after failure), then confirm
      # that the failure occurred.
      (
        sub.start()
           .register(topics=topic_regs)
           .wait_for_update(transient_topic_name, 1)
           .wait_for_failure(timeout=WAIT_FOR_FAILURE_TIMEOUT)
       )

    def verify_transient_entry_removed(sub, args):
      transient_delta = args.topic_deltas[transient_topic_name]
      assert len(transient_delta.topic_entries) == 0, args
      non_transient_delta = args.topic_deltas[non_transient_topic_name]
      # Non-transient update should include topic that was not removed
      assert len(non_transient_delta.topic_entries) == 1, args
      entry = non_transient_delta.topic_entries[0]
      assert entry.key == "k0"
      assert entry.value == "v0"
      assert not entry.deleted
      # Skip updates so that statestore will re-send non-transient entries and the above
      # assertions remain valid on subsequent callbacks.
      return TUpdateStateResponse(status=STATUS_OK, topic_updates=[], skipped=True)

    # Verify that the transient entry for the failed subscriber is not present.
    with StatestoreSubscriber(update_cb=verify_transient_entry_removed) as sub:
      (
        sub.start()
           .register(topics=topic_regs)
           .wait_for_update(transient_topic_name, 1)
       )