|
|
import uuidimport threadingimport timefrom unittest.mock import patch
from kazoo.exceptions import LockTimeoutfrom kazoo.testing import KazooTestCasefrom kazoo.recipe.partitioner import PartitionState
class SlowLockMock: """Emulates a slow ZooKeeper lock."""
default_delay_time = 3
def __init__(self, client, lock, delay_time=None): self._client = client self._lock = lock self.delay_time = ( self.default_delay_time if delay_time is None else delay_time )
def acquire(self, timeout=None): sleep = self._client.handler.sleep_func sleep(self.delay_time)
if timeout is None: return self._lock.acquire()
start_time = time.time()
while time.time() - start_time < timeout: if self._lock.acquire(False): return True
sleep(0.1)
raise LockTimeout("Mocked slow lock has timed out.")
def release(self): self._lock.release()
class KazooPartitionerTests(KazooTestCase): @staticmethod def make_event(): return threading.Event()
def setUp(self): super(KazooPartitionerTests, self).setUp() self.path = "/" + uuid.uuid4().hex self.__partitioners = []
def test_party_of_one(self): self.__create_partitioner(size=3) self.__wait_for_acquire() self.__assert_state(PartitionState.ACQUIRED) self.__assert_partitions([0, 1, 2]) self.__finish()
def test_party_of_two(self): for i in range(2): self.__create_partitioner(size=2, identifier=str(i))
self.__wait_for_acquire() self.__assert_partitions([0], [1])
self.__partitioners[0].finish() self.__wait() assert self.__partitioners[1].release self.__partitioners[1].finish()
def test_party_expansion(self): for i in range(2): self.__create_partitioner(size=3, identifier=str(i))
self.__wait_for_acquire() self.__assert_state(PartitionState.ACQUIRED) self.__assert_partitions([0, 2], [1])
for partitioner in self.__partitioners: partitioner.state_change_event.clear()
# Add another partition, wait till they settle self.__create_partitioner(size=3, identifier="2") self.__wait()
self.__assert_state( PartitionState.RELEASE, partitioners=self.__partitioners[:-1] ) for partitioner in self.__partitioners[-1]: assert partitioner.state_change_event.is_set() self.__release(self.__partitioners[:-1])
self.__wait_for_acquire() self.__assert_partitions([0], [1], [2])
self.__finish()
def test_more_members_than_set_items(self): for i in range(2): self.__create_partitioner(size=1, identifier=str(i))
self.__wait_for_acquire() self.__assert_state(PartitionState.ACQUIRED) self.__assert_partitions([0], [])
self.__finish()
def test_party_session_failure(self): partitioner = self.__create_partitioner(size=3) self.__wait_for_acquire() assert partitioner.state == PartitionState.ACQUIRED # simulate session failure partitioner._fail_out() partitioner.release_set() assert partitioner.failed is True
def test_connection_loss(self): self.__create_partitioner(identifier="0", size=3) self.__create_partitioner(identifier="1", size=3)
self.__wait_for_acquire() self.__assert_state(PartitionState.ACQUIRED) self.__assert_partitions([0, 2], [1])
# Emulate connection loss self.lose_connection(self.make_event) self.__assert_state(PartitionState.RELEASE) self.__release()
# Check that partitioners settle after connection loss self.__wait_for_acquire() self.__assert_state(PartitionState.ACQUIRED) self.__assert_partitions([0, 2], [1])
# Check that partitioners react on new events after connection loss self.__create_partitioner(identifier="2", size=3) self.__wait()
self.__assert_state( PartitionState.RELEASE, partitioners=self.__partitioners[:-1] ) self.__release(partitioners=self.__partitioners[:-1]) self.__wait_for_acquire() self.__assert_state(PartitionState.ACQUIRED) self.__assert_partitions([0], [1], [2])
def test_race_condition_new_partitioner_during_the_lock(self): locks = {}
def get_lock(path): lock = locks.setdefault(path, self.client.handler.lock_object()) return SlowLockMock(self.client, lock)
with patch.object(self.client, "Lock", side_effect=get_lock): # Create first partitioner. It will start to acquire the set # members. self.__create_partitioner(identifier="0", size=2)
# Wait until the first partitioner has acquired first lock and # started to acquire the second lock. self.client.handler.sleep_func(SlowLockMock.default_delay_time + 1)
# Create the second partitioner a the time when the first # partitioner is in the process of acquiring the lock that should # belong to the second partitioner. self.__create_partitioner(identifier="1", size=2)
# The first partitioner should acquire the both locks but then it # must notice that the party has changed and it must reacquire # the set. No deadlocks must happen. self.__wait_for_acquire()
self.__assert_state(PartitionState.ACQUIRED) self.__assert_partitions([0], [1])
def test_race_condition_new_partitioner_steals_the_lock(self): locks = {}
def get_lock(path): new_lock = self.client.handler.lock_object() lock = locks.setdefault(path, new_lock)
if lock is new_lock: # The first partitioner will be delayed delay_time = SlowLockMock.default_delay_time else: # The second partitioner won't be delayed delay_time = 0
return SlowLockMock(self.client, lock, delay_time=delay_time)
with patch.object(self.client, "Lock", side_effect=get_lock): # Create first partitioner. It will start to acquire the set # members. self.__create_partitioner(identifier="0", size=2)
# Wait until the first partitioner has acquired first lock and # started to acquire the second lock. self.client.handler.sleep_func(SlowLockMock.default_delay_time + 1)
# Create the second partitioner a the time when the first # partitioner is in the process of acquiring the lock that should # belong to the second partitioner. The second partitioner should # steal the lock because it won't be delayed. self.__create_partitioner(identifier="1", size=2)
# The first partitioner should fail to acquire the second lock and # must notice that the party has changed and it must reacquire the # set. No deadlocks must happen. self.__wait_for_acquire()
self.__assert_state(PartitionState.ACQUIRED) self.__assert_partitions([0], [1])
def __create_partitioner(self, size, identifier=None): partitioner = self.client.SetPartitioner( self.path, set=range(size), time_boundary=0.2, identifier=identifier, ) self.__partitioners.append(partitioner) return partitioner
def __wait_for_acquire(self): for partitioner in self.__partitioners: partitioner.wait_for_acquire(14)
def __assert_state(self, state, partitioners=None): if partitioners is None: partitioners = self.__partitioners
for partitioner in partitioners: assert partitioner.state == state
def __assert_partitions(self, *partitions): assert len(partitions) == len(self.__partitioners) for partitioner, own_partitions in zip( self.__partitioners, partitions ): assert list(partitioner) == own_partitions
def __wait(self): time.sleep(0.1)
def __release(self, partitioners=None): if partitioners is None: partitioners = self.__partitioners
for partitioner in partitioners: partitioner.release_set()
def __finish(self): for partitioner in self.__partitioners: partitioner.finish()
|
