Skip to content

Commit 1137c13

Browse files
liorsvexShinnRyuu
authored andcommitted
Python: improve UDS socket error handling (#4733)
fixed uds error handling Signed-off-by: Lior Sventitzky <liorsve@amazon.com>
1 parent 28ba81e commit 1137c13

File tree

3 files changed

+29
-2
lines changed

3 files changed

+29
-2
lines changed

python/glide-async/python/glide/glide_client.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -338,7 +338,7 @@ async def _write_or_buffer_request(self, request: TRequest):
338338
request.callback_idx if isinstance(request, CommandRequest) else 0
339339
)
340340
res_future = self._available_futures.pop(callback_idx, None)
341-
if res_future:
341+
if res_future and not res_future.done():
342342
res_future.set_exception(e)
343343
else:
344344
ClientLogger.log(
@@ -355,7 +355,10 @@ async def _write_buffered_requests_to_socket(self) -> None:
355355
b_arr = bytearray()
356356
for request in requests:
357357
ProtobufCodec.encode_delimited(b_arr, request)
358-
await self._stream.send(b_arr)
358+
try:
359+
await self._stream.send(b_arr)
360+
except (anyio.ClosedResourceError, anyio.EndOfStream):
361+
raise ClosingError("The communication layer was unexpectedly closed.")
359362

360363
def _encode_arg(self, arg: TEncodable) -> bytes:
361364
"""

python/tests/async_tests/test_async_client.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -412,6 +412,29 @@ async def connect_to_client():
412412
# Clean up the main client
413413
await client.close()
414414

415+
@pytest.mark.parametrize("cluster_mode", [True, False])
416+
@pytest.mark.parametrize("protocol", [ProtocolVersion.RESP2, ProtocolVersion.RESP3])
417+
async def test_UDS_socket_connection_failure(self, glide_client: TGlideClient):
418+
"""Test that the client's error handling during UDS socket connection failure"""
419+
assert await glide_client.set("test_key", "test_value") == OK
420+
assert await glide_client.get("test_key") == b"test_value"
421+
422+
# Force close the UDS connection to simulate socket failure
423+
await glide_client._stream.aclose()
424+
425+
# Verify a ClosingError is raised
426+
with pytest.raises(
427+
ClosingError, match="The communication layer was unexpectedly closed"
428+
):
429+
await glide_client.get("test_key")
430+
431+
# Verify the client is closed
432+
with pytest.raises(
433+
ClosingError,
434+
match="Unable to execute requests; the client is closed. Please create a new client.",
435+
):
436+
await glide_client.get("test_key")
437+
415438

416439
@pytest.mark.anyio
417440
class TestCommands:

python/tests/test_api_consistency.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@
6969
"test_inflight_request_limit",
7070
"test_statistics",
7171
"test_select",
72+
"test_UDS_socket_connection_failure",
7273
],
7374
"sync_only": ["test_sync_fork"],
7475
}

0 commit comments

Comments
 (0)