wehub-resource-sync
1164 lines
40 KiB
Python
1164 lines
40 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Unit tests for RegistryTree with fine-grained locking.
|
|
|
|
Tests verify that the object-based locking mechanism provides:
|
|
1. Thread-safe operations on instances, workers, and KV stores
|
|
2. Concurrent operations on different instances don't block each other
|
|
3. Data consistency under high concurrency
|
|
"""
|
|
|
|
# Standard
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import threading
|
|
import time
|
|
|
|
# First Party
|
|
from lmcache.v1.cache_controller.message import BatchedKVOperationMsg, KVOpEvent, OpType
|
|
from lmcache.v1.cache_controller.utils import (
|
|
InstanceNode,
|
|
RegistryTree,
|
|
WorkerNode,
|
|
)
|
|
|
|
|
|
class TestWorkerNodeLocking:
|
|
"""Test WorkerNode's internal locking for kv_store and seq_tracker."""
|
|
|
|
def test_concurrent_admit_kv(self):
|
|
"""Test concurrent KV admission on the same worker."""
|
|
worker = WorkerNode(
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
num_threads = 10
|
|
keys_per_thread = 100
|
|
|
|
def admit_keys(thread_id):
|
|
for i in range(keys_per_thread):
|
|
key = thread_id * keys_per_thread + i
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=worker.worker_id,
|
|
location=location,
|
|
operations=[
|
|
KVOpEvent(
|
|
op_type=OpType.ADMIT,
|
|
key=key,
|
|
seq_num=thread_id * keys_per_thread + i,
|
|
)
|
|
],
|
|
)
|
|
worker.handle_batched_kv_operations(msg)
|
|
|
|
threads = [
|
|
threading.Thread(target=admit_keys, args=(i,)) for i in range(num_threads)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# All keys should be admitted
|
|
assert worker.get_kv_count() == num_threads * keys_per_thread
|
|
|
|
def test_concurrent_admit_evict_kv(self):
|
|
"""Test concurrent admit and evict on the same worker."""
|
|
worker = WorkerNode(
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
num_threads = 5
|
|
operations_per_thread = 100
|
|
errors = []
|
|
|
|
def mixed_operations(thread_id):
|
|
for i in range(operations_per_thread):
|
|
key = i # Same key range for all threads
|
|
try:
|
|
if i % 2 == 0:
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=worker.worker_id,
|
|
location=location,
|
|
operations=[
|
|
KVOpEvent(
|
|
op_type=OpType.ADMIT,
|
|
key=key,
|
|
seq_num=thread_id * operations_per_thread + i,
|
|
)
|
|
],
|
|
)
|
|
else:
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=worker.worker_id,
|
|
location=location,
|
|
operations=[
|
|
KVOpEvent(
|
|
op_type=OpType.EVICT,
|
|
key=key,
|
|
seq_num=thread_id * operations_per_thread + i,
|
|
)
|
|
],
|
|
)
|
|
worker.handle_batched_kv_operations(msg)
|
|
except Exception as e:
|
|
errors.append("Thread %d error: %s" % (thread_id, e))
|
|
|
|
threads = [
|
|
threading.Thread(target=mixed_operations, args=(i,))
|
|
for i in range(num_threads)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
|
|
def test_concurrent_seq_num_update(self):
|
|
"""Test concurrent sequence number updates."""
|
|
worker = WorkerNode(
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
num_threads = 10
|
|
updates_per_thread = 100
|
|
|
|
def update_seq(thread_id):
|
|
for i in range(updates_per_thread):
|
|
worker.update_seq_num(location, thread_id * updates_per_thread + i)
|
|
|
|
threads = [
|
|
threading.Thread(target=update_seq, args=(i,)) for i in range(num_threads)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# Final seq_num should be set (exact value depends on thread order)
|
|
final_seq = worker.get_seq_num(location)
|
|
assert final_seq is not None
|
|
|
|
|
|
class TestInstanceNodeLocking:
|
|
"""Test InstanceNode's internal locking for workers dict."""
|
|
|
|
def test_concurrent_add_workers(self):
|
|
"""Test concurrent worker additions to same instance."""
|
|
instance = InstanceNode(instance_id="test_instance")
|
|
num_threads = 10
|
|
|
|
def add_worker(worker_id):
|
|
worker = WorkerNode(
|
|
worker_id=worker_id,
|
|
ip="127.0.0.1",
|
|
port=8000 + worker_id,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
instance.add_worker(worker)
|
|
|
|
threads = [
|
|
threading.Thread(target=add_worker, args=(i,)) for i in range(num_threads)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# All workers should be added
|
|
assert len(instance.get_worker_ids()) == num_threads
|
|
|
|
def test_concurrent_add_remove_workers(self):
|
|
"""Test concurrent worker add/remove on same instance."""
|
|
instance = InstanceNode(instance_id="test_instance")
|
|
errors = []
|
|
|
|
# Pre-add some workers
|
|
for i in range(50):
|
|
worker = WorkerNode(
|
|
worker_id=i,
|
|
ip="127.0.0.1",
|
|
port=8000 + i,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
instance.add_worker(worker)
|
|
|
|
def add_workers(start_id):
|
|
for i in range(10):
|
|
try:
|
|
worker = WorkerNode(
|
|
worker_id=start_id + i,
|
|
ip="127.0.0.1",
|
|
port=9000 + start_id + i,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
instance.add_worker(worker)
|
|
except Exception as e:
|
|
errors.append("Add error: %s" % e)
|
|
|
|
def remove_workers(start_id):
|
|
for i in range(10):
|
|
try:
|
|
instance.remove_worker(start_id + i)
|
|
except Exception as e:
|
|
errors.append("Remove error: %s" % e)
|
|
|
|
threads = []
|
|
# Add workers 100-199
|
|
for i in range(10):
|
|
threads.append(threading.Thread(target=add_workers, args=(100 + i * 10,)))
|
|
# Remove workers 0-49
|
|
for i in range(5):
|
|
threads.append(threading.Thread(target=remove_workers, args=(i * 10,)))
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
|
|
|
|
class TestRegistryTreeFineGrainedLocking:
|
|
"""Test RegistryTree's fine-grained object-based locking."""
|
|
|
|
def test_concurrent_operations_different_instances(self):
|
|
"""
|
|
Test that operations on different instances don't block each other.
|
|
This is the key benefit of fine-grained locking.
|
|
"""
|
|
registry = RegistryTree()
|
|
num_instances = 5
|
|
workers_per_instance = 20
|
|
errors = []
|
|
timing_results = []
|
|
|
|
def operate_on_instance(instance_idx):
|
|
nonlocal errors
|
|
instance_id = "instance_%d" % instance_idx
|
|
start_time = time.time()
|
|
|
|
for worker_id in range(workers_per_instance):
|
|
try:
|
|
# Register worker
|
|
registry.register_worker(
|
|
instance_id=instance_id,
|
|
worker_id=worker_id,
|
|
ip="192.168.%d.%d" % (instance_idx, worker_id),
|
|
port=8000 + worker_id,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
# KV operations
|
|
for kv_key in range(10):
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id=instance_id,
|
|
worker_id=worker_id,
|
|
location="location_%d" % worker_id,
|
|
operations=[
|
|
KVOpEvent(
|
|
op_type=OpType.ADMIT,
|
|
key=kv_key,
|
|
seq_num=kv_key,
|
|
)
|
|
],
|
|
)
|
|
registry.handle_batched_kv_operations(msg)
|
|
|
|
except Exception as e:
|
|
errors.append("Instance %d error: %s" % (instance_idx, e))
|
|
|
|
elapsed = time.time() - start_time
|
|
timing_results.append((instance_idx, elapsed))
|
|
|
|
threads = [
|
|
threading.Thread(target=operate_on_instance, args=(i,))
|
|
for i in range(num_instances)
|
|
]
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
|
|
# Verify all data
|
|
for i in range(num_instances):
|
|
instance_id = "instance_%d" % i
|
|
worker_ids = registry.get_worker_ids(instance_id)
|
|
assert len(worker_ids) == workers_per_instance, (
|
|
"Instance %d: expected %d workers, got %d"
|
|
% (
|
|
i,
|
|
workers_per_instance,
|
|
len(worker_ids),
|
|
)
|
|
)
|
|
|
|
def test_concurrent_register_deregister_same_instance(self):
|
|
"""Test concurrent register/deregister on the same instance."""
|
|
registry = RegistryTree()
|
|
instance_id = "test_instance"
|
|
num_threads = 10
|
|
operations_per_thread = 50
|
|
errors = []
|
|
|
|
def register_deregister(thread_id):
|
|
for i in range(operations_per_thread):
|
|
worker_id = thread_id * 1000 + i
|
|
try:
|
|
# Register
|
|
registry.register_worker(
|
|
instance_id=instance_id,
|
|
worker_id=worker_id,
|
|
ip="192.168.1.%d" % thread_id,
|
|
port=8000 + i,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
# Verify registration
|
|
worker = registry.get_worker(instance_id, worker_id)
|
|
if worker is None:
|
|
errors.append(
|
|
"Thread %d: Worker %d not found after register"
|
|
% (thread_id, worker_id)
|
|
)
|
|
continue
|
|
|
|
# Deregister
|
|
result = registry.deregister_worker(instance_id, worker_id)
|
|
if result is None:
|
|
errors.append(
|
|
"Thread %d: Deregister failed for worker %d"
|
|
% (thread_id, worker_id)
|
|
)
|
|
|
|
except Exception as e:
|
|
errors.append("Thread %d error: %s" % (thread_id, e))
|
|
|
|
threads = [
|
|
threading.Thread(target=register_deregister, args=(i,))
|
|
for i in range(num_threads)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
|
|
def test_concurrent_kv_operations_same_worker(self):
|
|
"""Test concurrent KV operations on the same worker."""
|
|
registry = RegistryTree()
|
|
instance_id = "test_instance"
|
|
worker_id = 0
|
|
|
|
registry.register_worker(
|
|
instance_id=instance_id,
|
|
worker_id=worker_id,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
num_threads = 10
|
|
keys_per_thread = 100
|
|
errors = []
|
|
|
|
def kv_operations(thread_id):
|
|
nonlocal errors
|
|
for i in range(keys_per_thread):
|
|
key = thread_id * keys_per_thread + i
|
|
try:
|
|
# Admit using batched operation message
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id=instance_id,
|
|
worker_id=worker_id,
|
|
location=location,
|
|
operations=[
|
|
KVOpEvent(
|
|
op_type=OpType.ADMIT,
|
|
key=key,
|
|
seq_num=thread_id * keys_per_thread + i,
|
|
)
|
|
],
|
|
)
|
|
result = registry.handle_batched_kv_operations(msg)
|
|
if not result:
|
|
errors.append(
|
|
"Thread %d: handle_batched_kv_operations failed for key %d"
|
|
% (thread_id, key)
|
|
)
|
|
except Exception as e:
|
|
errors.append("Thread %d error: %s" % (thread_id, e))
|
|
|
|
threads = [
|
|
threading.Thread(target=kv_operations, args=(i,))
|
|
for i in range(num_threads)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
|
|
# Verify total KV count
|
|
total = registry.get_total_kv_count()
|
|
assert total == num_threads * keys_per_thread
|
|
|
|
def test_find_kv_concurrent_with_admit(self):
|
|
"""Test find_kv while concurrent admit operations are happening."""
|
|
registry = RegistryTree()
|
|
num_instances = 3
|
|
workers_per_instance = 3
|
|
|
|
# Setup: register workers
|
|
for i in range(num_instances):
|
|
for w in range(workers_per_instance):
|
|
registry.register_worker(
|
|
instance_id="instance_%d" % i,
|
|
worker_id=w,
|
|
ip="192.168.%d.%d" % (i, w),
|
|
port=8000 + w,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
errors = []
|
|
found_keys = []
|
|
|
|
def admit_keys(instance_idx, worker_id):
|
|
nonlocal errors
|
|
for key in range(100):
|
|
try:
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="instance_%d" % instance_idx,
|
|
worker_id=worker_id,
|
|
location="location_%d" % worker_id,
|
|
operations=[
|
|
KVOpEvent(
|
|
op_type=OpType.ADMIT,
|
|
key=key + instance_idx * 1000,
|
|
seq_num=key,
|
|
)
|
|
],
|
|
)
|
|
registry.handle_batched_kv_operations(msg)
|
|
except Exception as e:
|
|
errors.append("Admit error: %s" % e)
|
|
|
|
def find_keys():
|
|
for _ in range(50):
|
|
for key in range(100):
|
|
try:
|
|
result = registry.find_kv(key)
|
|
if result:
|
|
found_keys.append(key)
|
|
except Exception as e:
|
|
errors.append("Find error: %s" % e)
|
|
time.sleep(0.001)
|
|
|
|
threads = []
|
|
# Admit threads
|
|
for i in range(num_instances):
|
|
for w in range(workers_per_instance):
|
|
threads.append(threading.Thread(target=admit_keys, args=(i, w)))
|
|
# Find threads
|
|
threads.append(threading.Thread(target=find_keys))
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
|
|
def test_heartbeat_update_concurrent(self):
|
|
"""Test concurrent heartbeat updates."""
|
|
registry = RegistryTree()
|
|
instance_id = "test_instance"
|
|
num_workers = 10
|
|
|
|
# Register workers
|
|
for w in range(num_workers):
|
|
registry.register_worker(
|
|
instance_id=instance_id,
|
|
worker_id=w,
|
|
ip="127.0.0.1",
|
|
port=8000 + w,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
errors = []
|
|
num_threads = 20
|
|
updates_per_thread = 100
|
|
|
|
def update_heartbeats():
|
|
for _ in range(updates_per_thread):
|
|
for w in range(num_workers):
|
|
try:
|
|
result = registry.update_heartbeat(instance_id, w, time.time())
|
|
if not result:
|
|
errors.append("Heartbeat update failed for worker %d" % w)
|
|
except Exception as e:
|
|
errors.append("Heartbeat error: %s" % e)
|
|
|
|
threads = [
|
|
threading.Thread(target=update_heartbeats) for _ in range(num_threads)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
|
|
def test_get_all_worker_infos_concurrent(self):
|
|
"""Test get_all_worker_infos during concurrent modifications."""
|
|
registry = RegistryTree()
|
|
errors = []
|
|
|
|
def register_workers():
|
|
for i in range(100):
|
|
try:
|
|
registry.register_worker(
|
|
instance_id="instance_%d" % (i % 5),
|
|
worker_id=i,
|
|
ip="192.168.1.%d" % i,
|
|
port=8000 + i,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
except Exception as e:
|
|
errors.append("Register error: %s" % e)
|
|
|
|
def get_infos():
|
|
for _ in range(50):
|
|
try:
|
|
infos = registry.get_all_worker_infos_cached()
|
|
# Just verify it returns a list without error
|
|
assert isinstance(infos, list)
|
|
except Exception as e:
|
|
errors.append("Get infos error: %s" % e)
|
|
time.sleep(0.001)
|
|
|
|
threads = [
|
|
threading.Thread(target=register_workers),
|
|
threading.Thread(target=get_infos),
|
|
threading.Thread(target=get_infos),
|
|
]
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
|
|
def test_instance_cleanup_on_last_worker_deregister(self):
|
|
"""Test that empty instances are cleaned up correctly."""
|
|
registry = RegistryTree()
|
|
instance_id = "test_instance"
|
|
num_workers = 10
|
|
|
|
# Register workers
|
|
for w in range(num_workers):
|
|
registry.register_worker(
|
|
instance_id=instance_id,
|
|
worker_id=w,
|
|
ip="127.0.0.1",
|
|
port=8000 + w,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
# Deregister all workers concurrently
|
|
errors = []
|
|
|
|
def deregister(worker_id):
|
|
try:
|
|
registry.deregister_worker(instance_id, worker_id)
|
|
except Exception as e:
|
|
errors.append("Deregister error: %s" % e)
|
|
|
|
threads = [
|
|
threading.Thread(target=deregister, args=(w,)) for w in range(num_workers)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
|
|
def test_high_contention_stress(self):
|
|
"""Stress test with high contention on a single instance."""
|
|
registry = RegistryTree()
|
|
instance_id = "stress_instance"
|
|
errors = []
|
|
num_threads = 30
|
|
operations_per_thread = 100
|
|
|
|
def stress_operations(thread_id):
|
|
nonlocal errors
|
|
for i in range(operations_per_thread):
|
|
worker_id = thread_id * operations_per_thread + i
|
|
try:
|
|
# Register
|
|
registry.register_worker(
|
|
instance_id=instance_id,
|
|
worker_id=worker_id,
|
|
ip="10.0.0.%d" % (thread_id % 256),
|
|
port=8000 + (i % 1000),
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
# KV operations
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id=instance_id,
|
|
worker_id=worker_id,
|
|
location="loc1",
|
|
operations=[
|
|
KVOpEvent(
|
|
op_type=OpType.ADMIT,
|
|
key=i,
|
|
seq_num=i,
|
|
)
|
|
],
|
|
)
|
|
registry.handle_batched_kv_operations(msg)
|
|
|
|
# Read operations
|
|
registry.get_worker(instance_id, worker_id)
|
|
|
|
# Deregister
|
|
registry.deregister_worker(instance_id, worker_id)
|
|
|
|
except Exception as e:
|
|
errors.append(
|
|
"Thread %d iteration %d error: %s" % (thread_id, i, e)
|
|
)
|
|
|
|
with ThreadPoolExecutor(max_workers=num_threads) as executor:
|
|
futures = [
|
|
executor.submit(stress_operations, i) for i in range(num_threads)
|
|
]
|
|
for f in futures:
|
|
f.result()
|
|
|
|
assert not errors, "Errors occurred (first 10): %s" % errors[:10]
|
|
|
|
def test_data_consistency_after_concurrent_ops(self):
|
|
"""Verify data consistency after heavy concurrent operations."""
|
|
registry = RegistryTree()
|
|
num_instances = 5
|
|
workers_per_instance = 10
|
|
kv_keys = 50
|
|
|
|
# Register all workers
|
|
for inst in range(num_instances):
|
|
for w in range(workers_per_instance):
|
|
registry.register_worker(
|
|
instance_id="inst_%d" % inst,
|
|
worker_id=w,
|
|
ip="10.%d.%d.1" % (inst, w),
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
errors = []
|
|
|
|
def concurrent_operations(inst_idx):
|
|
nonlocal errors
|
|
instance_id = "inst_%d" % inst_idx
|
|
for w in range(workers_per_instance):
|
|
for key in range(kv_keys):
|
|
try:
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id=instance_id,
|
|
worker_id=w,
|
|
location="loc",
|
|
operations=[
|
|
KVOpEvent(
|
|
op_type=OpType.ADMIT,
|
|
key=key,
|
|
seq_num=key,
|
|
)
|
|
],
|
|
)
|
|
registry.handle_batched_kv_operations(msg)
|
|
except Exception as e:
|
|
errors.append("Admit error: %s" % e)
|
|
|
|
threads = [
|
|
threading.Thread(target=concurrent_operations, args=(i,))
|
|
for i in range(num_instances)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors during concurrent ops: %s" % errors
|
|
|
|
# Verify final state
|
|
total_kv = registry.get_total_kv_count()
|
|
expected_kv = num_instances * workers_per_instance * kv_keys
|
|
assert total_kv == expected_kv, "Expected %d KV entries, got %d" % (
|
|
expected_kv,
|
|
total_kv,
|
|
)
|
|
|
|
for inst in range(num_instances):
|
|
worker_ids = registry.get_worker_ids("inst_%d" % inst)
|
|
assert len(worker_ids) == workers_per_instance, (
|
|
"Instance %d: expected %d workers, got %d"
|
|
% (
|
|
inst,
|
|
workers_per_instance,
|
|
len(worker_ids),
|
|
)
|
|
)
|
|
|
|
|
|
class TestBatchOperations:
|
|
"""Test batch KV operations for performance optimization."""
|
|
|
|
def test_batch_admit_kv_basic(self):
|
|
"""Test basic batch admit functionality using handle_batched_kv_operations."""
|
|
worker = WorkerNode(
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
keys = list(range(100))
|
|
|
|
# Create batch operation message
|
|
operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=key) for key in keys
|
|
]
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=worker.worker_id,
|
|
location=location,
|
|
operations=operations,
|
|
)
|
|
worker.handle_batched_kv_operations(msg)
|
|
|
|
assert worker.get_kv_count() == 100
|
|
for key in keys:
|
|
assert worker.has_kv(location, key)
|
|
|
|
def test_batch_evict_kv_basic(self):
|
|
"""Test basic batch evict functionality using handle_batched_kv_operations."""
|
|
worker = WorkerNode(
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
keys = list(range(100))
|
|
|
|
# First admit all keys using batch operation - start with seq_num=0
|
|
admit_operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=i)
|
|
for i, key in enumerate(keys)
|
|
]
|
|
admit_msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=worker.worker_id,
|
|
location=location,
|
|
operations=admit_operations,
|
|
)
|
|
worker.handle_batched_kv_operations(admit_msg)
|
|
assert worker.get_kv_count() == 100
|
|
|
|
# Evict half of them using batch operation - continue sequence
|
|
evict_keys = list(range(50))
|
|
evict_operations = [
|
|
KVOpEvent(op_type=OpType.EVICT, key=key, seq_num=100 + i)
|
|
for i, key in enumerate(evict_keys)
|
|
]
|
|
evict_msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=worker.worker_id,
|
|
location=location,
|
|
operations=evict_operations,
|
|
)
|
|
worker.handle_batched_kv_operations(evict_msg)
|
|
|
|
assert worker.get_kv_count() == 50
|
|
|
|
# Check remaining keys
|
|
for key in range(50, 100):
|
|
assert worker.has_kv(location, key)
|
|
|
|
def test_batch_evict_nonexistent_keys(self):
|
|
"""Test batch evict with some non-existent keys
|
|
using handle_batched_kv_operations."""
|
|
worker = WorkerNode(
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
|
|
# Admit only keys 0-49 using batch operation - start with seq_num=0
|
|
admit_keys = list(range(50))
|
|
admit_operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=i)
|
|
for i, key in enumerate(admit_keys)
|
|
]
|
|
admit_msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=worker.worker_id,
|
|
location=location,
|
|
operations=admit_operations,
|
|
)
|
|
worker.handle_batched_kv_operations(admit_msg)
|
|
assert worker.get_kv_count() == 50
|
|
|
|
# Try to evict keys 0-99 (half don't exist)
|
|
# using batch operation - continue sequence
|
|
evict_keys = list(range(100))
|
|
evict_operations = [
|
|
KVOpEvent(op_type=OpType.EVICT, key=key, seq_num=50 + i)
|
|
for i, key in enumerate(evict_keys)
|
|
]
|
|
evict_msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=worker.worker_id,
|
|
location=location,
|
|
operations=evict_operations,
|
|
)
|
|
worker.handle_batched_kv_operations(evict_msg)
|
|
|
|
# All admitted keys should be evicted
|
|
assert worker.get_kv_count() == 0
|
|
|
|
def test_registry_batch_operations(self):
|
|
"""Test batch operations through RegistryTree
|
|
using handle_batched_kv_operations."""
|
|
registry = RegistryTree()
|
|
|
|
registry.register_worker(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
keys = list(range(1000))
|
|
|
|
# Batch admit using handle_batched_kv_operations - start with seq_num=0
|
|
admit_operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=i)
|
|
for i, key in enumerate(keys)
|
|
]
|
|
admit_msg = BatchedKVOperationMsg(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
location=location,
|
|
operations=admit_operations,
|
|
)
|
|
result = registry.handle_batched_kv_operations(admit_msg)
|
|
assert result is not False # Should succeed
|
|
assert registry.get_total_kv_count() == 1000
|
|
|
|
# Batch evict using handle_batched_kv_operations - continue sequence
|
|
evict_keys = keys[:500]
|
|
evict_operations = [
|
|
KVOpEvent(op_type=OpType.EVICT, key=key, seq_num=1000 + i)
|
|
for i, key in enumerate(evict_keys)
|
|
]
|
|
evict_msg = BatchedKVOperationMsg(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
location=location,
|
|
operations=evict_operations,
|
|
)
|
|
registry.handle_batched_kv_operations(evict_msg)
|
|
assert registry.get_total_kv_count() == 500
|
|
|
|
def test_batch_operations_nonexistent_worker(self):
|
|
"""Test batch operations on non-existent worker
|
|
using handle_batched_kv_operations."""
|
|
registry = RegistryTree()
|
|
|
|
# Batch admit to non-existent worker using handle_batched_kv_operations
|
|
admit_operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=key) for key in [1, 2, 3]
|
|
]
|
|
admit_msg = BatchedKVOperationMsg(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
location="loc",
|
|
operations=admit_operations,
|
|
)
|
|
result = registry.handle_batched_kv_operations(admit_msg)
|
|
# Should return None or indicate failure when worker doesn't exist
|
|
assert result is None or result is False
|
|
|
|
# Batch evict from non-existent worker using handle_batched_kv_operations
|
|
evict_operations = [
|
|
KVOpEvent(op_type=OpType.EVICT, key=key, seq_num=key) for key in [1, 2, 3]
|
|
]
|
|
evict_msg = BatchedKVOperationMsg(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
location="loc",
|
|
operations=evict_operations,
|
|
)
|
|
result = registry.handle_batched_kv_operations(evict_msg)
|
|
# Should return None or indicate failure when worker doesn't exist
|
|
assert result is None or result is False
|
|
|
|
def test_concurrent_batch_operations(self):
|
|
"""Test concurrent batch operations on same worker
|
|
using handle_batched_kv_operations."""
|
|
worker = WorkerNode(
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
last_heartbeat_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
num_threads = 10
|
|
keys_per_batch = 100
|
|
errors = []
|
|
|
|
def batch_admit(thread_id):
|
|
keys = [thread_id * keys_per_batch + i for i in range(keys_per_batch)]
|
|
try:
|
|
operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=key)
|
|
for key in keys
|
|
]
|
|
msg = BatchedKVOperationMsg(
|
|
instance_id="test_instance",
|
|
worker_id=worker.worker_id,
|
|
location=location,
|
|
operations=operations,
|
|
)
|
|
worker.handle_batched_kv_operations(msg)
|
|
except Exception as e:
|
|
errors.append("Thread %d error: %s" % (thread_id, e))
|
|
|
|
threads = [
|
|
threading.Thread(target=batch_admit, args=(i,)) for i in range(num_threads)
|
|
]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, "Errors occurred: %s" % errors
|
|
assert worker.get_kv_count() == num_threads * keys_per_batch
|
|
|
|
def test_registry_batch_with_seq_check(self):
|
|
"""Test batch operations with sequence check through RegistryTree
|
|
using handle_batched_kv_operations."""
|
|
registry = RegistryTree()
|
|
|
|
registry.register_worker(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
discontinuity_detected = False
|
|
|
|
def on_discontinuity():
|
|
nonlocal discontinuity_detected
|
|
discontinuity_detected = True
|
|
|
|
# First batch admit with seq check (continuous)
|
|
batch1_operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=seq_num)
|
|
for key, seq_num in [(1, 0), (2, 1), (3, 2)]
|
|
]
|
|
batch1_msg = BatchedKVOperationMsg(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
location=location,
|
|
operations=batch1_operations,
|
|
)
|
|
result = registry.handle_batched_kv_operations(batch1_msg)
|
|
# Check discontinuity count instead
|
|
discontinuity_count = registry.get_seq_discontinuity_count()
|
|
assert result is not False and discontinuity_count == 0, (
|
|
"First batch should succeed and be continuous"
|
|
)
|
|
assert registry.get_total_kv_count() == 3
|
|
|
|
# Second batch admit (continuous)
|
|
batch2_operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=seq_num)
|
|
for key, seq_num in [(4, 3), (5, 4), (6, 5)]
|
|
]
|
|
batch2_msg = BatchedKVOperationMsg(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
location=location,
|
|
operations=batch2_operations,
|
|
)
|
|
result = registry.handle_batched_kv_operations(batch2_msg)
|
|
discontinuity_count = registry.get_seq_discontinuity_count()
|
|
assert result is not False and discontinuity_count == 0, (
|
|
"Second batch should succeed and be continuous"
|
|
)
|
|
assert registry.get_total_kv_count() == 6
|
|
|
|
# Third batch with gap
|
|
batch3_operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=seq_num)
|
|
for key, seq_num in [(10, 10), (11, 11)]
|
|
]
|
|
batch3_msg = BatchedKVOperationMsg(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
location=location,
|
|
operations=batch3_operations,
|
|
)
|
|
result = registry.handle_batched_kv_operations(batch3_msg)
|
|
discontinuity_count = registry.get_seq_discontinuity_count()
|
|
assert result is not False, "Batch with gap should still succeed"
|
|
assert discontinuity_count > 0, "Sequence discontinuity should be detected"
|
|
assert registry.get_total_kv_count() == 8
|
|
|
|
def test_registry_batch_evict_with_seq_check(self):
|
|
"""Test batch evict with sequence check through RegistryTree
|
|
using handle_batched_kv_operations."""
|
|
registry = RegistryTree()
|
|
|
|
registry.register_worker(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
ip="127.0.0.1",
|
|
port=8000,
|
|
peer_init_url=None,
|
|
socket=None,
|
|
registration_time=time.time(),
|
|
)
|
|
|
|
location = "test_location"
|
|
|
|
# First admit some keys (continuous sequence starting from 0)
|
|
admit_operations = [
|
|
KVOpEvent(op_type=OpType.ADMIT, key=key, seq_num=i)
|
|
for i, key in enumerate(range(100))
|
|
]
|
|
admit_msg = BatchedKVOperationMsg(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
location=location,
|
|
operations=admit_operations,
|
|
)
|
|
result = registry.handle_batched_kv_operations(admit_msg)
|
|
discontinuity_count = registry.get_seq_discontinuity_count()
|
|
assert result is not False and discontinuity_count == 0, (
|
|
"Admit should succeed and be continuous"
|
|
)
|
|
assert registry.get_total_kv_count() == 100
|
|
|
|
# Batch evict with seq check (continuous) - continue sequence from 100
|
|
evict_operations = [
|
|
KVOpEvent(op_type=OpType.EVICT, key=key, seq_num=100 + i)
|
|
for i, key in enumerate(range(50))
|
|
]
|
|
evict_msg = BatchedKVOperationMsg(
|
|
instance_id="inst_0",
|
|
worker_id=0,
|
|
location=location,
|
|
operations=evict_operations,
|
|
)
|
|
result = registry.handle_batched_kv_operations(evict_msg)
|
|
discontinuity_count_after_evict = registry.get_seq_discontinuity_count()
|
|
assert discontinuity_count_after_evict == 0, (
|
|
"Evict should succeed and be continuous"
|
|
)
|
|
assert result is not False, "Evict should succeed"
|
|
# Note: Evict operations don't check sequence continuity,
|
|
# so count shouldn't increase
|
|
assert registry.get_total_kv_count() == 50
|