New state for quic connection: disconnecting

When a quic connection is disconnecting, close the
send loop and the incoming udp connection.

quic/connection.nim

@@ -14,7 +14,6 @@ type
     udp: DatagramTransport
     quic: QuicConnection
     loop: Future[void]
-    closed: bool
     onClose: proc()
   IncomingConnection = ref object of Connection
   OutgoingConnection = ref object of Connection
@@ -40,18 +39,36 @@ proc startSending(connection: Connection, remote: TransportAddress) =
 proc stopSending(connection: Connection) {.async.} =
   await connection.loop.cancelAndWait()
 
+method closeUdp(connection: Connection) {.async, base.} =
+  discard
+
+method closeUdp(connection: OutgoingConnection) {.async.} =
+  await connection.udp.closeWait()
+
+proc disconnect(connection: Connection) {.async.} =
+  await connection.stopSending()
+  await connection.closeUdp()
+  if connection.onClose != nil:
+    connection.onClose()
+
 proc newIncomingConnection*(udp: DatagramTransport,
                            remote: TransportAddress): Connection =
   let datagram = Datagram(data: udp.getMessage())
   let quic = newQuicServerConnection(udp.localAddress, remote, datagram)
-  result = IncomingConnection(udp: udp, quic: quic)
-  result.startSending(remote)
+  let connection = IncomingConnection(udp: udp, quic: quic)
+  quic.disconnect = proc {.async.} =
+    await connection.disconnect()
+  connection.startSending(remote)
+  connection
 
 proc newOutgoingConnection*(udp: DatagramTransport,
                            remote: TransportAddress): Connection =
   let quic = newQuicClientConnection(udp.localAddress, remote)
-  result = OutgoingConnection(udp: udp, quic: quic)
-  result.startSending(remote)
+  let connection = OutgoingConnection(udp: udp, quic: quic)
+  quic.disconnect = proc {.async.} =
+    await connection.disconnect()
+  connection.startSending(remote)
+  connection
 
 proc startHandshake*(connection: Connection) =
   connection.quic.send()
@@ -66,21 +83,8 @@ proc openStream*(connection: Connection): Future[Stream] {.async.} =
 proc incomingStream*(connection: Connection): Future[Stream] {.async.} =
   result = await connection.quic.incomingStream()
 
-method closeUdp(connection: Connection) {.async, base.} =
-  discard
-
-method closeUdp(connection: OutgoingConnection) {.async.} =
-  await connection.udp.closeWait()
-
 proc drop*(connection: Connection) {.async.} =
-  if not connection.closed:
-    connection.closed = true
-    await connection.stopSending()
-    await connection.closeUdp()
-    if connection.onClose != nil:
-      connection.onClose()
-    connection.quic.drop()
+  await connection.quic.drop()
 
 proc close*(connection: Connection) {.async.} =
   await connection.quic.close()
-  await connection.drop()

quic/transport/ngtcp2/connection/closedstate.nim

@@ -31,7 +31,7 @@ method openStream(state: ClosedConnection): Future[Stream] {.async.} =
 method close(state: ClosedConnection) {.async.} =
   discard
 
-method drop(state: ClosedConnection) =
+method drop(state: ClosedConnection) {.async.} =
   discard
 
 {.pop.}

quic/transport/ngtcp2/connection/disconnectingstate.nim

@@ -0,0 +1,51 @@
+import pkg/chronos
+import ../../../udp/datagram
+import ../../quicconnection
+import ../../connectionid
+import ../../stream
+import ./closedstate
+
+type
+  DisconnectingConnection* = ref object of ConnectionState
+    connection: QuicConnection
+    disconnect: Future[void]
+    ids: seq[ConnectionId]
+
+proc newDisconnectingConnection*(ids: seq[ConnectionId]):
+                                            DisconnectingConnection =
+  DisconnectingConnection(ids: ids)
+
+proc callDisconnect(connection: QuicConnection) {.async.} =
+  if connection.disconnect != nil:
+    await connection.disconnect()
+
+{.push locks: "unknown".}
+
+method ids*(state: DisconnectingConnection): seq[ConnectionId] =
+  state.ids
+
+method enter(state: DisconnectingConnection, connection: QuicConnection) =
+  state.connection = connection
+  state.disconnect = callDisconnect(connection)
+
+method leave(state: DisconnectingConnection) =
+  state.connection = nil
+
+method send(state: DisconnectingConnection) =
+  raise newException(ClosedConnectionError, "connection is disconnecting")
+
+method receive(state: DisconnectingConnection, datagram: Datagram) =
+  discard
+
+method openStream(state: DisconnectingConnection): Future[Stream] {.async.} =
+  raise newException(ClosedConnectionError, "connection is disconnecting")
+
+method close(state: DisconnectingConnection) {.async.} =
+  await state.disconnect
+  state.connection.switch(newClosedConnection())
+
+method drop(state: DisconnectingConnection) {.async.} =
+  await state.disconnect
+  state.connection.switch(newClosedConnection())
+
+{.pop.}

