Asyncio Module Ported more tests [4/6]

Makefile

@@ -5,10 +5,10 @@ check:
 	black --check --config black.toml .
 
 test:
-	pytest -m "not enterprise"
+	pytest --verbose -m "not enterprise"
 
 test-enterprise:
-	pytest
+	pytest --verbose
 
 test-cover:
 	pytest --cov=hazelcast --cov-report=xml

hazelcast/asyncio/client.py

@@ -30,7 +30,7 @@
 from hazelcast.internal.asyncio_reactor import AsyncioReactor
 from hazelcast.serialization import SerializationServiceV1
 from hazelcast.sql import SqlService, _InternalSqlService
-from hazelcast.statistics import Statistics
+from hazelcast.internal.asyncio_statistics import Statistics
 from hazelcast.types import KeyType, ValueType, ItemType, MessageType
 from hazelcast.util import AtomicInteger, RoundRobinLB
 
@@ -162,7 +162,6 @@ def _init_context(self):
         )
 
     async def _start(self):
-        self._reactor.start()
         try:
             self._internal_lifecycle_service.start()
             self._invocation_service.start()
@@ -177,7 +176,7 @@ async def _start(self):
             self._listener_service.start()
             await self._invocation_service.add_backup_listener()
             self._load_balancer.init(self._cluster_service)
-            self._statistics.start()
+            await self._statistics.start()
         except Exception:
             await self.shutdown()
             raise
@@ -250,7 +249,6 @@ async def shutdown(self) -> None:
                 await self._connection_manager.shutdown()
                 self._invocation_service.shutdown()
                 self._statistics.shutdown()
-                self._reactor.shutdown()
                 self._internal_lifecycle_service.fire_lifecycle_event(LifecycleState.SHUTDOWN)
 
     @property

hazelcast/internal/asyncio_connection.py

@@ -185,6 +185,9 @@ def __init__(
         self._use_public_ip = (
             isinstance(address_provider, DefaultAddressProvider) and config.use_public_ip
         )
+        # asyncio tasks are weakly referenced
+        # storing tasks here in order not to lose them midway
+        self._tasks = set()
 
     def add_listener(self, on_connection_opened=None, on_connection_closed=None):
         """Registers a ConnectionListener.
@@ -315,22 +318,21 @@ async def on_connection_close(self, closed_connection):
         disconnected = False
         removed = False
         trigger_reconnection = False
-        async with self._lock:
-            connection = self.active_connections.get(remote_uuid, None)
-            if connection == closed_connection:
-                self.active_connections.pop(remote_uuid, None)
-                removed = True
-                _logger.info(
-                    "Removed connection to %s:%s, connection: %s",
-                    remote_address,
-                    remote_uuid,
-                    connection,
-                )
+        connection = self.active_connections.get(remote_uuid, None)
+        if connection == closed_connection:
+            self.active_connections.pop(remote_uuid, None)
+            removed = True
+            _logger.info(
+                "Removed connection to %s:%s, connection: %s",
+                remote_address,
+                remote_uuid,
+                connection,
+            )
 
-                if not self.active_connections:
-                    trigger_reconnection = True
-                    if self._client_state == ClientState.INITIALIZED_ON_CLUSTER:
-                        disconnected = True
+            if not self.active_connections:
+                trigger_reconnection = True
+                if self._client_state == ClientState.INITIALIZED_ON_CLUSTER:
+                    disconnected = True
 
         if disconnected:
             self._lifecycle_service.fire_lifecycle_event(LifecycleState.DISCONNECTED)
@@ -809,6 +811,9 @@ def __init__(self, connection_manager, client, config, reactor, invocation_servi
         self._heartbeat_timeout = config.heartbeat_timeout
         self._heartbeat_interval = config.heartbeat_interval
         self._heartbeat_task: asyncio.Task | None = None
+        # asyncio tasks are weakly referenced
+        # storing tasks here in order not to lose them midway
+        self._tasks = set()
 
     def start(self):
         """Starts sending periodic HeartBeat operations."""
@@ -848,7 +853,9 @@ async def _check_connection(self, now, connection):
         if (now - connection.last_write_time) > self._heartbeat_interval:
             request = client_ping_codec.encode_request()
             invocation = Invocation(request, connection=connection, urgent=True)
-            asyncio.create_task(self._invocation_service.ainvoke(invocation))
+            task = asyncio.create_task(self._invocation_service.ainvoke(invocation))
+            self._tasks.add(task)
+            task.add_done_callback(self._tasks.discard)
 
 
 _frame_header = struct.Struct("<iH")

hazelcast/internal/asyncio_listener.py

@@ -175,10 +175,8 @@ async def _register_on_connection(
         self, user_registration_id, listener_registration, connection
     ):
         registration_map = listener_registration.connection_registrations
-
         if connection in registration_map:
             return
-
         registration_request = listener_registration.registration_request.copy()
         invocation = Invocation(
             registration_request,
@@ -187,26 +185,20 @@ async def _register_on_connection(
             response_handler=lambda m: m,
             urgent=True,
         )
-        self._invocation_service.invoke(invocation)
-
-        def callback(f):
-            try:
-                response = f.result()
-                server_registration_id = listener_registration.decode_register_response(response)
-                correlation_id = registration_request.get_correlation_id()
-                registration = _EventRegistration(server_registration_id, correlation_id)
-                registration_map[connection] = registration
-            except Exception as e:
-                if connection.live:
-                    _logger.exception(
-                        "Listener %s can not be added to a new connection: %s",
-                        user_registration_id,
-                        connection,
-                    )
-                raise e
-
-        invocation.future.add_done_callback(callback)
-        return await invocation.future
+        response = await self._invocation_service.ainvoke(invocation)
+        try:
+            server_registration_id = listener_registration.decode_register_response(response)
+            correlation_id = registration_request.get_correlation_id()
+            registration = _EventRegistration(server_registration_id, correlation_id)
+            registration_map[connection] = registration
+        except Exception as e:
+            if connection.live:
+                _logger.exception(
+                    "Listener %s can not be added to a new connection: %s",
+                    user_registration_id,
+                    connection,
+                )
+            raise e
 
     async def _connection_added(self, connection):
         async with self._registration_lock:

hazelcast/internal/asyncio_proxy/base.py

@@ -41,10 +41,13 @@ async def destroy(self) -> bool:
             ``True`` if this proxy is destroyed successfully, ``False``
             otherwise.
         """
-        self._on_destroy()
-        return await self._context.proxy_manager.destroy_proxy(self.service_name, self.name)
+        async with asyncio.TaskGroup() as tg:  # type: ignore[attr-defined]
+            tg.create_task(self._on_destroy())
+            return await tg.create_task(
+                self._context.proxy_manager.destroy_proxy(self.service_name, self.name)
+            )
 
-    def _on_destroy(self):
+    async def _on_destroy(self):
         pass
 
     def __repr__(self) -> str:

hazelcast/internal/asyncio_proxy/manager.py

@@ -8,7 +8,10 @@
 
 MAP_SERVICE = "hz:impl:mapService"
 
-_proxy_init: typing.Dict[str, typing.Callable[[str, str, typing.Any], Proxy]] = {
+_proxy_init: typing.Dict[
+    str,
+    typing.Callable[[str, str, typing.Any], typing.Coroutine[typing.Any, typing.Any, typing.Any]],
+] = {
     MAP_SERVICE: create_map_proxy,
 }
 
@@ -34,7 +37,7 @@ async def _create_proxy(self, service_name, name, create_on_remote) -> Proxy:
             invocation_service = self._context.invocation_service
             await invocation_service.ainvoke(invocation)
 
-        return _proxy_init[service_name](service_name, name, self._context)
+        return await _proxy_init[service_name](service_name, name, self._context)
 
     async def destroy_proxy(self, service_name, name, destroy_on_remote=True):
         ns = (service_name, name)

hazelcast/internal/asyncio_proxy/map.py

@@ -65,6 +65,7 @@
     map_set_with_max_idle_codec,
     map_remove_interceptor_codec,
     map_remove_all_codec,
+    map_add_near_cache_invalidation_listener_codec,
 )
 from hazelcast.internal.asyncio_proxy.base import (
     Proxy,
@@ -971,8 +972,177 @@ def handler(message):
         return self._invoke_on_key(request, key_data, handler)
 
 
-def create_map_proxy(service_name, name, context):
+class MapFeatNearCache(Map[KeyType, ValueType]):
+    """Map proxy implementation featuring Near Cache"""
+
+    def __init__(self, service_name, name, context):
+        super(MapFeatNearCache, self).__init__(service_name, name, context)
+        self._invalidation_listener_id = None
+        self._near_cache = context.near_cache_manager.get_or_create_near_cache(name)
+
+    async def clear(self):
+        self._near_cache._clear()
+        return await super(MapFeatNearCache, self).clear()
+
+    async def evict_all(self):
+        self._near_cache.clear()
+        return await super(MapFeatNearCache, self).evict_all()
+
+    async def load_all(self, keys=None, replace_existing_values=True):
+        if keys is None and replace_existing_values:
+            self._near_cache.clear()
+        return await super(MapFeatNearCache, self).load_all(keys, replace_existing_values)
+
+    async def _on_destroy(self):
+        await self._remove_near_cache_invalidation_listener()
+        self._near_cache.clear()
+        await super(MapFeatNearCache, self)._on_destroy()
+
+    async def _add_near_cache_invalidation_listener(self):
+        codec = map_add_near_cache_invalidation_listener_codec
+        request = codec.encode_request(self.name, EntryEventType.INVALIDATION, self._is_smart)
+        self._invalidation_listener_id = await self._register_listener(
+            request,
+            lambda r: codec.decode_response(r),
+            lambda reg_id: map_remove_entry_listener_codec.encode_request(self.name, reg_id),
+            lambda m: codec.handle(m, self._handle_invalidation, self._handle_batch_invalidation),
+        )
+
+    async def _remove_near_cache_invalidation_listener(self):
+        if self._invalidation_listener_id:
+            await self.remove_entry_listener(self._invalidation_listener_id)
+
+    def _handle_invalidation(self, key, source_uuid, partition_uuid, sequence):
+        # key is always ``Data``
+        # null key means near cache has to remove all entries in it.
+        # see MapAddNearCacheEntryListenerMessageTask.
+        if key is None:
+            self._near_cache._clear()
+        else:
+            self._invalidate_cache(key)
+
+    def _handle_batch_invalidation(self, keys, source_uuids, partition_uuids, sequences):
+        # key_list is always list of ``Data``
+        for key_data in keys:
+            self._invalidate_cache(key_data)
+
+    def _invalidate_cache(self, key_data):
+        self._near_cache._invalidate(key_data)
+
+    def _invalidate_cache_batch(self, key_data_list):
+        for key_data in key_data_list:
+            self._near_cache._invalidate(key_data)
+
+    # internals
+    async def _contains_key_internal(self, key_data):
+        try:
+            return self._near_cache[key_data]
+        except KeyError:
+            return await super(MapFeatNearCache, self)._contains_key_internal(key_data)
+
+    async def _get_internal(self, key_data):
+        try:
+            return self._near_cache[key_data]
+        except KeyError:
+            value = await super(MapFeatNearCache, self)._get_internal(key_data)
+            self._near_cache.__setitem__(key_data, value)
+            return value
+
+    async def _get_all_internal(self, partition_to_keys, tasks=None):
+        tasks = tasks or []
+        for key_dic in partition_to_keys.values():
+            for key in list(key_dic.keys()):
+                try:
+                    key_data = key_dic[key]
+                    value = self._near_cache[key_data]
+                    future = asyncio.Future()
+                    future.set_result((key, value))
+                    tasks.append(future)
+                    del key_dic[key]
+                except KeyError:
+                    pass
+        return await super(MapFeatNearCache, self)._get_all_internal(partition_to_keys, tasks)
+
+    def _try_remove_internal(self, key_data, timeout):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._try_remove_internal(key_data, timeout)
+
+    def _try_put_internal(self, key_data, value_data, timeout):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._try_put_internal(key_data, value_data, timeout)
+
+    def _set_internal(self, key_data, value_data, ttl, max_idle):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._set_internal(key_data, value_data, ttl, max_idle)
+
+    def _set_ttl_internal(self, key_data, ttl):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._set_ttl_internal(key_data, ttl)
+
+    def _replace_internal(self, key_data, value_data):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._replace_internal(key_data, value_data)
+
+    def _replace_if_same_internal(self, key_data, old_value_data, new_value_data):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._replace_if_same_internal(
+            key_data, old_value_data, new_value_data
+        )
+
+    def _remove_internal(self, key_data):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._remove_internal(key_data)
+
+    def _remove_all_internal(self, predicate_data):
+        self._near_cache.clear()
+        return super(MapFeatNearCache, self)._remove_all_internal(predicate_data)
+
+    def _remove_if_same_internal_(self, key_data, value_data):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._remove_if_same_internal_(key_data, value_data)
+
+    def _put_transient_internal(self, key_data, value_data, ttl, max_idle):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._put_transient_internal(
+            key_data, value_data, ttl, max_idle
+        )
+
+    def _put_internal(self, key_data, value_data, ttl, max_idle):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._put_internal(key_data, value_data, ttl, max_idle)
+
+    def _put_if_absent_internal(self, key_data, value_data, ttl, max_idle):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._put_if_absent_internal(
+            key_data, value_data, ttl, max_idle
+        )
+
+    def _load_all_internal(self, key_data_list, replace_existing_values):
+        self._invalidate_cache_batch(key_data_list)
+        return super(MapFeatNearCache, self)._load_all_internal(
+            key_data_list, replace_existing_values
+        )
+
+    def _execute_on_key_internal(self, key_data, entry_processor_data):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._execute_on_key_internal(
+            key_data, entry_processor_data
+        )
+
+    def _evict_internal(self, key_data):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._evict_internal(key_data)
+
+    def _delete_internal(self, key_data):
+        self._invalidate_cache(key_data)
+        return super(MapFeatNearCache, self)._delete_internal(key_data)
+
+
+async def create_map_proxy(service_name, name, context):
     near_cache_config = context.config.near_caches.get(name, None)
     if near_cache_config is None:
         return Map(service_name, name, context)
-    raise InvalidConfigurationError("near cache is not supported")
+    nc = MapFeatNearCache(service_name, name, context)
+    if nc._near_cache.invalidate_on_change:
+        await nc._add_near_cache_invalidation_listener()
+    return nc