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

@@ -53,6 +53,7 @@ proc newOutgoingConnection*(
   )
   conn.ensureClosedFut = conn.ensureClosed()
   conn.quicConn.onClose = proc() {.raises: [].} =
+    echo "OUTGOING CONN FIRE ONCLOSE"
     conn.closed.fire()
   conn
 
@@ -69,6 +70,7 @@ proc dial*(
     retFut.fail(newException(DialError, "could not dial: " & error))
     nil
   connection.quicConn.onClose = proc() {.raises: [].} =
+    echo "OUTGOING CONN FIRE DIAL"
     connection.closed.fire()
   retFut
 
@@ -91,6 +93,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 +113,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,14 @@ 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()
       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,16 +4,19 @@ 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
 
+  echo "CLOSING STREAM!!!!!!", lsquic_stream_id(stream)
   let streamCtx = cast[Stream](ctx)
   if not streamCtx.closeWrite:
     streamCtx.isEof = true
+    echo "FIRING CLOSE 3"
     streamCtx.closed.fire()
-    streamCtx.abortPendingWrites("stream closed")
+    streamCtx.abortPendingWrites("stream closed 4")
+  GC_unref(streamCtx)
 
 type StreamReadContext = object
   stream: ptr lsquic_stream_t
@@ -29,6 +32,7 @@ proc readCtxCb(
     copyMem(addr s[0], data, len)
     streamCtx.incoming.putNoWait(s)
   if fin != 0:
+    echo "FIN"
     streamCtx.incoming.putNoWait(@[])
   len
 
@@ -40,12 +44,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 +58,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()
@@ -88,6 +93,9 @@ proc onWrite*(stream: ptr lsquic_stream_t, ctx: ptr lsquic_stream_ctx_t) {.cdecl
       return
 
   if streamCtx.toWrite.len == 0:
+    echo streamCtx.shouldClose.isNil
+    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:
@@ -33,27 +36,38 @@ proc abortPendingWrites*(stream: Stream, reason: string = "") =
 proc abort*(stream: Stream) =
   if stream.closeWrite and stream.isEof:
     if not stream.closed.isSet():
+      echo "FIRING CLOSE 1"
       stream.closed.fire()
     stream.abortPendingWrites("stream aborted")
     return
 
+  echo "CLOSING ON ABORT"
   let ret = lsquic_stream_close(stream.quicStream)
   if ret != 0:
     error "could not abort stream", streamId = lsquic_stream_id(stream.quicStream)
 
   stream.closeWrite = true
   stream.isEof = true
   stream.abortPendingWrites("stream aborted")
+  echo "FIRING CLOSE 2"
   stream.closed.fire()
 
-proc close*(stream: Stream) =
+proc close*(stream: Stream) {.async: (raises: [StreamError, CancelledError]).} =
   if stream.closeWrite:
     return
+  echo "CALLED CLOSE!!!!!!!!!!!!!!!!!!"
+  if stream.toWrite.len != 0:
+    stream.shouldClose = Future[void].Raising([CancelledError]).init()
+
+    echo "AWAITING FOR CLOSE"
+    await stream.shouldClose
 
   # Closing only the write side
+  echo "SHUTDOWN WRITE"
   let ret = lsquic_stream_shutdown(stream.quicStream, 1)
   if ret == 0:
     if stream.isEof:
+      echo "CLOSE DUE TO EOF"
       if lsquic_stream_close(stream.quicStream) != 0:
         stream.abort()
         raise newException(StreamError, "could not close the stream")
@@ -68,22 +82,26 @@ proc read*(
     return @[]
 
   if lsquic_stream_wantread(stream.quicStream, 1) == -1:
+    echo "ABORTING"
     stream.abort()
     raise newException(StreamError, "could not set wantread")
 
   let incomingFut = stream.incoming.get()
   let closedFut = stream.closed.wait()
   let raceFut = await race(closedFut, incomingFut)
   if raceFut == closedFut:
+    echo "STREAM GOT CLOSED"
     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:
+    echo "INCOMING LEN IS 0"
     if stream.closeWrite:
       # We were already closed for write. Close the stream completely
+      echo "CLOSING AFTER READ IS EOF, and WRITE WAS CALLED BEFORE"
       if lsquic_stream_close(stream.quicStream) != 0:
         stream.abort()
         raise newException(StreamError, "could not close the stream")
@@ -95,15 +113,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