quic/transport/ngtcp2/connection/drainingstate.nim

@@ -4,6 +4,7 @@ import ../../quicconnection
 import ../../connectionid
 import ../../stream
 import ../../timeout
+import ./disconnectingstate
 import ./closedstate
 
 type
@@ -54,8 +55,13 @@ method openStream(state: DrainingConnection): Future[Stream] {.async.} =
 
 method close(state: DrainingConnection) {.async.} =
   await state.done.wait()
+  let disconnecting = newDisconnectingConnection(state.ids)
+  state.connection.switch(disconnecting)
+  await disconnecting.close()
 
-method drop(state: DrainingConnection) =
-  state.connection.switch(newClosedConnection())
+method drop(state: DrainingConnection) {.async.} =
+  let disconnecting = newDisconnectingConnection(state.ids)
+  state.connection.switch(disconnecting)
+  await disconnecting.drop()
 
 {.pop.}

quic/transport/ngtcp2/connection/openstate.nim

@@ -6,7 +6,7 @@ import ../../stream
 import ../connection
 import ../streams
 import ./closingstate
-import ./closedstate
+import ./disconnectingstate
 
 type
   OpenConnection* = ref object of ConnectionState
@@ -46,8 +46,10 @@ method close(state: OpenConnection) {.async.} =
   state.quicConnection.switch(closing)
   await closing.close()
 
-method drop(state: OpenConnection) =
-  state.quicConnection.switch(newClosedConnection())
+method drop(state: OpenConnection) {.async.} =
+  let disconnecting = newDisconnectingConnection(state.ids)
+  state.quicConnection.switch(disconnecting)
+  await disconnecting.drop()
 
 method `onNewId=`*(state: OpenConnection, callback: IdCallback) =
   state.ngtcp2Connection.onNewId = callback

quic/transport/quicconnection.nim

@@ -10,6 +10,7 @@ type
     incoming*: AsyncQueue[Stream]
     handshake*: AsyncEvent
     closed*: AsyncEvent
+    disconnect*: proc(): Future[void] {.gcsafe.}
   ConnectionState* = ref object of RootObj
   IdCallback* = proc(id: ConnectionId)
   ConnectionError* = object of IOError
@@ -34,7 +35,7 @@ method receive*(state: ConnectionState, datagram: Datagram) =
 method openStream*(state: ConnectionState): Future[Stream] =
   doAssert false # override this method
 
-method drop*(state: ConnectionState) =
+method drop*(state: ConnectionState): Future[void] =
   doAssert false # override this method
 
 method close*(state: ConnectionState): Future[void] =
@@ -88,5 +89,5 @@ proc incomingStream*(connection: QuicConnection): Future[Stream] =
 proc close*(connection: QuicConnection): Future[void] =
   connection.state.close()
 
-proc drop*(connection: QuicConnection) =
+proc drop*(connection: QuicConnection): Future[void] =
   connection.state.drop()

tests/quic/testQuicConnection.nim

@@ -12,20 +12,20 @@ suite "quic connection":
 
   asynctest "sends outgoing datagrams":
     let client = newQuicClientConnection(zeroAddress, zeroAddress)
-    defer: client.drop()
+    defer: await client.drop()
     client.send()
     let datagram = await client.outgoing.get()
     check datagram.len > 0
 
   asynctest "processes received datagrams":
     let client = newQuicClientConnection(zeroAddress, zeroAddress)
-    defer: client.drop()
+    defer: await client.drop()
 
     client.send()
     let datagram = await client.outgoing.get()
 
     let server = newQuicServerConnection(zeroAddress, zeroAddress, datagram)
-    defer: server.drop()
+    defer: await server.drop()
 
     server.receive(datagram)
 
