loaderiface: fix race in poll register

handleRead can register through fetch, so this doesn't work out as
nicely as in loader (where we control all register/unregister calls).
So now we queue up register events first, and only process them after
the "events" iterator exits.

3 files changed, 36 insertions, 11 deletions

diff --git a/src/local/pager.nim b/src/local/pager.nim
index 1db7f15b..e5510b23 100644
--- a/src/local/pager.nim
+++ b/src/local/pager.nim
@@ -3269,6 +3269,7 @@ proc inputLoop(pager: Pager) =
   while true:
     let timeout = pager.timeouts.sortAndGetTimeout()
     pager.pollData.poll(timeout)
+    pager.loader.blockRegister()
     for event in pager.pollData.events:
       let efd = int(event.fd)
       if (event.revents and POLLIN) != 0:
@@ -3285,8 +3286,9 @@ proc inputLoop(pager: Pager) =
       let container = pager.consoleWrapper.container
       if container != nil:
         container.tailOnLoad = true
-    pager.runJSJobs()
+    pager.loader.unblockRegister()
     pager.loader.unregistered.setLen(0)
+    pager.runJSJobs()
     pager.acceptBuffers()
     pager.runCommand()
     if pager.container == nil and pager.lineedit == nil:
@@ -3316,6 +3318,7 @@ proc headlessLoop(pager: Pager) =
   while pager.hasSelectFds():
     let timeout = pager.timeouts.sortAndGetTimeout()
     pager.pollData.poll(timeout)
+    pager.loader.blockRegister()
     for event in pager.pollData.events:
       let efd = int(event.fd)
       if (event.revents and POLLIN) != 0:
@@ -3324,9 +3327,10 @@ proc headlessLoop(pager: Pager) =
         pager.handleWrite(efd)
       if (event.revents and POLLERR) != 0 or (event.revents and POLLHUP) != 0:
         pager.handleError(efd)
+    pager.loader.unblockRegister()
+    pager.loader.unregistered.setLen(0)
     discard pager.timeouts.run(pager.console.err)
     pager.runJSJobs()
-    pager.loader.unregistered.setLen(0)
     pager.acceptBuffers()
 
 proc dumpBuffers(pager: Pager) =
diff --git a/src/server/buffer.nim b/src/server/buffer.nim
index b4b6f206..96ca918f 100644
--- a/src/server/buffer.nim
+++ b/src/server/buffer.nim
@@ -1895,7 +1895,7 @@ proc handleRead(buffer: Buffer; fd: int): bool =
     assert false
   true
 
-proc handleError(buffer: Buffer; fd: int): bool =
+proc handleError(buffer: Buffer; fd: int; event: TPollfd): bool =
   if fd == buffer.rfd:
     # Connection reset by peer, probably. Close the buffer.
     return false
@@ -1923,20 +1923,22 @@ proc runBuffer(buffer: Buffer) =
   while alive:
     let timeout = buffer.getPollTimeout()
     buffer.pollData.poll(timeout)
+    buffer.loader.blockRegister()
     for event in buffer.pollData.events:
       if (event.revents and POLLIN) != 0:
         if not buffer.handleRead(event.fd):
           alive = false
           break
       if (event.revents and POLLERR) != 0 or (event.revents and POLLHUP) != 0:
-        if not buffer.handleError(event.fd):
+        if not buffer.handleError(event.fd, event):
           alive = false
           break
+    buffer.loader.unregistered.setLen(0)
+    buffer.loader.unblockRegister()
     if buffer.config.scripting != smFalse:
       if buffer.window.timeouts.run(buffer.estream):
         buffer.window.runJSJobs()
         buffer.maybeReshape()
-    buffer.loader.unregistered.setLen(0)
 
 proc cleanup(buffer: Buffer) =
   buffer.pstream.sclose()
diff --git a/src/server/loaderiface.nim b/src/server/loaderiface.nim
index 5eb0f0f6..204d0cc5 100644
--- a/src/server/loaderiface.nim
+++ b/src/server/loaderiface.nim
@@ -34,6 +34,10 @@ type
     sockDir*: string
     # (FreeBSD only) fd for the socket directory so we can connectat() on it
     sockDirFd*: cint
+    # A mechanism to queue up new fds being added to the poll data
+    # inside the events iterator.
+    registerBlocked: bool
+    registerQueue: seq[ConnectData]
 
   ConnectDataState = enum
     cdsBeforeResult, cdsBeforeStatus, cdsBeforeHeaders
@@ -163,11 +167,28 @@ proc unset*(loader: FileLoader; data: MapData) =
   if loader.get(fd) != nil:
     loader.unset(fd)
 
+proc register(loader: FileLoader; data: ConnectData) =
+  if loader.registerBlocked:
+    loader.registerQueue.add(data)
+  else:
+    loader.registerFun(int(data.stream.fd))
+    loader.put(data)
+
+proc blockRegister*(loader: FileLoader) =
+  assert not loader.registerBlocked
+  loader.registerBlocked = true
+
+proc unblockRegister*(loader: FileLoader) =
+  assert loader.registerBlocked
+  loader.registerBlocked = false
+  for it in loader.registerQueue:
+    loader.register(it)
+  loader.registerQueue.setLen(0)
+
 proc fetch0(loader: FileLoader; input: Request; promise: FetchPromise;
     redirectNum: int) =
   let stream = loader.startRequest(input)
-  loader.registerFun(int(stream.fd))
-  loader.put(ConnectData(
+  loader.register(ConnectData(
     promise: promise,
     request: input,
     stream: stream,
@@ -182,13 +203,11 @@ proc fetch*(loader: FileLoader; input: Request): FetchPromise =
 proc reconnect*(loader: FileLoader; data: ConnectData) =
   data.stream.sclose()
   let stream = loader.startRequest(data.request)
-  let data = ConnectData(
+  loader.register(ConnectData(
     promise: data.promise,
     request: data.request,
     stream: stream
-  )
-  loader.put(data)
-  loader.registerFun(data.fd)
+  ))
 
 proc suspend*(loader: FileLoader; fds: seq[int]) =
   let stream = loader.connect()