From d4d970581f2cf69c435d04b2550adb38c7fb3905 Mon Sep 17 00:00:00 2001
From: bhartnett <51288821+bhartnett@users.noreply.github.com>
Date: Fri, 4 Oct 2024 16:39:58 +0800
Subject: [PATCH] Fix socket server close.

---
 json_rpc/clients/socketclient.nim | 7 ++++++-
 tests/testserverclient.nim        | 4 ++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/json_rpc/clients/socketclient.nim b/json_rpc/clients/socketclient.nim
index fbc6340..71722d0 100644
--- a/json_rpc/clients/socketclient.nim
+++ b/json_rpc/clients/socketclient.nim
@@ -80,6 +80,7 @@ method callBatch*(client: RpcSocketClient,
 proc processData(client: RpcSocketClient) {.async: (raises: []).} =
   while true:
     var localException: ref JsonRpcError
+    var stopping = false
     while true:
       try:
         var value = await client.transport.readLine(defaultMaxRequestLength)
@@ -99,6 +100,7 @@ proc processData(client: RpcSocketClient) {.async: (raises: []).} =
         break
       except CancelledError as exc:
         localException = newException(JsonRpcError, exc.msg)
+        stopping = true
         await client.transport.closeWait()
         break
 
@@ -108,6 +110,9 @@ proc processData(client: RpcSocketClient) {.async: (raises: []).} =
       if client.batchFut.isNil.not and not client.batchFut.completed():
         client.batchFut.fail(localException)
 
+    if stopping:
+      break
+
     # async loop reconnection and waiting
     try:
       info "Reconnect to server", address=`$`(client.address)
@@ -136,5 +141,5 @@ method isConnected*(client: RpcSocketClient): bool =
 method close*(client: RpcSocketClient) {.async.} =
   await client.loop.cancelAndWait()
   if not client.transport.isNil:
-    client.transport.close()
+    await client.transport.closeWait()
     client.transport = nil
diff --git a/tests/testserverclient.nim b/tests/testserverclient.nim
index 83f0e0b..a36cfab 100644
--- a/tests/testserverclient.nim
+++ b/tests/testserverclient.nim
@@ -52,8 +52,8 @@ suite "Socket Server/Client RPC":
   test "Client close and isConnected":
     check client.isConnected() == true
     # Is socket server close broken?
-    # waitFor client.close()
-    # check client.isConnected() == false
+    waitFor client.close()
+    check client.isConnected() == false
 
   srv.stop()
   waitFor srv.closeWait()