@@ -37,8 +37,8 @@ suite "quic connection":
 
   asynctest "performs handshake":
     let (client, server) = await performHandshake()
-    defer: client.drop()
-    defer: server.drop()
+    defer: await client.drop()
+    defer: await server.drop()
 
     check client.handshake.isSet()
     check server.handshake.isSet()
@@ -46,8 +46,8 @@ suite "quic connection":
   asynctest "performs handshake multiple times":
     for i in 1..100:
       let (client, server) = await performHandshake()
-      client.drop()
-      server.drop()
+      await client.drop()
+      await server.drop()
 
   asynctest "returns the current connection ids":
     let (client, server) = await setupConnection()
@@ -57,8 +57,8 @@ suite "quic connection":
 
   asynctest "notifies about id changes":
     let (client, server) = await setupConnection()
-    defer: client.drop
-    defer: server.drop
+    defer: await client.drop
+    defer: await server.drop
 
     var newId: ConnectionId
     server.onNewId = proc (id: ConnectionId) =
@@ -72,12 +72,12 @@ suite "quic connection":
 
   asynctest "fires event when closed":
     let client = newQuicClientConnection(zeroAddress, zeroAddress)
-    client.drop()
+    await client.drop()
     check client.closed.isSet()
 
   asynctest "raises ConnectionError when closed":
     let connection = newQuicClientConnection(zeroAddress, zeroAddress)
-    connection.drop()
+    await connection.drop()
 
     expect ConnectionError:
       connection.send()
@@ -88,8 +88,8 @@ suite "quic connection":
     expect ConnectionError:
       discard await connection.openStream()
 
-  test "has empty list of ids when closed":
+  asynctest "has empty list of ids when closed":
     let connection = newQuicClientConnection(zeroAddress, zeroAddress)
-    connection.drop()
+    await connection.drop()
 
     check connection.ids.len == 0

tests/quic/testStreams.nim

@@ -15,8 +15,8 @@ suite "streams":
     let (client, server) = await performHandshake()
     check client.openStream() != client.openStream()
 
-    client.drop()
-    server.drop()
+    await client.drop()
+    await server.drop()
 
   asynctest "closes stream":
     let (client, server) = await performHandshake()
@@ -24,8 +24,8 @@ suite "streams":
 
     await stream.close()
 
-    client.drop()
-    server.drop()
+    await client.drop()
+    await server.drop()
 
   asynctest "writes to stream":
     let (client, server) = await performHandshake()
@@ -35,8 +35,8 @@ suite "streams":
 
     check client.outgoing.anyIt(it.data.contains(message))
 
-    client.drop()
-    server.drop()
+    await client.drop()
+    await server.drop()
 
   asynctest "writes zero-length message":
     let (client, server) = await performHandshake()
@@ -46,8 +46,8 @@ suite "streams":
 
     check datagram.len > 0
 
-    client.drop()
-    server.drop()
+    await client.drop()
+    await server.drop()
 
   asynctest "raises when reading from or writing to closed stream":
     let (client, server) = await performHandshake()
@@ -60,8 +60,8 @@ suite "streams":
     expect IOError:
       await stream.write(@[1'u8, 2'u8, 3'u8])
 
-    client.drop()
-    server.drop()
+    await client.drop()
+    await server.drop()
 
   asynctest "accepts incoming streams":
     let (client, server) = await performHandshake()
@@ -74,8 +74,8 @@ suite "streams":
     check clientStream.id == serverStream.id
 
     await simulation.cancelAndWait()
-    client.drop()
-    server.drop()
+    await client.drop()
+    await server.drop()
 
   asynctest "reads from stream":
     let (client, server) = await performHandshake()
@@ -92,8 +92,8 @@ suite "streams":
     check incoming == message
 
     await simulation.cancelAndWait()
-    client.drop()
-    server.drop()
+    await client.drop()
+    await server.drop()
 
   asynctest "writes long messages to stream":
     let (client, server) = await performHandshake()
@@ -108,8 +108,8 @@ suite "streams":
       discard await incoming.read()
 
     await simulation.cancelAndWait()
-    client.drop()
-    server.drop()
+    await client.drop()
+    await server.drop()
 
 
   asynctest "handles packet loss":
@@ -126,5 +126,5 @@ suite "streams":
     check incoming == message
 
     await simulation.cancelAndWait()
-    client.drop()
-    server.drop()
+    await client.drop()
+    await server.drop()