File descriptors are now removed from fds list explicitly in close().

Fixes tasyncawait on linux.

3 files changed, 25 insertions, 23 deletions

diff --git a/lib/pure/asyncio2.nim b/lib/pure/asyncio2.nim
index c37370b7b..0ea4d7e79 100644
--- a/lib/pure/asyncio2.nim
+++ b/lib/pure/asyncio2.nim
@@ -508,6 +508,10 @@ else:
     result = socket(domain, typ, protocol)
     disp.register(result)
   
+  proc close*(disp: PDispatcher, sock: TSocketHandle) =
+    sock.close()
+    disp.selector.unregister(sock)
+
   proc addRead(p: PDispatcher, sock: TSocketHandle, cb: TCallback) =
     if sock notin p.selector:
       raise newException(EInvalidValue, "File descriptor not registered.")
@@ -892,7 +896,7 @@ proc recvLine*(p: PDispatcher, socket: TSocketHandle): PFuture[string] {.async.}
 when isMainModule:
   
   var p = newDispatcher()
-  var sock = socket()
+  var sock = p.socket()
   sock.setBlocking false
 
 
diff --git a/lib/pure/selectors.nim b/lib/pure/selectors.nim
index a113e3362..085344e3e 100644
--- a/lib/pure/selectors.nim
+++ b/lib/pure/selectors.nim
@@ -60,7 +60,6 @@ when defined(linux) or defined(nimdoc):
       events: set[TEvent]): PSelectorKey {.discardable.} =
     ## Updates the events which ``fd`` wants notifications for.
     if s.fds[fd].events != events:
-      echo("Update ", fd.cint, " to ", events)
       var event = createEventStruct(events, fd)
       
       s.fds[fd].events = events
@@ -70,7 +69,9 @@ when defined(linux) or defined(nimdoc):
   
   proc unregister*(s: PSelector, fd: TSocketHandle): PSelectorKey {.discardable.} =
     if epoll_ctl(s.epollFD, EPOLL_CTL_DEL, fd, nil) != 0:
-      OSError(OSLastError())
+      let err = OSLastError()
+      if err.cint notin {ENOENT, EBADF}: # TODO: Why do we sometimes get an EBADF? Is this normal?
+        OSError(err)
     result = s.fds[fd]
     s.fds.del(fd)
 
@@ -78,6 +79,15 @@ when defined(linux) or defined(nimdoc):
     if s.epollFD.close() != 0: OSError(OSLastError())
     dealloc(addr s.events) # TODO: Test this
   
+  proc epollHasFd(s: PSelector, fd: TSocketHandle): bool =
+    result = true
+    var event = createEventStruct(s.fds[fd].events, fd)
+    if epoll_ctl(s.epollFD, EPOLL_CTL_MOD, fd, addr(event)) != 0:
+      let err = osLastError()
+      if err.cint in {ENOENT, EBADF}:
+        return false
+      OSError(OSLastError())
+  
   proc select*(s: PSelector, timeout: int): seq[TReadyInfo] =
     ##
     ## The ``events`` field of the returned ``key`` contains the original events
@@ -85,24 +95,19 @@ when defined(linux) or defined(nimdoc):
     ## of the ``TReadyInfo`` tuple which determines which events are ready
     ## on the ``fd``.
     result = @[]
-    
     let evNum = epoll_wait(s.epollFD, addr s.events[0], 64.cint, timeout.cint)
     if evNum < 0: OSError(OSLastError())
     if evNum == 0: return @[]
     for i in 0 .. <evNum:
+      let fd = s.events[i].data.fd.TSocketHandle
+    
       var evSet: set[TEvent] = {}
       if (s.events[i].events and EPOLLIN) != 0: evSet = evSet + {EvRead}
       if (s.events[i].events and EPOLLOUT) != 0: evSet = evSet + {EvWrite}
-      let selectorKey = s.fds[s.events[i].data.fd.TSocketHandle]
+      let selectorKey = s.fds[fd]
       assert selectorKey != nil
       result.add((selectorKey, evSet))
-  
-      if (s.events[i].events and EPOLLHUP) != 0 or
-         (s.events[i].events and EPOLLRDHUP) != 0:
-        # fd closed
-        #echo("fd closed ", s.events[i].data.fd)
-        s.unregister(s.events[i].data.fd.TSocketHandle)
-  
+
       #echo("Epoll: ", result[i].key.fd, " ", result[i].events, " ", result[i].key.events)
   
   proc newSelector*(): PSelector =
@@ -116,15 +121,8 @@ when defined(linux) or defined(nimdoc):
   proc contains*(s: PSelector, fd: TSocketHandle): bool =
     ## Determines whether selector contains a file descriptor.
     if s.fds.hasKey(fd):
-      result = true
-      
       # Ensure the underlying epoll instance still contains this fd.
-      var event = createEventStruct(s.fds[fd].events, fd)
-      if epoll_ctl(s.epollFD, EPOLL_CTL_MOD, fd, addr(event)) != 0:
-        let err = osLastError()
-        if err.cint in {ENOENT, EBADF}:
-          return false
-        OSError(OSLastError())
+      result = epollHasFd(s, fd)
     else:
       return false
 
diff --git a/tests/async/tasyncawait.nim b/tests/async/tasyncawait.nim
index fea0783a0..34dcdc773 100644
--- a/tests/async/tasyncawait.nim
+++ b/tests/async/tasyncawait.nim
@@ -26,19 +26,19 @@ proc launchSwarm(disp: PDispatcher, port: TPort): PFuture[int] {.async.} =
     discard await disp.connect(sock, "localhost", port)
     when true:
       discard await sendMessages(disp, sock)
-      sock.close()
+      disp.close(sock)
     else:
       # Issue #932: https://github.com/Araq/Nimrod/issues/932
       var msgFut = sendMessages(disp, sock)
       msgFut.callback =
         proc () =
-          sock.close()
+          disp.close(sock)
 
 proc readMessages(disp: PDispatcher, client: TSocketHandle): PFuture[int] {.async.} =
   while true:
     var line = await disp.recvLine(client)
     if line == "":
-      client.close()
+      disp.close(client)
       clientCount.inc
       break
     else: