From c3404f6eea9a62a0e17957ae155b3fb56474ae8e Mon Sep 17 00:00:00 2001
From: Giovanni Petrantoni <sinkingsugar@gmail.com>
Date: Wed, 22 Jul 2020 00:03:41 +0900
Subject: [PATCH] Handle cancellation in timeoutMonitor (#283)

* Handle cancellation in timeoutMonitor

* refactor lpchannel timeout as suggested by cheatfate
---
 libp2p/muxers/mplex/lpchannel.nim | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/libp2p/muxers/mplex/lpchannel.nim b/libp2p/muxers/mplex/lpchannel.nim
index 6fb6221d1..1c4aa8e52 100644
--- a/libp2p/muxers/mplex/lpchannel.nim
+++ b/libp2p/muxers/mplex/lpchannel.nim
@@ -57,7 +57,6 @@ type
     msgCode*: MessageType         # cached in/out message code
     closeCode*: MessageType       # cached in/out close code
     resetCode*: MessageType       # cached in/out reset code
-    timerFut: Future[void]        # the current timer instanse
     timerTaskFut: Future[void]    # the current timer instanse
 
 proc open*(s: LPChannel) {.async, gcsafe.}
@@ -82,11 +81,8 @@ template withEOFExceptions(body: untyped): untyped =
 
 proc cleanupTimer(s: LPChannel) {.async.} =
   ## cleanup timers
-  ##
-  if not(isNil(s.timerFut)) and
-    not(s.timerFut.finished):
-    s.timerFut.cancel()
-    await s.timerTaskFut
+  if not s.timerTaskFut.finished:
+    await s.timerTaskFut.cancelAndWait()
 
 proc closeMessage(s: LPChannel) {.async.} =
   logScope:
@@ -248,13 +244,10 @@ proc timeoutMonitor(s: LPChannel) {.async.} =
   ## be reset
   ##
 
-  if not(isNil(s.timerFut)):
-    return
-
   try:
     while true:
-      s.timerFut = sleepAsync(s.timeout)
-      await s.timerFut
+      await sleepAsync(s.timeout)
+
       if s.closed or s.atEof:
         return
 
@@ -267,6 +260,8 @@ proc timeoutMonitor(s: LPChannel) {.async.} =
     # reset channel on innactivity timeout
     trace "channel timed out, resetting"
     await s.reset()
+  except CancelledError as exc:
+    raise exc
   except CatchableError as exc:
     trace "exception in timeout", exc = exc.msg
 
@@ -328,7 +323,6 @@ proc init*(
       await conn.writeMsg(chann.id,
                           chann.msgCode,
                           data)
-
     except CatchableError as exc:
       trace "exception in lpchannel write handler", exc = exc.msg
       await chann.reset()