feat(wstransport): add support for concurrent accept

tests/helpers.nim

@@ -37,31 +37,56 @@ proc waitForClose*(ws: WSSession) {.async.} =
   except CatchableError:
     trace "Closing websocket"
 
-proc createServer*(
+proc createServerNoStart*(
   address = initTAddress("127.0.0.1:8888"),
   tlsPrivateKey = TLSPrivateKey.init(SecureKey),
   tlsCertificate = TLSCertificate.init(SecureCert),
-  handler: HttpAsyncCallback = nil,
   flags: set[ServerFlags] = {ServerFlags.TcpNoDelay, ServerFlags.ReuseAddr},
   tlsFlags: set[TLSFlags] = {},
   tlsMinVersion = TLSVersion.TLS12,
-  tlsMaxVersion = TLSVersion.TLS12): HttpServer
-  {.raises: [].} =
+  tlsMaxVersion = TLSVersion.TLS12
+): HttpServer {.raises: [].} =
   try:
-    let server = when defined secure:
+    let server = when defined(secure):
       HttpServer.create(
         address = address,
         tlsPrivateKey = tlsPrivateKey,
         tlsCertificate = tlsCertificate,
         flags = flags,
         tlsFlags = tlsFlags,
         tlsMinVersion = tlsMinVersion,
-        tlsMaxVersion = tlsMaxVersion)
+        tlsMaxVersion = tlsMaxVersion,
+      )
     else:
       HttpServer.create(
         address = address,
-        flags = flags)
+        flags = flags,
+      )
+    return server
+  except CatchableError as exc:
+    raise newException(Defect, exc.msg)
 
+proc createServer*(
+  address = initTAddress("127.0.0.1:8888"),
+  tlsPrivateKey = TLSPrivateKey.init(SecureKey),
+  tlsCertificate = TLSCertificate.init(SecureCert),
+  handler: HttpAsyncCallback = nil,
+  flags: set[ServerFlags] = {ServerFlags.TcpNoDelay, ServerFlags.ReuseAddr},
+  tlsFlags: set[TLSFlags] = {},
+  tlsMinVersion = TLSVersion.TLS12,
+  tlsMaxVersion = TLSVersion.TLS12
+): HttpServer {.raises: [].} =
+  let server = createServerNoStart(
+    address = address,
+    tlsPrivateKey = tlsPrivateKey,
+    tlsCertificate = tlsCertificate,
+    flags = flags,
+    tlsFlags = tlsFlags,
+    tlsMinVersion = tlsMinVersion,
+    tlsMaxVersion = tlsMaxVersion
+  )
+
+  try:
     when defined accepts:
       proc accepts() {.async: (raises: []).} =
         try:

tests/testwebsockets.nim

@@ -1191,3 +1191,117 @@ suite "Partial frames":
 
   asyncTest "receiver frameSize greater than sender":
     await lowLevelRecv(7, 10, 5)
