Handle cancellation in close (#143)

* Handle cancellation in close

* Using a future instead of sleepAsync

tests/testwebsockets.nim

@@ -170,6 +170,37 @@ suite "Test transmission":
     check string.fromBytes(clientRes) == testString
     await waitForClose(session)
 
+  asyncTest "Close handle cancellation":
+    let testString = "Hello!"
+    let cancelSignal = newFuture[void]()
+
+    proc handle(request: HttpRequest) {.async.} =
+      check request.uri.path == WSPath
+
+      let server = WSServer.new(protos = ["proto"])
+      let ws = await server.handleRequest(request)
+      let servRes = await ws.recvMsg()
+
+      check string.fromBytes(servRes) == testString
+      await ws.waitForClose()
+
+    server = createServer(
+      address = address,
+      handler = handle,
+      flags = {ReuseAddr})
+
+    proc client() {.async, gcsafe.} =
+      let session = await connectClient()
+      await session.send(testString)
+      cancelSignal.complete()
+      expect CancelledError:
+        await session.close()
+
+    let task = client()
+
+    await cancelSignal
+    await task.cancelAndWait()
+
 suite "Test ping-pong":
   setup:
     var

websock/session.nim

@@ -518,11 +518,16 @@ proc close*(
     try:
       while ws.readyState != ReadyState.Closed:
         discard await ws.readFrame()
+    except CancelledError as exc:
+      raise exc
     except CatchableError as exc:
       discard # most likely EOF
   try:
     ws.readyState = ReadyState.Closing
     await gentleCloser(ws, prepareCloseBody(code, reason)).wait(10.seconds)
+  except CancelledError as exc:
+    trace "Cancellation when closing!", exc = exc.msg
+    raise exc
   except CatchableError as exc:
     trace "Exception closing", exc = exc.msg
   finally: