Asyncio Module Compact Serialization [6/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

@@ -5,13 +5,14 @@
 
 from hazelcast.internal.asyncio_cluster import ClusterService, _InternalClusterService
 from hazelcast.internal.asyncio_compact import CompactSchemaService
-from hazelcast.config import Config
+from hazelcast.config import Config, IndexConfig
 from hazelcast.internal.asyncio_connection import ConnectionManager, DefaultAddressProvider
 from hazelcast.core import DistributedObjectEvent, DistributedObjectInfo
 from hazelcast.cp import CPSubsystem, ProxySessionManager
 from hazelcast.discovery import HazelcastCloudAddressProvider
 from hazelcast.errors import IllegalStateError, InvalidConfigurationError
 from hazelcast.internal.asyncio_invocation import InvocationService, Invocation
+from hazelcast.internal.asyncio_proxy.vector_collection import VectorCollection
 from hazelcast.lifecycle import LifecycleService, LifecycleState, _InternalLifecycleService
 from hazelcast.internal.asyncio_listener import ClusterViewListenerService, ListenerService
 from hazelcast.near_cache import NearCacheManager
@@ -20,17 +21,19 @@
     client_add_distributed_object_listener_codec,
     client_get_distributed_objects_codec,
     client_remove_distributed_object_listener_codec,
+    dynamic_config_add_vector_collection_config_codec,
 )
 from hazelcast.internal.asyncio_proxy.manager import (
     MAP_SERVICE,
     ProxyManager,
+    VECTOR_SERVICE,
 )
 from hazelcast.internal.asyncio_proxy.base import Proxy
 from hazelcast.internal.asyncio_proxy.map import Map
 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,10 +165,9 @@ def _init_context(self):
         )
 
     async def _start(self):
-        self._reactor.start()
         try:
             self._internal_lifecycle_service.start()
-            self._invocation_service.start()
+            await self._invocation_service.start()
             membership_listeners = self._config.membership_listeners
             self._internal_cluster_service.start(self._connection_manager, membership_listeners)
             self._cluster_view_listener.start()
@@ -177,7 +179,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
@@ -186,6 +188,37 @@ async def _start(self):
     async def get_map(self, name: str) -> Map[KeyType, ValueType]:
         return await self._proxy_manager.get_or_create(MAP_SERVICE, name)
 
+    async def create_vector_collection_config(
+        self,
+        name: str,
+        indexes: typing.List[IndexConfig],
+        backup_count: int = 1,
+        async_backup_count: int = 0,
+        split_brain_protection_name: typing.Optional[str] = None,
+        merge_policy: str = "PutIfAbsentMergePolicy",
+        merge_batch_size: int = 100,
+    ) -> None:
+        # check that indexes have different names
+        if indexes:
+            index_names = set(index.name for index in indexes)
+            if len(index_names) != len(indexes):
+                raise AssertionError("index names must be unique")
+
+        request = dynamic_config_add_vector_collection_config_codec.encode_request(
+            name,
+            indexes,
+            backup_count,
+            async_backup_count,
+            split_brain_protection_name,
+            merge_policy,
+            merge_batch_size,
+        )
+        invocation = Invocation(request, response_handler=lambda m: m)
+        await self._invocation_service.ainvoke(invocation)
+
+    async def get_vector_collection(self, name: str) -> VectorCollection:
+        return await self._proxy_manager.get_or_create(VECTOR_SERVICE, name)
+
     async def add_distributed_object_listener(
         self, listener_func: typing.Callable[[DistributedObjectEvent], None]
     ) -> str:
@@ -250,7 +283,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_compact.py

@@ -57,77 +57,61 @@ def fetch_schema(self, schema_id: int) -> asyncio.Future:
         self._invocation_service.invoke(fetch_schema_invocation)
         return fetch_schema_invocation.future
 
-    def send_schema_and_retry(
+    async def send_schema_and_retry(
         self,
         error: "SchemaNotReplicatedError",
         func: typing.Callable[..., asyncio.Future],
         *args: typing.Any,
         **kwargs: typing.Any,
-    ) -> asyncio.Future:
+    ) -> None:
         schema = error.schema
         clazz = error.clazz
         request = client_send_schema_codec.encode_request(schema)
 
-        def callback():
+        async def callback():
             self._has_replicated_schemas = True
             self._compact_serializer.register_schema_to_type(schema, clazz)
-            return func(*args, **kwargs)
+            return await func(*args, **kwargs)
 
