loader: use recvData instead of readData

recvData is a new method for PosixStream that does less weird magic than
readData.

Also, allow duplicates in unregWrite/unregRead; it's simpler to live
with them than to prevent them.

3 files changed, 35 insertions, 16 deletions

diff --git a/src/io/posixstream.nim b/src/io/posixstream.nim
index 04fe0e5c..66b2d0d9 100644
--- a/src/io/posixstream.nim
+++ b/src/io/posixstream.nim
@@ -62,6 +62,16 @@ proc psReadData(s: Stream, buffer: pointer, len: int): int =
   if result == -1:
     raisePosixIOError()
 
+method recvData*(s: PosixStream, buffer: pointer, len: int): int {.base.} =
+  let n = read(s.fd, buffer, len)
+  if n < 0:
+    raisePosixIOError()
+  if n == 0:
+    if unlikely(s.isend):
+      raise newException(EOFError, "eof")
+    s.isend = true
+  return n
+
 method sendData*(s: PosixStream, buffer: pointer, len: int): int {.base.} =
   #TODO use sendData instead
   let n = write(s.fd, buffer, len)
diff --git a/src/io/socketstream.nim b/src/io/socketstream.nim
index fb378083..38c43a84 100644
--- a/src/io/socketstream.nim
+++ b/src/io/socketstream.nim
@@ -49,6 +49,16 @@ proc sockWriteData(s: Stream, buffer: pointer, len: int) =
       raisePosixIOError()
     i += n
 
+method recvData*(s: SocketStream, buffer: pointer, len: int): int =
+  let n = s.source.recv(buffer, len)
+  if n < 0:
+    raisePosixIOError()
+  if n == 0:
+    if unlikely(s.isend):
+      raise newException(EOFError, "eof")
+    s.isend = true
+  return n
+
 method sendData*(s: SocketStream, buffer: pointer, len: int): int =
   let n = s.source.send(buffer, len)
   if n < 0:
diff --git a/src/loader/loader.nim b/src/loader/loader.nim
index 3b4e016f..491ca095 100644
--- a/src/loader/loader.nim
+++ b/src/loader/loader.nim
@@ -286,7 +286,7 @@ proc runFileLoader*(fd: cint, config: LoaderConfig) =
           while true:
             let buffer = newLoaderBuffer()
             try:
-              buffer.len = handle.istream.readData(addr buffer[0], buffer.cap)
+              buffer.len = handle.istream.recvData(addr buffer[0], buffer.cap)
               if buffer.len == 0:
                 dealloc(buffer)
                 break
@@ -318,31 +318,30 @@ proc runFileLoader*(fd: cint, config: LoaderConfig) =
           except ErrorBrokenPipe: # receiver died; stop streaming
             unregWrite.add(handle)
             break
-        if handle.istream == nil and handle.currentBuffer == nil and
-            (unregWrite.len == 0 or unregWrite[^1] != handle):
+        if handle.istream == nil and handle.currentBuffer == nil:
           # after EOF, but not appended in this send cycle
           unregWrite.add(handle)
       if Error in event.events:
         assert event.fd != ctx.fd
         let handle = ctx.handleMap[event.fd]
-        if handle.fd == event.fd:
-          if unregWrite.len == 0 or unregWrite[^1] != handle: # ostream died
-            unregWrite.add(handle)
+        if handle.fd == event.fd: # ostream died
+          unregWrite.add(handle)
         else: # istream died
           unregRead.add(handle)
+    # Unregister handles queued for unregistration.
+    # It is possible for both unregRead and unregWrite to contain duplicates. To
+    # avoid double-close/double-unregister, we set the istream/ostream of
+    # unregistered handles to nil.
     for handle in unregRead:
-      ctx.selector.unregister(handle.istream.fd)
-      ctx.handleMap.del(handle.istream.fd)
-      handle.istream.close()
-      handle.istream = nil
-      if handle.currentBuffer == nil:
-        unregWrite.add(handle)
-      #TODO TODO TODO what to do about sostream
+      if handle.istream != nil:
+        ctx.selector.unregister(handle.istream.fd)
+        ctx.handleMap.del(handle.istream.fd)
+        handle.istream.close()
+        handle.istream = nil
+        if handle.currentBuffer == nil:
+          unregWrite.add(handle)
     for handle in unregWrite:
       if handle.ostream != nil:
-        # if the previous loop adds its handle to this one, it is possible that
-        # we try to unregister the same handle twice
-        #TODO this is kind of a mess
         ctx.selector.unregister(handle.fd)
         ctx.handleMap.del(handle.fd)
         handle.ostream.close()