valkey-glide 2.0.0rc3__cp310-cp310-macosx_11_0_arm64.whl → 2.2.3__cp310-cp310-macosx_11_0_arm64.whl

glide/glide_client.py

@@ -1,62 +1,121 @@
 # Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
 
-import asyncio
 import sys
 import threading
-from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
-
-from glide.async_commands.cluster_commands import ClusterCommands
-from glide.async_commands.command_args import ObjectType
-from glide.async_commands.core import CoreCommands
-from glide.async_commands.standalone_commands import StandaloneCommands
-from glide.config import BaseClientConfiguration, ServerCredentials
-from glide.constants import DEFAULT_READ_BYTES_SIZE, OK, TEncodable, TRequest, TResult
-from glide.exceptions import (
-    ClosingError,
-    ConfigurationError,
-    ConnectionError,
-    ExecAbortError,
-    RequestError,
-    TimeoutError,
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+    cast,
 )
-from glide.logger import Level as LogLevel
-from glide.logger import Logger as ClientLogger
-from glide.protobuf.command_request_pb2 import Command, CommandRequest, RequestType
-from glide.protobuf.connection_request_pb2 import ConnectionRequest
-from glide.protobuf.response_pb2 import RequestErrorType, Response
-from glide.protobuf_codec import PartialMessageException, ProtobufCodec
-from glide.routes import Route, set_protobuf_route
-
-from .glide import (
+
+import anyio
+import sniffio
+from anyio import to_thread
+from glide.glide import (
     DEFAULT_TIMEOUT_IN_MILLISECONDS,
     MAX_REQUEST_ARGS_LEN,
     ClusterScanCursor,
     create_leaked_bytes_vec,
+    create_otel_span,
+    drop_otel_span,
     get_statistics,
     start_socket_listener_external,
     value_from_pointer,
 )
+from glide_shared.commands.command_args import ObjectType
+from glide_shared.commands.core_options import PubSubMsg
+from glide_shared.config import BaseClientConfiguration, ServerCredentials
+from glide_shared.constants import (
+    DEFAULT_READ_BYTES_SIZE,
+    OK,
+    TEncodable,
+    TRequest,
+    TResult,
+)
+from glide_shared.exceptions import (
+    ClosingError,
+    ConfigurationError,
+    ConnectionError,
+    get_request_error_class,
+)
+from glide_shared.protobuf.command_request_pb2 import (
+    Command,
+    CommandRequest,
+    RefreshIamToken,
+    RequestType,
+)
+from glide_shared.protobuf.connection_request_pb2 import ConnectionRequest
+from glide_shared.protobuf.response_pb2 import Response
+from glide_shared.protobuf_codec import PartialMessageException, ProtobufCodec
+from glide_shared.routes import Route, set_protobuf_route
+
+from .async_commands.cluster_commands import ClusterCommands
+from .async_commands.core import CoreCommands
+from .async_commands.standalone_commands import StandaloneCommands
+from .logger import Level as LogLevel
+from .logger import Logger as ClientLogger
+from .opentelemetry import OpenTelemetry
 
 if sys.version_info >= (3, 11):
-    import asyncio as async_timeout
     from typing import Self
 else:
-    import async_timeout
     from typing_extensions import Self
 
+if TYPE_CHECKING:
+    import asyncio
+
+    import trio
+
+    TTask = Union[asyncio.Task[None], trio.lowlevel.Task]
+    TFuture = Union[asyncio.Future[Any], "_CompatFuture"]
+
+
+class _CompatFuture:
+    """anyio shim for asyncio.Future-like functionality"""
+
+    def __init__(self) -> None:
+        self._is_done = anyio.Event()
+        self._result: Any = None
+        self._exception: Optional[Exception] = None
+
+    def set_result(self, result: Any) -> None:
+        self._result = result
+        self._is_done.set()
+
+    def set_exception(self, exception: Exception) -> None:
+        self._exception = exception
+        self._is_done.set()
+
+    def done(self) -> bool:
+        return self._is_done.is_set()
+
+    def __await__(self):
+        return self._is_done.wait().__await__()
 
-def get_request_error_class(
-    error_type: Optional[RequestErrorType.ValueType],
-) -> Type[RequestError]:
-    if error_type == RequestErrorType.Disconnect:
-        return ConnectionError
-    if error_type == RequestErrorType.ExecAbort:
-        return ExecAbortError
-    if error_type == RequestErrorType.Timeout:
-        return TimeoutError
-    if error_type == RequestErrorType.Unspecified:
-        return RequestError
-    return RequestError
+    def result(self) -> Any:
+        if self._exception:
+            raise self._exception
+
+        return self._result
+
+
+def _get_new_future_instance() -> "TFuture":
+    if sniffio.current_async_library() == "asyncio":
+        import asyncio
+
+        return asyncio.get_running_loop().create_future()
+
+    # _CompatFuture is also compatible with asyncio, but is not as closely integrated
+    # into the asyncio event loop and thus introduces a noticeable performance
+    # degradation. so we only use it for trio
+    return _CompatFuture()
 
 
 class BaseClient(CoreCommands):
@@ -65,17 +124,47 @@ class BaseClient(CoreCommands):
         To create a new client, use the `create` classmethod
         """
         self.config: BaseClientConfiguration = config
-        self._available_futures: Dict[int, asyncio.Future] = {}
+        self._available_futures: Dict[int, "TFuture"] = {}
         self._available_callback_indexes: List[int] = list()
         self._buffered_requests: List[TRequest] = list()
         self._writer_lock = threading.Lock()
         self.socket_path: Optional[str] = None
-        self._reader_task: Optional[asyncio.Task] = None
+        self._reader_task: Optional["TTask"] = None
         self._is_closed: bool = False
-        self._pubsub_futures: List[asyncio.Future] = []
+        self._pubsub_futures: List["TFuture"] = []
         self._pubsub_lock = threading.Lock()
         self._pending_push_notifications: List[Response] = list()
 
+        self._pending_tasks: Optional[Set[Awaitable[None]]] = None
+        """asyncio-only to avoid gc on pending write tasks"""
+
+    def _create_task(self, task, *args, **kwargs):
+        """framework agnostic free-floating task shim"""
+        framework = sniffio.current_async_library()
+        if framework == "trio":
+            from functools import partial
+
+            import trio
+
+            return trio.lowlevel.spawn_system_task(partial(task, **kwargs), *args)
+        elif framework == "asyncio":
+            import asyncio
+
+            # the asyncio event loop holds weak refs to tasks, so it's recommended to
+            # hold strong refs to them during their lifetime to prevent garbage
+            # collection
+            t = asyncio.create_task(task(*args, **kwargs))
+
+            if self._pending_tasks is None:
+                self._pending_tasks = set()
+
+            self._pending_tasks.add(t)
+            t.add_done_callback(self._pending_tasks.discard)
+
+            return t
+
+        raise RuntimeError(f"Unsupported async framework {framework}")
+
     @classmethod
     async def create(cls, config: BaseClientConfiguration) -> Self:
         """Creates a Glide client.
@@ -148,8 +237,8 @@ class BaseClient(CoreCommands):
         """
         config = config
         self = cls(config)
-        init_future: asyncio.Future = asyncio.Future()
-        loop = asyncio.get_event_loop()
+
+        init_event: threading.Event = threading.Event()
 
         def init_callback(socket_path: Optional[str], err: Optional[str]):
             if err is not None:
@@ -161,7 +250,7 @@ class BaseClient(CoreCommands):
             else:
                 # Received socket path
                 self.socket_path = socket_path
-                loop.call_soon_threadsafe(init_future.set_result, True)
+                init_event.set()
 
         start_socket_listener_external(init_callback=init_callback)
 
@@ -169,36 +258,27 @@ class BaseClient(CoreCommands):
         # level or higher
         ClientLogger.log(LogLevel.INFO, "connection info", "new connection established")
         # Wait for the socket listener to complete its initialization
-        await init_future
+        await to_thread.run_sync(init_event.wait)
         # Create UDS connection
         await self._create_uds_connection()
+
         # Start the reader loop as a background task
-        self._reader_task = asyncio.create_task(self._reader_loop())
+        self._reader_task = self._create_task(self._reader_loop)
+
         # Set the client configurations
         await self._set_connection_configurations()
+
         return self
 
     async def _create_uds_connection(self) -> None:
         try:
             # Open an UDS connection
-            async with async_timeout.timeout(DEFAULT_TIMEOUT_IN_MILLISECONDS):
-                reader, writer = await asyncio.open_unix_connection(
-                    path=self.socket_path
+            with anyio.fail_after(DEFAULT_TIMEOUT_IN_MILLISECONDS):
+                self._stream = await anyio.connect_unix(
+                    path=cast(str, self.socket_path)
                 )
-            self._reader = reader
-            self._writer = writer
         except Exception as e:
-            await self.close(f"Failed to create UDS connection: {e}")
-            raise
-
-    def __del__(self) -> None:
-        try:
-            if self._reader_task:
-                self._reader_task.cancel()
-        except RuntimeError as e:
-            if "no running event loop" in str(e):
-                # event loop already closed
-                pass
+            raise ClosingError("Failed to create UDS connection") from e
 
     async def close(self, err_message: Optional[str] = None) -> None:
         """
@@ -210,25 +290,24 @@ class BaseClient(CoreCommands):
             closing all open futures.
             Defaults to None.
         """
-        self._is_closed = True
-        for response_future in self._available_futures.values():
-            if not response_future.done():
-                err_message = "" if err_message is None else err_message
-                response_future.set_exception(ClosingError(err_message))
-        try:
-            self._pubsub_lock.acquire()
-            for pubsub_future in self._pubsub_futures:
-                if not pubsub_future.done() and not pubsub_future.cancelled():
-                    pubsub_future.set_exception(ClosingError(""))
-        finally:
-            self._pubsub_lock.release()
+        if not self._is_closed:
+            self._is_closed = True
+            err_message = "" if err_message is None else err_message
+            for response_future in self._available_futures.values():
+                if not response_future.done():
+                    response_future.set_exception(ClosingError(err_message))
+            try:
+                self._pubsub_lock.acquire()
+                for pubsub_future in self._pubsub_futures:
+                    if not pubsub_future.done():
+                        pubsub_future.set_exception(ClosingError(err_message))
+            finally:
+                self._pubsub_lock.release()
 
-        self._writer.close()
-        await self._writer.wait_closed()
-        self.__del__()
+            await self._stream.aclose()
 
-    def _get_future(self, callback_idx: int) -> asyncio.Future:
-        response_future: asyncio.Future = asyncio.Future()
+    def _get_future(self, callback_idx: int) -> "TFuture":
+        response_future: "TFuture" = _get_new_future_instance()
         self._available_futures.update({callback_idx: response_future})
         return response_future
 
@@ -237,14 +316,15 @@ class BaseClient(CoreCommands):
 
     async def _set_connection_configurations(self) -> None:
         conn_request = self._get_protobuf_conn_request()
-        response_future: asyncio.Future = self._get_future(0)
-        await self._write_or_buffer_request(conn_request)
+        response_future: "TFuture" = self._get_future(0)
+        self._create_write_task(conn_request)
         await response_future
-        if response_future.result() is not OK:
-            raise ClosingError(response_future.result())
+        res = response_future.result()
+        if res is not OK:
+            raise ClosingError(res)
 
     def _create_write_task(self, request: TRequest):
-        asyncio.create_task(self._write_or_buffer_request(request))
+        self._create_task(self._write_or_buffer_request, request)
 
     async def _write_or_buffer_request(self, request: TRequest):
         self._buffered_requests.append(request)
@@ -252,7 +332,21 @@ class BaseClient(CoreCommands):
             try:
                 while len(self._buffered_requests) > 0:
                     await self._write_buffered_requests_to_socket()
-
+            except Exception as e:
+                # trio system tasks cannot raise exceptions, so gracefully propagate
+                # any error to the pending future instead
+                callback_idx = (
+                    request.callback_idx if isinstance(request, CommandRequest) else 0
+                )
+                res_future = self._available_futures.pop(callback_idx, None)
+                if res_future and not res_future.done():
+                    res_future.set_exception(e)
+                else:
+                    ClientLogger.log(
+                        LogLevel.WARN,
+                        "unhandled response error",
+                        f"Unhandled response error for unknown request: {callback_idx}",
+                    )
             finally:
                 self._writer_lock.release()
 
@@ -262,8 +356,10 @@ class BaseClient(CoreCommands):
         b_arr = bytearray()
         for request in requests:
             ProtobufCodec.encode_delimited(b_arr, request)
-        self._writer.write(b_arr)
-        await self._writer.drain()
+        try:
+            await self._stream.send(b_arr)
+        except (anyio.ClosedResourceError, anyio.EndOfStream):
+            raise ClosingError("The communication layer was unexpectedly closed.")
 
     def _encode_arg(self, arg: TEncodable) -> bytes:
         """
@@ -314,6 +410,13 @@ class BaseClient(CoreCommands):
             raise ClosingError(
                 "Unable to execute requests; the client is closed. Please create a new client."
             )
+
+        # Create span if OpenTelemetry is configured and sampling indicates we should trace
+        span = None
+        if OpenTelemetry.should_sample():
+            command_name = RequestType.Name(request_type)
+            span = create_otel_span(command_name)
+
         request = CommandRequest()
         request.callback_idx = self._get_callback_index()
         request.single_command.request_type = request_type
@@ -328,6 +431,11 @@ class BaseClient(CoreCommands):
             request.single_command.args_vec_pointer = create_leaked_bytes_vec(
                 encoded_args
             )
+
+        # Add span pointer to request if span was created
+        if span:
+            request.root_span_ptr = span
+
         set_protobuf_route(request, route)
         return await self._write_request_await_response(request)
 
@@ -345,6 +453,14 @@ class BaseClient(CoreCommands):
             raise ClosingError(
                 "Unable to execute requests; the client is closed. Please create a new client."
             )
+
+        # Create span if OpenTelemetry is configured and sampling indicates we should trace
+        span = None
+
+        if OpenTelemetry.should_sample():
+            # Use "Batch" as span name for batches
+            span = create_otel_span("Batch")
+
         request = CommandRequest()
         request.callback_idx = self._get_callback_index()
         batch_commands = []
@@ -366,6 +482,11 @@ class BaseClient(CoreCommands):
             request.batch.timeout = timeout
         request.batch.retry_server_error = retry_server_error
         request.batch.retry_connection_error = retry_connection_error
+
+        # Add span pointer to request if span was created
+        if span:
+            request.root_span_ptr = span
+
         set_protobuf_route(request, route)
         return await self._write_request_await_response(request)
 
@@ -400,7 +521,7 @@ class BaseClient(CoreCommands):
         set_protobuf_route(request, route)
         return await self._write_request_await_response(request)
 
-    async def get_pubsub_message(self) -> CoreCommands.PubSubMsg:
+    async def get_pubsub_message(self) -> PubSubMsg:
         if self._is_closed:
             raise ClosingError(
                 "Unable to execute requests; the client is closed. Please create a new client."
@@ -417,16 +538,17 @@ class BaseClient(CoreCommands):
             )
 
         # locking might not be required
-        response_future: asyncio.Future = asyncio.Future()
+        response_future: "TFuture" = _get_new_future_instance()
         try:
             self._pubsub_lock.acquire()
             self._pubsub_futures.append(response_future)
             self._complete_pubsub_futures_safe()
         finally:
             self._pubsub_lock.release()
-        return await response_future
+        await response_future
+        return response_future.result()
 
-    def try_get_pubsub_message(self) -> Optional[CoreCommands.PubSubMsg]:
+    def try_get_pubsub_message(self) -> Optional[PubSubMsg]:
         if self._is_closed:
             raise ClosingError(
                 "Unable to execute requests; the client is closed. Please create a new client."
@@ -443,7 +565,7 @@ class BaseClient(CoreCommands):
             )
 
         # locking might not be required
-        msg: Optional[CoreCommands.PubSubMsg] = None
+        msg: Optional[PubSubMsg] = None
         try:
             self._pubsub_lock.acquire()
             self._complete_pubsub_futures_safe()
@@ -457,12 +579,11 @@ class BaseClient(CoreCommands):
     def _cancel_pubsub_futures_with_exception_safe(self, exception: ConnectionError):
         while len(self._pubsub_futures):
             next_future = self._pubsub_futures.pop(0)
-            if not next_future.cancelled():
-                next_future.set_exception(exception)
+            next_future.set_exception(exception)
 
     def _notification_to_pubsub_message_safe(
         self, response: Response
-    ) -> Optional[CoreCommands.PubSubMsg]:
+    ) -> Optional[PubSubMsg]:
         pubsub_message = None
         push_notification = cast(
             Dict[str, Any], value_from_pointer(response.resp_pointer)
@@ -481,11 +602,11 @@ class BaseClient(CoreCommands):
         ):
             values: List = push_notification["values"]
             if message_kind == "PMessage":
