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/context/client.nim

@@ -52,6 +52,7 @@ proc onConnClosed(conn: ptr lsquic_conn_t) {.cdecl.} =
       )
     quicClientConn.cancelPending()
     quicClientConn.onClose()
+    GC_unref(quicClientConn)
   lsquic_conn_set_ctx(conn, nil)
 
 proc onNewStream(
@@ -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,

lsquic/context/server.nim

@@ -19,6 +19,7 @@ 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)

lsquic/context/stream.nim

@@ -12,8 +12,10 @@ proc onClose*(stream: ptr lsquic_stream_t, ctx: ptr lsquic_stream_ctx_t) {.cdecl
   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
@@ -40,12 +42,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"
 

lsquic/stream.nim

@@ -18,11 +18,13 @@ type Stream* = ref object
   toWrite*: seq[WriteTask]
 
 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,6 +35,7 @@ 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
@@ -44,6 +47,7 @@ proc abort*(stream: Stream) =
   stream.closeWrite = true
   stream.isEof = true
   stream.abortPendingWrites("stream aborted")
+  echo "FIRING CLOSE 2"
   stream.closed.fire()
 
 proc close*(stream: Stream) =
@@ -78,7 +82,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 +99,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

tests/test_connection.nim

@@ -11,6 +11,8 @@ import ./helpers/certificate
 import lsquic/certificateverifier
 import lsquic/stream
 import lsquic/lsquic_ffi
+import stew/endians2
+import sequtils
 
 proc logging(ctx: pointer, buf: cstring, len: csize_t): cint {.cdecl.} =
   echo $buf
@@ -21,10 +23,104 @@ proc certificateCb(
 ): bool {.gcsafe.} =
   return derCertificates.len > 0
 
-suite "connections":
-  setup:
-    let address = initTAddress("127.0.0.1:12345")
+let address = initTAddress("127.0.0.1:12345")
+
+const
+  runs = 1
+  uploadSize = 100000 # 100KB
+  downloadSize = 100000000 # 100MB
+  chunkSize = 65536 # 64KB chunks like perf
+
+proc runPerf(): Future[Duration] {.async.} =
+  let customCertVerif: CertificateVerifier =
+    CustomCertificateVerifier.init(certificateCb)
+  let clientTLSConfig = TLSConfig.new(
+    testCertificate(),
+    testPrivateKey(),
+    @["test"].toHashSet(),
+    Opt.some(customCertVerif),
+  )
+  let serverTLSConfig = TLSConfig.new(
+    testCertificate(),
+    testPrivateKey(),
+    @["test"].toHashSet(),
+    Opt.some(customCertVerif),
+  )
+  let client = QuicClient.new(clientTLSConfig)
+  let server = QuicServer.new(serverTLSConfig)
+  let listener = server.listen(address)
+  let accepting = listener.accept()
+  let dialing = client.dial(address)
+
+  let outgoingConn = await dialing
+  let incomingConn = await accepting
+
+  let serverHandler = proc() {.async.} =
+    let stream = await incomingConn.incomingStream()
+
+    # Step 1: Read download size (8 bytes) 
+    let clientDownloadSize = await stream.read()
+
+    # Step 2: Read upload data until EOF
+    var totalBytesRead = 0
+    while true:
+      let chunk = await stream.read()
+      if chunk.len == 0:
+        break
+      totalBytesRead += chunk.len
+
+    # Step 3: Send download data back
+    var remainingToSend = uint64.fromBytesBE(clientDownloadSize)
+    while remainingToSend > 0:
+      let toSend = min(remainingToSend, chunkSize)
+      await stream.write(newSeq[byte](toSend))
+      remainingToSend -= toSend
+
+    stream.close()
+
+  # Start server handler
+  asyncSpawn serverHandler()
+
+  let startTime = Moment.now()
+
+  # Step 1: Send download size, activate stream first
+  let clientStream = await outgoingConn.openStream()
+  await clientStream.write(toSeq(downloadSize.uint64.toBytesBE()))
+
+  # Step 2: Send upload data in chunks
+  var remainingToSend = uploadSize
+  while remainingToSend > 0:
+    let toSend = min(remainingToSend, chunkSize)
+    await clientStream.write(newSeq[byte](toSend))
+    remainingToSend -= toSend
+
+  # Step 3: Close write side
+  clientStream.close()
+
+  # Step 4: Start reading download data
+  var totalDownloaded = 0
+  while totalDownloaded < downloadSize:
+    let chunk = await clientStream.read()
+    totalDownloaded += chunk.len
+
+  let duration = Moment.now() - startTime
+
+  await listener.stop()
+  await client.stop()
+
+  return duration
+
+suite "perf protocol simulation":
+  asyncTest "test":
+    var total: Duration
+    for i in 0 ..< runs:
+      let duration = await runPerf()
+      total += duration
+      echo "\trun #" & $(i + 1) & " duration: " & $duration
+
+    echo "\tavrg duration: " & $(total div runs)
 
+suite "tests":
   asyncTest "test":
     let logger = struct_lsquic_logger_if(log_buf: logging)
     discard lsquic_set_log_level("debug")
@@ -76,7 +172,6 @@ suite "connections":
 
     let incomingBehaviour = proc() {.async.} =
       try:
-        echo "HERE"
         let stream = await incomingConn.incomingStream()
         echo "Received stream in server"
 
@@ -107,12 +202,15 @@ suite "connections":
     outgoingConn.close()
     incomingConn.close()
 
+    # Cannot create a stream once closed
+    expect ConnectionError:
+      discard await outgoingConn.openStream()
+
     await sleepAsync(1.seconds)
 
     await client.stop()
     await listener.stop()
 
-    # TODO: perf example
     # TODO: destructors: (nice to have:)
     # - lsquic_global_cleanup() to free global resources. 
     # - lsquic_engine_destroy(engine)