[DON'T MERGE] Asyncio Module VectorCollection Support [5/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,7 +165,6 @@ def _init_context(self):
         )
 
     async def _start(self):
-        self._reactor.start()
         try:
             self._internal_lifecycle_service.start()
             self._invocation_service.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_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)
@@ -340,7 +342,6 @@ async def on_connection_close(self, closed_connection):
 
         if removed:
             async with asyncio.TaskGroup() as tg:
-                # TODO: see on_connection_open
                 for _, on_connection_closed in self._connection_listeners:
                     if on_connection_closed:
                         try:
@@ -395,13 +396,12 @@ async def _get_or_connect_to_member(self, member):
 
         translated = self._translate_member_address(member)
         connection = await self._create_connection(translated)
-        response = await self._authenticate(connection)  # .continue_with(self._on_auth, connection)
+        response = await self._authenticate(connection)
         await self._on_auth(response, connection)
         return connection
 
     async def _create_connection(self, address):
-        factory = self._reactor.connection_factory
-        return await factory(
+        return await self._reactor.connection_factory(
             self,
             self._connection_id_generator.get_and_increment(),
             address,
@@ -473,7 +473,6 @@ async def run():
                     connecting_uuids.add(member_uuid)
                     if not self._lifecycle_service.running:
                         break
-                    # TODO: ERROR:asyncio:Task was destroyed but it is pending!
                     tg.create_task(self._get_or_connect_to_member(member))
                     member_uuids.append(member_uuid)
 
@@ -706,8 +705,6 @@ async def _handle_successful_auth(self, response, connection):
             for on_connection_opened, _ in self._connection_listeners:
                 if on_connection_opened:
                     try:
-                        # TODO: creating the task may not throw the exception
-                        # TODO: protect the loop against exceptions, so all handlers run
                         maybe_coro = on_connection_opened(connection)
                         if isinstance(maybe_coro, Coroutine):
                             tg.create_task(maybe_coro)
@@ -814,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."""
@@ -853,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

@@ -1,15 +1,24 @@
 import typing
 
+from hazelcast.internal.asyncio_proxy.vector_collection import (
+    VectorCollection,
+    create_vector_collection_proxy,
+)
 from hazelcast.protocol.codec import client_create_proxy_codec, client_destroy_proxy_codec
 from hazelcast.internal.asyncio_invocation import Invocation
 from hazelcast.internal.asyncio_proxy.base import Proxy
 from hazelcast.internal.asyncio_proxy.map import create_map_proxy
 from hazelcast.util import to_list
 
 MAP_SERVICE = "hz:impl:mapService"
+VECTOR_SERVICE = "hz:service:vector"
 
-_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,
+    VECTOR_SERVICE: create_vector_collection_proxy,
 }
 
 
@@ -34,7 +43,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)