-        return self._replicate_schema(
-            schema, request, CompactSchemaService._SEND_SCHEMA_RETRY_COUNT, callback
+        return await self._replicate_schema(
+            schema, request, CompactSchemaService._SEND_SCHEMA_RETRY_COUNT, callback()
         )
 
-    def _replicate_schema(
+    async def _replicate_schema(
         self,
         schema: "Schema",
         request: "OutboundMessage",
         remaining_retries: int,
-        callback: typing.Callable[..., asyncio.Future],
-    ) -> asyncio.Future:
-        def continuation(future: asyncio.Future):
-            replicated_members = future.result()
+        callback: typing.Coroutine[typing.Any, typing.Any, typing.Any],
+    ) -> None:
+        while remaining_retries >= 2:
+            replicated_members = await self._send_schema_replication_request(request)
             members = self._cluster_service.get_members()
             for member in members:
                 if member.uuid not in replicated_members:
                     break
             else:
                 # Loop completed normally.
                 # All members in our member list all known to have the schema
-                return callback()
+                return await callback
 
             # There is a member in our member list that the schema
             # is not known to be replicated yet. We should retry
             # sending it in a random member.
-            if remaining_retries <= 1:
-                # We tried to send it a couple of times, but the member list
-                # in our local and the member list returned by the initiator
-                # nodes did not match.
-                raise IllegalStateError(
-                    f"The schema {schema} cannot be replicated in the cluster, "
-                    f"after {CompactSchemaService._SEND_SCHEMA_RETRY_COUNT} retries. "
-                    f"It might be the case that the client is connected to the two "
-                    f"halves of the cluster that is experiencing a split-brain, "
-                    f"and continue putting the data associated with that schema "
-                    f"might result in data loss. It might be possible to replicate "
-                    f"the schema after some time, when the cluster is healed."
-                )
-
-            delayed_future: asyncio.Future = asyncio.get_running_loop().create_future()
-            self._reactor.add_timer(
-                self._invocation_retry_pause,
-                lambda: delayed_future.set_result(None),
-            )
-
-            def retry(_):
-                return self._replicate_schema(
-                    schema, request.copy(), remaining_retries - 1, callback
-                )
-
-            return delayed_future.add_done_callback(retry)
-
-        fut = self._send_schema_replication_request(request)
-        fut.add_done_callback(continuation)
-        return fut
+            await asyncio.sleep(self._invocation_retry_pause)
+
+        # We tried to send it a couple of times, but the member list
+        # in our local and the member list returned by the initiator
+        # nodes did not match.
+        raise IllegalStateError(
+            f"The schema {schema} cannot be replicated in the cluster, "
+            f"after {CompactSchemaService._SEND_SCHEMA_RETRY_COUNT} retries. "
+            f"It might be the case that the client is connected to the two "
+            f"halves of the cluster that is experiencing a split-brain, "
+            f"and continue putting the data associated with that schema "
+            f"might result in data loss. It might be possible to replicate "
+            f"the schema after some time, when the cluster is healed."
+        )
 
     def _send_schema_replication_request(self, request: "OutboundMessage") -> asyncio.Future:
         invocation = Invocation(request, response_handler=client_send_schema_codec.decode_response)

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_invocation.py

@@ -96,7 +96,7 @@ def __init__(self, client, config, reactor):
         self._backup_ack_to_client_enabled = smart_routing and config.backup_ack_to_client_enabled
         self._fail_on_indeterminate_state = config.fail_on_indeterminate_operation_state
         self._backup_timeout = config.operation_backup_timeout
-        self._clean_resources_timer = None
+        self._clean_resources_task = None
         self._shutdown = False
         self._compact_schema_service = None
 
@@ -107,8 +107,8 @@ def init(self, partition_service, connection_manager, listener_service, compact_
         self._check_invocation_allowed_fn = connection_manager.check_invocation_allowed
         self._compact_schema_service = compact_schema_service
 
-    def start(self):
-        self._start_clean_resources_timer()
+    async def start(self):
+        await self._start_clean_resources_timer()
 
     async def add_backup_listener(self):
         if self._backup_ack_to_client_enabled:
@@ -152,8 +152,8 @@ def shutdown(self):
             return
 
         self._shutdown = True
-        if self._clean_resources_timer:
-            self._clean_resources_timer.cancel()
+        if self._clean_resources_task:
+            self._clean_resources_task.cancel()
         for invocation in list(self._pending.values()):
             self._notify_error(invocation, HazelcastClientNotActiveError())
 
@@ -400,8 +400,9 @@ def _notify_backup_complete(self, invocation):
 
         self._complete(invocation, invocation.pending_response)
 
-    def _start_clean_resources_timer(self):
-        def run():
+    async def _start_clean_resources_timer(self):
+        async def run():
+            await asyncio.sleep(self._CLEAN_RESOURCES_PERIOD)
             if self._shutdown:
                 return
 
@@ -419,9 +420,9 @@ def run():
                 if self._backup_ack_to_client_enabled:
                     self._detect_and_handle_backup_timeout(invocation, now)
 
-            self._clean_resources_timer = self._reactor.add_timer(self._CLEAN_RESOURCES_PERIOD, run)
+            self._clean_resources_task = asyncio.create_task(run())
 
-        self._clean_resources_timer = self._reactor.add_timer(self._CLEAN_RESOURCES_PERIOD, run)
+        self._clean_resources_task = asyncio.create_task(run())
 
     def _detect_and_handle_backup_timeout(self, invocation, now):
         if not invocation.pending_response:

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: