blob: 90c71b255bc804d48869b1a49ec42486c4f2c551 [file] [log] [blame]
# 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.
import asyncio
import contextlib
from collections import OrderedDict, deque
import random
import pytest
from pyignite import GenericObjectMeta, AioClient, Client
from pyignite.aio_cache import AioCache
from pyignite.connection import Connection, AioConnection
from pyignite.constants import PROTOCOL_BYTE_ORDER
from pyignite.datatypes import String, LongObject
from pyignite.datatypes.cache_config import CacheMode
from pyignite.datatypes.prop_codes import PROP_NAME, PROP_BACKUPS_NUMBER, PROP_CACHE_KEY_CONFIGURATION, PROP_CACHE_MODE
from tests.util import wait_for_condition, wait_for_condition_async, start_ignite, kill_process_tree
try:
from contextlib import asynccontextmanager
except ImportError:
from async_generator import asynccontextmanager
requests = deque()
old_send = Connection.send
old_send_async = AioConnection._send
def patched_send(self, *args, **kwargs):
"""Patched send function that push to queue idx of server to which request is routed."""
buf = args[0]
if buf and len(buf) >= 6:
op_code = int.from_bytes(buf[4:6], byteorder=PROTOCOL_BYTE_ORDER)
# Filter only caches operation.
if 1000 <= op_code < 1100:
requests.append(self.port % 100)
return old_send(self, *args, **kwargs)
async def patched_send_async(self, *args, **kwargs):
"""Patched send function that push to queue idx of server to which request is routed."""
buf = args[0]
if buf and len(buf) >= 6:
op_code = int.from_bytes(buf[4:6], byteorder=PROTOCOL_BYTE_ORDER)
# Filter only caches operation.
if 1000 <= op_code < 1100:
requests.append(self.port % 100)
return await old_send_async(self, *args, **kwargs)
def setup_function():
requests.clear()
Connection.send = patched_send
AioConnection._send = patched_send_async
def teardown_function():
Connection.send = old_send
AioConnection.send = old_send_async
def wait_for_affinity_distribution(cache, key, node_idx, timeout=30):
real_node_idx = 0
def check_grid_idx():
nonlocal real_node_idx
try:
cache.get(key)
real_node_idx = requests.pop()
except (OSError, IOError):
return False
return real_node_idx == node_idx
res = wait_for_condition(check_grid_idx, timeout=timeout)
if not res:
raise TimeoutError(f"failed to wait for affinity distribution, expected node_idx {node_idx},"
f"got {real_node_idx} instead")
async def wait_for_affinity_distribution_async(cache, key, node_idx, timeout=30):
real_node_idx = 0
async def check_grid_idx():
nonlocal real_node_idx
try:
await cache.get(key)
real_node_idx = requests.pop()
except (OSError, IOError):
return False
return real_node_idx == node_idx
res = await wait_for_condition_async(check_grid_idx, timeout=timeout)
if not res:
raise TimeoutError(f"failed to wait for affinity distribution, expected node_idx {node_idx},"
f"got {real_node_idx} instead")
@pytest.mark.parametrize("key,grid_idx", [(1, 1), (2, 2), (3, 3), (4, 1), (5, 1), (6, 2), (11, 1), (13, 1), (19, 1)])
@pytest.mark.parametrize("backups", [0, 1, 2, 3])
def test_cache_operation_on_primitive_key_routes_request_to_primary_node(request, key, grid_idx, backups, client):
cache = client.get_or_create_cache({
PROP_NAME: request.node.name + str(backups),
PROP_BACKUPS_NUMBER: backups,
})
try:
__perform_operations_on_primitive_key(client, cache, key, grid_idx)
finally:
cache.destroy()
@pytest.mark.parametrize("key,grid_idx", [(1, 1), (2, 2), (3, 3), (4, 1), (5, 1), (6, 2), (11, 1), (13, 1), (19, 1)])
@pytest.mark.parametrize("backups", [0, 1, 2, 3])
@pytest.mark.asyncio
async def test_cache_operation_on_primitive_key_routes_request_to_primary_node_async(
request, key, grid_idx, backups, async_client):
cache = await async_client.get_or_create_cache({
PROP_NAME: request.node.name + str(backups),
PROP_BACKUPS_NUMBER: backups,
})
try:
await __perform_operations_on_primitive_key(async_client, cache, key, grid_idx)
finally:
await cache.destroy()
def __perform_operations_on_primitive_key(client, cache, key, grid_idx):
operations = [
('get', 1), ('put', 2), ('replace', 2), ('clear_key', 1), ('contains_key', 1), ('get_and_put', 2),
('get_and_put_if_absent', 2), ('put_if_absent', 2), ('get_and_remove', 1), ('get_and_replace', 2),
('remove_key', 1), ('remove_if_equals', 2), ('replace', 2), ('replace_if_equals', 3)
]
def inner():
cache.put(key, key)
wait_for_affinity_distribution(cache, key, grid_idx)
for op_name, param_nums in operations:
op = getattr(cache, op_name)
args = [random.randint(-100, 100) for _ in range(0, param_nums - 1)]
op(key, *args)
assert requests.pop() == grid_idx
async def inner_async():
await cache.put(key, key)
await wait_for_affinity_distribution_async(cache, key, grid_idx)
for op_name, param_nums in operations:
op = getattr(cache, op_name)
args = [random.randint(-100, 100) for _ in range(0, param_nums - 1)]
await op(key, *args)
assert requests.pop() == grid_idx
return inner_async() if isinstance(client, AioClient) else inner()
@pytest.mark.skip(reason="Custom key objects are not supported yet")
def test_cache_operation_on_complex_key_routes_request_to_primary_node():
pass
@pytest.mark.parametrize("key,grid_idx", [(1, 2), (2, 1), (3, 1), (4, 2), (5, 2), (6, 3)])
@pytest.mark.skip(reason="Custom key objects are not supported yet")
def test_cache_operation_on_custom_affinity_key_routes_request_to_primary_node(request, client, key, grid_idx):
class AffinityTestType1(
metaclass=GenericObjectMeta,
type_name='AffinityTestType1',
schema=OrderedDict([
('test_str', String),
('test_int', LongObject)
])
):
pass
cache_config = {
PROP_NAME: request.node.name,
PROP_CACHE_KEY_CONFIGURATION: [
{
'type_name': 'AffinityTestType1',
'affinity_key_field_name': 'test_int',
},
],
}
cache = client.create_cache(cache_config)
# noinspection PyArgumentList
key_obj = AffinityTestType1(
test_str="abc",
test_int=key
)
cache.put(key_obj, 1)
cache.put(key_obj, 2)
assert requests.pop() == grid_idx
client_routed_connection_string = [('127.0.0.1', 10800 + idx) for idx in range(1, 5)]
@pytest.fixture
def client_routed():
client = Client(partition_aware=True)
try:
client.connect(client_routed_connection_string)
yield client
finally:
client.close()
@pytest.fixture
def client_routed_cache(client_routed, request):
yield client_routed.get_or_create_cache(request.node.name)
@pytest.fixture
async def async_client_routed():
client = AioClient(partition_aware=True)
try:
await client.connect(client_routed_connection_string)
yield client
finally:
await client.close()
@pytest.fixture
async def async_client_routed_cache(async_client_routed, request):
cache = await async_client_routed.get_or_create_cache(request.node.name)
yield cache
def test_cache_operation_routed_to_new_cluster_node(client_routed_cache):
__perform_cache_operation_routed_to_new_node(client_routed_cache)
@pytest.mark.asyncio
async def test_cache_operation_routed_to_new_cluster_node_async(async_client_routed_cache):
await __perform_cache_operation_routed_to_new_node(async_client_routed_cache)
def __perform_cache_operation_routed_to_new_node(cache):
key = 12
def inner():
wait_for_affinity_distribution(cache, key, 3)
cache.put(key, key)
cache.put(key, key)
assert requests.pop() == 3
srv = start_ignite(idx=4)
try:
# Wait for rebalance and partition map exchange
wait_for_affinity_distribution(cache, key, 4)
# Response is correct and comes from the new node
res = cache.get_and_remove(key)
assert res == key
assert requests.pop() == 4
finally:
kill_process_tree(srv.pid)
async def inner_async():
await wait_for_affinity_distribution_async(cache, key, 3)
await cache.put(key, key)
await cache.put(key, key)
assert requests.pop() == 3
srv = start_ignite(idx=4)
try:
# Wait for rebalance and partition map exchange
await wait_for_affinity_distribution_async(cache, key, 4)
# Response is correct and comes from the new node
res = await cache.get_and_remove(key)
assert res == key
assert requests.pop() == 4
finally:
kill_process_tree(srv.pid)
return inner_async() if isinstance(cache, AioCache) else inner()
@pytest.fixture
def replicated_cache(request, client):
cache = client.get_or_create_cache({
PROP_NAME: request.node.name,
PROP_CACHE_MODE: CacheMode.REPLICATED,
})
try:
yield cache
finally:
cache.destroy()
@pytest.fixture
async def async_replicated_cache(request, async_client):
cache = await async_client.get_or_create_cache({
PROP_NAME: request.node.name,
PROP_CACHE_MODE: CacheMode.REPLICATED,
})
try:
yield cache
finally:
await cache.destroy()
def test_replicated_cache_operation_routed_to_random_node(replicated_cache):
verify_random_node(replicated_cache)
@pytest.mark.asyncio
async def test_replicated_cache_operation_routed_to_random_node_async(async_replicated_cache):
await verify_random_node(async_replicated_cache)
def verify_random_node(cache):
key = 1
def inner():
cache.put(key, key)
idx1 = requests.pop()
idx2 = idx1
# Try 10 times - random node may end up being the same
for _ in range(1, 10):
cache.put(key, key)
idx2 = requests.pop()
if idx2 != idx1:
break
assert idx1 != idx2
async def inner_async():
await cache.put(key, key)
idx1 = requests.pop()
idx2 = idx1
# Try 10 times - random node may end up being the same
for _ in range(1, 10):
await cache.put(key, key)
idx2 = requests.pop()
if idx2 != idx1:
break
assert idx1 != idx2
return inner_async() if isinstance(cache, AioCache) else inner()
@contextlib.contextmanager
def create_caches(client):
caches = []
try:
caches = [client.create_cache(f'test_cache_{i}') for i in range(0, 10)]
yield caches
finally:
for cache in caches:
try:
cache.destroy()
except: # noqa: 13
cache.destroy() # Retry if connection failed.
pass
@asynccontextmanager
async def create_caches_async(client):
caches = []
try:
caches = await asyncio.gather(*[client.create_cache(f'test_cache_{i}') for i in range(0, 10)])
yield caches
finally:
for cache in caches:
try:
await cache.destroy()
except: # noqa: 13
await cache.destroy() # Retry if connection failed.
pass
def test_new_registered_cache_affinity(client):
with create_caches(client) as caches:
key = 12
test_cache = random.choice(caches)
test_cache.put(key, key)
wait_for_affinity_distribution(test_cache, key, 3)
caches.append(client.create_cache('new_cache'))
for cache in caches:
cache.get(key)
assert requests.pop() == 3
@pytest.mark.asyncio
async def test_new_registered_cache_affinity_async(async_client):
async with create_caches_async(async_client) as caches:
key = 12
test_cache = random.choice(caches)
await test_cache.put(key, key)
await wait_for_affinity_distribution_async(test_cache, key, 3)
caches.append(await async_client.create_cache('new_cache'))
for cache in caches:
await cache.get(key)
assert requests.pop() == 3
def test_all_registered_cache_updated_on_new_server(client_routed):
with create_caches(client_routed) as caches:
key = 12
test_cache = random.choice(caches)
wait_for_affinity_distribution(test_cache, key, 3)
test_cache.put(key, key)
assert requests.pop() == 3
srv = start_ignite(idx=4)
try:
# Wait for rebalance and partition map exchange
wait_for_affinity_distribution(test_cache, key, 4)
for cache in caches:
cache.get(key)
assert requests.pop() == 4
finally:
kill_process_tree(srv.pid)
@pytest.mark.asyncio
async def test_all_registered_cache_updated_on_new_server_async(async_client_routed):
async with create_caches_async(async_client_routed) as caches:
key = 12
test_cache = random.choice(caches)
await wait_for_affinity_distribution_async(test_cache, key, 3)
await test_cache.put(key, key)
assert requests.pop() == 3
srv = start_ignite(idx=4)
try:
# Wait for rebalance and partition map exchange
await wait_for_affinity_distribution_async(test_cache, key, 4)
for cache in caches:
await cache.get(key)
assert requests.pop() == 4
finally:
kill_process_tree(srv.pid)
@pytest.mark.asyncio
async def test_update_affinity_concurrently(async_client):
async with create_caches_async(async_client) as caches:
key = 12
await asyncio.gather(*[cache.put(key, key) for cache in caches])
for cache in caches:
await cache.get(key)
assert requests.pop() == 3