-                pubsub_message = BaseClient.PubSubMsg(
+                pubsub_message = PubSubMsg(
                     message=values[2], channel=values[1], pattern=values[0]
                 )
             else:
-                pubsub_message = BaseClient.PubSubMsg(
+                pubsub_message = PubSubMsg(
                     message=values[1], channel=values[0], pattern=None
                 )
         elif (
@@ -538,10 +659,16 @@ class BaseClient(CoreCommands):
                 if response.HasField("closing_error")
                 else f"Client Error - closing due to unknown error. callback index:  {response.callback_idx}"
             )
+            exc = ClosingError(err_msg)
             if res_future is not None:
-                res_future.set_exception(ClosingError(err_msg))
-            await self.close(err_msg)
-            raise ClosingError(err_msg)
+                res_future.set_exception(exc)
+            else:
+                ClientLogger.log(
+                    LogLevel.WARN,
+                    "unhandled response error",
+                    f"Unhandled response error for unknown request: {response.callback_idx}",
+                )
+            raise exc
         else:
             self._available_callback_indexes.append(response.callback_idx)
             if response.HasField("request_error"):
@@ -554,6 +681,10 @@ class BaseClient(CoreCommands):
             else:
                 res_future.set_result(None)
 
+        # Clean up span if it was created
+        if response.HasField("root_span_ptr"):
+            drop_otel_span(response.root_span_ptr)
+
     async def _process_push(self, response: Response) -> None:
         if response.HasField("closing_error") or not response.HasField("resp_pointer"):
             err_msg = (
@@ -561,9 +692,7 @@ class BaseClient(CoreCommands):
                 if response.HasField("closing_error")
                 else "Client Error - push notification without resp_pointer"
             )
-            await self.close(err_msg)
             raise ClosingError(err_msg)
-
         try:
             self._pubsub_lock.acquire()
             callback, context = self.config._get_pubsub_callback_and_context()
@@ -579,30 +708,36 @@ class BaseClient(CoreCommands):
 
     async def _reader_loop(self) -> None:
         # Socket reader loop
-        remaining_read_bytes = bytearray()
-        while True:
-            read_bytes = await self._reader.read(DEFAULT_READ_BYTES_SIZE)
-            if len(read_bytes) == 0:
-                err_msg = "The communication layer was unexpectedly closed."
-                await self.close(err_msg)
-                raise ClosingError(err_msg)
-            read_bytes = remaining_read_bytes + bytearray(read_bytes)
-            read_bytes_view = memoryview(read_bytes)
-            offset = 0
-            while offset <= len(read_bytes):
+        try:
+            remaining_read_bytes = bytearray()
+            while True:
                 try:
-                    response, offset = ProtobufCodec.decode_delimited(
-                        read_bytes, read_bytes_view, offset, Response
+                    read_bytes = await self._stream.receive(DEFAULT_READ_BYTES_SIZE)
+                except (anyio.ClosedResourceError, anyio.EndOfStream):
+                    raise ClosingError(
+                        "The communication layer was unexpectedly closed."
                     )
-                except PartialMessageException:
-                    # Received only partial response, break the inner loop
-                    remaining_read_bytes = read_bytes[offset:]
-                    break
-                response = cast(Response, response)
-                if response.is_push:
-                    await self._process_push(response=response)
-                else:
-                    await self._process_response(response=response)
+                read_bytes = remaining_read_bytes + bytearray(read_bytes)
+                read_bytes_view = memoryview(read_bytes)
+                offset = 0
+                while offset <= len(read_bytes):
+                    try:
+                        response, offset = ProtobufCodec.decode_delimited(
+                            read_bytes, read_bytes_view, offset, Response
+                        )
+                    except PartialMessageException:
+                        # Received only partial response, break the inner loop
+                        remaining_read_bytes = read_bytes[offset:]
+                        break
+                    response = cast(Response, response)
+                    if response.is_push:
+                        await self._process_push(response=response)
+                    else:
+                        await self._process_response(response=response)
+        except Exception as e:
+            # close and stop reading at terminal exceptions from incoming responses or
+            # stream closures
+            await self.close(str(e))
 
     async def get_statistics(self) -> dict:
         return get_statistics()
@@ -623,6 +758,15 @@ class BaseClient(CoreCommands):
                 self.config.credentials.password = password or ""
         return response
 
+    async def _refresh_iam_token(self) -> TResult:
+        request = CommandRequest()
+        request.callback_idx = self._get_callback_index()
+        request.refresh_iam_token.CopyFrom(
+            RefreshIamToken()
+        )  # Empty message, just triggers the refresh
+        response = await self._write_request_await_response(request)
+        return response
+
 
 class GlideClusterClient(BaseClient, ClusterCommands):
     """