+
+suite "Test accept":
+  proc runEchoClient(address: TransportAddress, id: string): Future[void] {.async.} =
+    let client = await connectClient(address).wait(2.seconds)
+    try:
+      let payload = "client-" & id
+      await client.send(payload)
+      let res = await client.recvMsg().wait(2.seconds)
+      check string.fromBytes(res) == payload
+    finally:
+      await client.close()
+
+  proc runManyClients(address: TransportAddress, count: int): Future[void] {.async.} =
+    var clientFuts: seq[Future[void]]
+    for i in 0 ..< count:
+      clientFuts.add(runEchoClient(address, $i))
+    await allFutures(clientFuts).wait(5.seconds)
+
+  template withTestServer(body: untyped) =
+    let server {.inject.} = createServerNoStart(address)
+    let wsServer {.inject.} = WSServer.new()
+    try:
+      body
+    finally:
+      server.stop()
+      await server.closeWait()
+
+  asyncTest "basic acceptStream test":
+    withTestServer:
+      proc runServer() {.async.} =
+        let clientStream = await server.acceptStream()
+        try:
+          let req = await server.processHttpRequest(clientStream)
+          check req.uri.path == WSPath
+          let wsSession = await wsServer.handleRequest(req)
+          try:
+            let msg = await wsSession.recvMsg()
+            await wsSession.send(msg, Opcode.Text)
+          finally:
+            await wsSession.close()
+        except CatchableError as exc:
+          await clientStream.closeWait()
+          raise exc
+
+      let serverTask = runServer()
+
+      try:
+        await runEchoClient(address, "single")
+      finally:
+        if not serverTask.finished:
+          await serverTask.cancelAndWait()
+        else:
+          await serverTask
+
+  asyncTest "concurrent acceptStream test":
+    withTestServer:
+      proc worker(stream: AsyncStream) {.async.} =
+        var wsSession: WSSession
+        try:
+          let req = await server.processHttpRequest(stream)
+          wsSession = await wsServer.handleRequest(req)
+          let msg = await wsSession.recvMsg()
+          await wsSession.send(msg, Opcode.Text)
+        finally:
+          if not isNil(wsSession):
+            await wsSession.close()
+          else:
+            await stream.closeWait()
+
+      proc runServer() {.async.} =
+        while true:
+          try:
+            let clientStream = await server.acceptStream()
+            asyncSpawn worker(clientStream)
+          except CancelledError:
+            break
+
+      let serverTask = runServer()
+
+      try:
+        await runManyClients(address, 100)
+      finally:
+        await serverTask.cancelAndWait()
+        if serverTask.failed:
+          raise serverTask.error
+
+  asyncTest "accept test":
+    withTestServer:
+      proc worker(req: HttpRequest) {.async.} =
+        var wsSession: WSSession
+        try:
+          wsSession = await wsServer.handleRequest(req)
+          let msg = await wsSession.recvMsg()
+          await wsSession.send(msg, Opcode.Text)
+        finally:
+          if not isNil(wsSession):
+            await wsSession.close()
+
+      proc runServer() {.async.} =
+        while true:
+          try:
+            let req = await server.accept()
+            asyncSpawn worker(req)
+          except CancelledError:
+            break
+
+      let serverTask = runServer()
+
+      try:
+        await runManyClients(address, 10)
+      finally:
+        await serverTask.cancelAndWait()
+        if serverTask.failed:
+          raise serverTask.error

websock/http/server.nim

@@ -86,6 +86,13 @@ proc readHttpRequest(
     headers: request.toHttpTable(), stream: stream, uri: request.uri().parseUri()
   )
 
+proc processHttpRequest*(
+  server: HttpServer, stream: AsyncStream
+): Future[HttpRequest] {.
+    async: (raises: [CancelledError, AsyncStreamError, HttpError])
+.} =
+  return await readHttpRequest(stream, server.headersTimeout)
+
 proc openAsyncStream(
     server: HttpServer, transp: StreamTransport
 ): Result[AsyncStream, string] =
@@ -129,9 +136,7 @@ proc handleConnCb(
   finally:
     await stream.closeWait()
 
-# TODO async raises not implemented for accept because it breaks libp2p (1.13.0
-#      at the time of writing)
-proc accept*(server: HttpServer): Future[HttpRequest] {.async.} =
+proc acceptStream*(server: HttpServer): Future[AsyncStream] {.async.} =
   if not isNil(server.handler):
     raise newException(
       HttpError, "Callback already registered - cannot mix callback and accepts styles!"
@@ -145,6 +150,13 @@ proc accept*(server: HttpServer): Future[HttpRequest] {.async.} =
       raise (ref HttpError)(msg: error)
 
   trace "Got new request", isTls = server.secure
+  return stream
+
+# TODO async raises not implemented for accept because it breaks libp2p (1.13.0
+#      at the time of writing)
+proc accept*(server: HttpServer): Future[HttpRequest] {.async.} =
+  let stream = await acceptStream(server)
+
   try:
     await stream.readHttpRequest(server.headersTimeout)
   except CancelledError as exc: