feat: perf tests and fix: pin connections and streams

Author: richard-ramos

lsquic.nimble

@@ -21,6 +21,7 @@ task format, "Format nim code using nph":
 
 task test, "Run tests":
   when defined(windows):
-    exec "nim c --mm:refc -d:nimDebugDlOpen -r --threads:on tests/test_connection.nim"
+    exec "nim c --mm:refc -d:nimDebugDlOpen --threads:on tests/test_connection.nim"
   else:
-    exec "nim c --mm:refc -r --threads:on tests/test_connection.nim"
+    exec "nim c --mm:refc --threads:on tests/test_connection.nim"
+  exec "./tests/test_connection --output-level=VERBOSE"

lsquic/connection.nim

@@ -91,6 +91,9 @@ method incomingStream*(
 method incomingStream*(
     connection: IncomingConnection
 ): Future[Stream] {.async: (raises: [CancelledError, ConnectionError]).} =
+  if connection.isClosed:
+    raise newException(ConnectionError, "connection is closed")
+
   let closedFut = connection.closed.wait()
   let incomingFut = connection.quicConn.incomingStream()
   let raceFut = await race(closedFut, incomingFut)
@@ -108,6 +111,8 @@ method openStream*(
 method openStream*(
     connection: OutgoingConnection
 ): Future[Stream] {.async: (raises: [CancelledError, ConnectionError]).} =
+  if connection.isClosed:
+    raise newException(ConnectionError, "connection is closed")
   let s = Stream.new()
   let created = connection.quicConn.addPendingStream(s)
   connection.quicContext.makeStream(connection.quicConn)

lsquic/connectionmanager.nim

@@ -46,14 +46,17 @@ proc localAddress*(
   connman.udp.localAddress()
 
 proc startSending*(connman: ConnectionManager) =
-  trace "Starting sending loop"
+  debug "Starting sending loop"
 
   proc send() {.async: (raises: [CancelledError]).} =
     try:
       let datagram = await connman.outgoing.get()
+      if datagram.len == 0:
+        # In windows, empty datagrams will make the peer stop receiving data
+        return
       await connman.udp.sendTo(datagram.taddr, datagram.data)
     except TransportError as e:
-      trace "Failed to send datagram", errorMsg = e.msg
+      debug "Failed to send datagram", errorMsg = e.msg
 
   connman.loop = asyncLoop(send)
 

lsquic/context/client.nim

@@ -9,14 +9,14 @@ import ../helpers/sequninit
 proc onNewConn(
     stream_if_ctx: pointer, conn: ptr lsquic_conn_t
 ): ptr lsquic_conn_ctx_t {.cdecl.} =
-  trace "New connection established: client"
+  debug "New connection established: client"
   let conn_ctx = lsquic_conn_get_ctx(conn)
   cast[ptr lsquic_conn_ctx_t](conn_ctx)
 
 proc onHandshakeDone(
     conn: ptr lsquic_conn_t, status: enum_lsquic_hsk_status
 ) {.cdecl.} =
-  trace "Handshake done", status
+  debug "Handshake done", status
   let conn_ctx = lsquic_conn_get_ctx(conn)
   if conn_ctx.isNil:
     debug "conn_ctx is nil in onHandshakeDone"
@@ -34,7 +34,7 @@ proc onHandshakeDone(
     quicClientConn.connectedFut.complete()
 
 proc onConnClosed(conn: ptr lsquic_conn_t) {.cdecl.} =
-  trace "Connection closed: client"
+  debug "Connection closed: client"
   let conn_ctx = lsquic_conn_get_ctx(conn)
   if not conn_ctx.isNil:
     let quicClientConn = cast[QuicClientConn](conn_ctx)
@@ -52,12 +52,13 @@ proc onConnClosed(conn: ptr lsquic_conn_t) {.cdecl.} =
       )
     quicClientConn.cancelPending()
     quicClientConn.onClose()
+    GC_unref(quicClientConn)
   lsquic_conn_set_ctx(conn, nil)
 
 proc onNewStream(
     stream_if_ctx: pointer, stream: ptr lsquic_stream_t
 ): ptr lsquic_stream_ctx_t {.cdecl.} =
-  trace "New stream created: client"
+  debug "New stream created: client"
   let conn = lsquic_stream_conn(stream)
   let conn_ctx = lsquic_conn_get_ctx(conn)
   if conn_ctx.isNil:
@@ -90,7 +91,7 @@ method dial*(
 
   let quicClientConn =
     QuicClientConn(connectedFut: connectedFut, local: local, remote: remote)
-
+  GC_ref(quicClientConn) # Keep it pinned until on_conn_closed is called
   let conn = lsquic_engine_connect(
     ctx.engine,
     N_LSQVER,
@@ -163,5 +164,5 @@ proc new*(
 method makeStream*(
     ctx: ClientContext, quicConn: QuicConnection
 ) {.gcsafe, raises: [].} =
-  trace "Creating stream: client"
+  debug "Creating stream: client"
   lsquic_conn_make_stream(quicConn.lsquicConn)

lsquic/context/server.nim

@@ -9,7 +9,7 @@ import ../helpers/[sequninit, transportaddr]
 proc onNewConn(
     stream_if_ctx: pointer, conn: ptr lsquic_conn_t
 ): ptr lsquic_conn_ctx_t {.cdecl.} =
-  trace "New connection established: server"
+  debug "New connection established: server"
   var local: ptr SockAddr
   var remote: ptr SockAddr
   discard lsquic_conn_get_sockaddr(conn, addr local, addr remote)
@@ -19,12 +19,13 @@ proc onNewConn(
     remote: remote.toTransportAddress(),
     lsquicConn: conn,
   )
+  GC_ref(quicServerConn) # Keep it pinned until on_conn_closed is called
   let serverCtx = cast[ServerContext](stream_if_ctx)
   serverCtx.incoming.putNoWait(quicServerConn)
   cast[ptr lsquic_conn_ctx_t](quicServerConn)
 
 proc onConnClosed(conn: ptr lsquic_conn_t) {.cdecl.} =
-  trace "Connection closed: server"
+  debug "Connection closed: server"
   let conn_ctx = lsquic_conn_get_ctx(conn)
   if not conn_ctx.isNil:
     let quicConn = cast[QuicConnection](conn_ctx)
@@ -34,7 +35,7 @@ proc onConnClosed(conn: ptr lsquic_conn_t) {.cdecl.} =
 proc onNewStream(
     stream_if_ctx: pointer, stream: ptr lsquic_stream_t
 ): ptr lsquic_stream_ctx_t {.cdecl.} =
-  trace "New stream created: server"
+  debug "New stream created: server"
   let conn = lsquic_stream_conn(stream)
   let conn_ctx = lsquic_conn_get_ctx(conn)
   if conn_ctx.isNil:

lsquic/context/stream.nim

@@ -4,7 +4,7 @@ import ../[lsquic_ffi, stream]
 import ../helpers/sequninit
 
 proc onClose*(stream: ptr lsquic_stream_t, ctx: ptr lsquic_stream_ctx_t) {.cdecl.} =
-  trace "Stream closed"
+  debug "Stream closed"
   if ctx.isNil:
     debug "stream_ctx is nil onClose"
     return
@@ -13,7 +13,8 @@ proc onClose*(stream: ptr lsquic_stream_t, ctx: ptr lsquic_stream_ctx_t) {.cdecl
   if not streamCtx.closeWrite:
     streamCtx.isEof = true
     streamCtx.closed.fire()
-    streamCtx.abortPendingWrites("stream closed")
+    streamCtx.abortPendingWrites("stream closed 4")
+  GC_unref(streamCtx)
 
 type StreamReadContext = object
   stream: ptr lsquic_stream_t
@@ -40,12 +41,11 @@ proc onRead*(stream: ptr lsquic_stream_t, ctx: ptr lsquic_stream_ctx_t) {.cdecl.
   if nread < 0:
     error "could not read from stream", nread, streamId = lsquic_stream_id(stream)
     streamCtx.abort()
-    
+
   if lsquic_stream_wantread(stream, 0) == -1:
     error "could not set stream wantread", streamId = lsquic_stream_id(stream)
     streamCtx.abort()
 
-
 proc onWrite*(stream: ptr lsquic_stream_t, ctx: ptr lsquic_stream_ctx_t) {.cdecl.} =
   trace "onWrite"
 
@@ -55,6 +55,8 @@ proc onWrite*(stream: ptr lsquic_stream_t, ctx: ptr lsquic_stream_ctx_t) {.cdecl
 
   let streamCtx = cast[Stream](ctx)
   if streamCtx.toWrite.len == 0:
+    if not streamCtx.shouldClose.isNil and not streamCtx.shouldClose.finished:
+      streamCtx.shouldClose.complete()
     if lsquic_stream_wantwrite(stream, 0) == -1:
       error "could not set stream wantwrite", streamId = lsquic_stream_id(stream)
       streamCtx.abort()
@@ -71,7 +73,6 @@ proc onWrite*(stream: ptr lsquic_stream_t, ctx: ptr lsquic_stream_ctx_t) {.cdecl
     let p = w.data[w.offset].addr
     let nAvail = (w.data.len - w.offset).csize_t
     let n: ssize_t = lsquic_stream_write(stream, p, nAvail)
-
     if n > 0:
       w.offset += n.int
       if w.offset >= w.data.len:
@@ -82,12 +83,15 @@ proc onWrite*(stream: ptr lsquic_stream_t, ctx: ptr lsquic_stream_ctx_t) {.cdecl
           return
         streamCtx.toWrite.delete(0)
     elif n == 0:
+      # Nothing to write
       break
     else:
       streamCtx.abortPendingWrites("write failed")
       return
 
   if streamCtx.toWrite.len == 0:
+    if not streamCtx.shouldClose.isNil and not streamCtx.shouldClose.finished:
+      streamCtx.shouldClose.complete()
     if lsquic_stream_wantwrite(stream, 0) == -1:
       error "could not set stream wantwrite", streamId = lsquic_stream_id(stream)
       streamCtx.abort()

lsquic/stream.nim

@@ -16,13 +16,16 @@ type Stream* = ref object
   closed*: AsyncEvent # This is called when on_close callback is executed
   isEof*: bool # Received a FIN from remote
   toWrite*: seq[WriteTask]
+  shouldClose*: Future[void].Raising([CancelledError])
 
 proc new*(T: typedesc[Stream], quicStream: ptr lsquic_stream_t = nil): T =
-  Stream(
+  let s = Stream(
     quicStream: quicStream,
     incoming: newAsyncQueue[seq[byte]](),
     closed: newAsyncEvent(),
   )
+  GC_ref(s) # Keep it pinned until stream_if.on_close is executed
+  s
 
 proc abortPendingWrites*(stream: Stream, reason: string = "") =
   for pendingWrite in stream.toWrite.mitems:
@@ -46,10 +49,14 @@ proc abort*(stream: Stream) =
   stream.abortPendingWrites("stream aborted")
   stream.closed.fire()
 
-proc close*(stream: Stream) =
+proc close*(stream: Stream) {.async: (raises: [StreamError, CancelledError]).} =
   if stream.closeWrite:
     return
 
+  if stream.toWrite.len != 0:
+    stream.shouldClose = Future[void].Raising([CancelledError]).init()
+    await stream.shouldClose
+
   # Closing only the write side
   let ret = lsquic_stream_shutdown(stream.quicStream, 1)
   if ret == 0:
@@ -78,7 +85,7 @@ proc read*(
     await incomingFut.cancelAndWait()
     stream.isEof = true
     stream.closeWrite = true
-    raise newException(StreamError, "stream closed")
+    raise newException(StreamError, "stream closed 1")
 
   let incoming = await incomingFut
   if incoming.len == 0:
@@ -95,15 +102,16 @@ proc write*(
     stream: Stream, data: seq[byte]
 ) {.async: (raises: [CancelledError, StreamError]).} =
   if stream.closeWrite:
-    raise newException(StreamError, "stream is closed")
+    raise newException(StreamError, "stream closed 3")
 
   let closedFut = stream.closed.wait()
   let doneFut = Future[void].Raising([CancelledError, StreamError]).init()
   stream.toWrite.add(WriteTask(data: data, doneFut: doneFut))
   discard lsquic_stream_wantwrite(stream.quicStream, 1)
   let raceFut = await race(closedFut, doneFut)
   if raceFut == closedFut:
-    doneFut.fail(newException(StreamError, "stream closed"))
+    if not doneFut.finished:
+      doneFut.fail(newException(StreamError, "stream closed 2"))
     stream.closeWrite = true
 
   await doneFut