From 5583168965c5bc6db4de3625ae0d20b4b4b29b30 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Mon, 18 May 2020 09:08:10 -0600
Subject: [PATCH] Track all connections (#175)

* adding to do for future refactor

* no mix conn and transport stopping or it will race

* don't use intermediary vars
---
 libp2p/switch.nim                  | 14 +++++++-------
 libp2p/transports/tcptransport.nim |  5 +++--
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/libp2p/switch.nim b/libp2p/switch.nim
index cf32bcc8e..2f284ab1a 100644
--- a/libp2p/switch.nim
+++ b/libp2p/switch.nim
@@ -335,16 +335,16 @@ proc stop*(s: Switch) {.async.} =
   # we want to report erros but we do not want to fail
   # or crash here, cos we need to clean possibly MANY items
   # and any following conn/transport won't be cleaned up
-  var futs = newSeq[Future[void]]()
-
   if s.pubSub.isSome:
-    futs &= s.pubSub.get().stop()
+    await s.pubSub.get().stop()
 
-  futs &= toSeq(s.connections.values).mapIt(s.cleanupConn(it))
-  futs &= s.transports.mapIt(it.close())
+  checkFutures(
+    await allFinished(
+    toSeq(s.connections.values).mapIt(s.cleanupConn(it))))
 
-  futs = await allFinished(futs)
-  checkFutures(futs)
+  checkFutures(
+    await allFinished(
+    s.transports.mapIt(it.close())))
 
   trace "switch stopped"
 
diff --git a/libp2p/transports/tcptransport.nim b/libp2p/transports/tcptransport.nim
index c9c886616..fd2d7c538 100644
--- a/libp2p/transports/tcptransport.nim
+++ b/libp2p/transports/tcptransport.nim
@@ -71,8 +71,9 @@ proc connHandler*(t: TcpTransport,
     if not isNil(t.handler):
       t.handlers &= t.handler(conn)
 
-    t.connections.add(conn)
-    t.cleanups &= t.cleanup(conn)
+  # TODO: store the streamtransport client here
+  t.connections.add(conn)
+  t.cleanups &= t.cleanup(conn)
 
   result = conn