Loader: use curl_multi

Note: for now it's only used for http requests.
The doRequest API still needs an async rework.

2 files changed, 195 insertions, 108 deletions

diff --git a/src/io/http.nim b/src/io/http.nim
index 36cd6473..a6a1781d 100644
--- a/src/io/http.nim
+++ b/src/io/http.nim
@@ -9,15 +9,31 @@ import types/url
 import utils/twtstr
 
 type
-  HeaderOpaque* = ref object
+  HandleData* = ref HandleDataObj
+  HandleDataObj = object
+    curl*: CURL
     statusline: bool
     headers: HeaderList
-    curl: CURL
     request: Request
-    ostream: Stream
-
-func newHeaderOpaque(curl: CURL, request: Request, ostream: Stream): HeaderOpaque =
-  HeaderOpaque(headers: newHeaderList(), curl: curl, ostream: ostream, request: request)
+    ostream*: Stream
+    mime: curl_mime
+    slist: curl_slist
+
+func newHandleData(curl: CURL, request: Request, ostream: Stream): HandleData =
+  let handleData = HandleData(
+    headers: newHeaderList(),
+    curl: curl,
+    ostream: ostream,
+    request: request
+  )
+  return handleData
+
+proc cleanup*(handleData: HandleData) =
+  if handleData.mime != nil:
+    curl_mime_free(handleData.mime)
+  if handleData.slist != nil:
+    curl_slist_free_all(handleData.slist)
+  curl_easy_cleanup(handleData.curl)
 
 template setopt(curl: CURL, opt: CURLoption, arg: typed) =
   discard curl_easy_setopt(curl, opt, arg)
@@ -33,7 +49,7 @@ proc curlWriteHeader(p: cstring, size: csize_t, nitems: csize_t, userdata: point
   for i in 0..<nitems:
     line[i] = p[i]
 
-  let op = cast[HeaderOpaque](userdata)
+  let op = cast[HandleData](userdata)
   if not op.statusline:
     op.statusline = true
     op.ostream.swrite(int(CURLE_OK))
@@ -54,80 +70,72 @@ proc curlWriteHeader(p: cstring, size: csize_t, nitems: csize_t, userdata: point
   return nitems
 
 proc curlWriteBody(p: cstring, size: csize_t, nmemb: csize_t, userdata: pointer): csize_t {.cdecl.} =
-  let stream = cast[Stream](userdata)
+  let handleData = cast[HandleData](userdata)
   if nmemb > 0:
-    stream.writeData(p, int(nmemb))
-    stream.flush()
+    handleData.ostream.writeData(p, int(nmemb))
+    handleData.ostream.flush()
   return nmemb
 
-proc loadHttp*(request: Request, ostream: Stream) =
+proc applyPostBody(curl: CURL, request: Request, handleData: HandleData) =
+  if request.multipart.issome:
+    handleData.mime = curl_mime_init(curl)
+    if handleData.mime == nil:
+      # fail (TODO: raise?)
+      handleData.ostream.swrite(-1)
+      handleData.ostream.flush()
+      return
+    for entry in request.multipart.get.content:
+      let part = curl_mime_addpart(handleData.mime)
+      if part == nil:
+        # fail (TODO: raise?)
+        handleData.ostream.swrite(-1)
+        handleData.ostream.flush()
+        return
+      curl_mime_name(part, cstring(entry.name))
+      if entry.isFile:
+        if entry.isStream:
+          curl_mime_filedata(part, cstring(entry.filename))
+        else:
+          let fd = readFile(entry.filename)
+          curl_mime_data(part, cstring(fd), csize_t(fd.len))
+        # may be overridden by curl_mime_filedata, so set it here
+        curl_mime_filename(part, cstring(entry.filename))
+      else:
+        curl_mime_data(part, cstring(entry.content), csize_t(entry.content.len))
+    curl.setopt(CURLOPT_MIMEPOST, handleData.mime)
+  elif request.body.issome:
+    curl.setopt(CURLOPT_POSTFIELDS, cstring(request.body.get))
+    curl.setopt(CURLOPT_POSTFIELDSIZE, request.body.get.len)
+
+proc loadHttp*(curlm: CURLM, request: Request, ostream: Stream): HandleData =
   let curl = curl_easy_init()
-
   if curl == nil:
     ostream.swrite(-1)
     ostream.flush()
     return # fail
-
   let surl = request.url.serialize()
   curl.setopt(CURLOPT_URL, surl)
-
-  curl.setopt(CURLOPT_WRITEDATA, ostream)
+  let handleData = curl.newHandleData(request, ostream)
+  curl.setopt(CURLOPT_WRITEDATA, handleData)
   curl.setopt(CURLOPT_WRITEFUNCTION, curlWriteBody)
-
-  let headerres = curl.newHeaderOpaque(request, ostream)
-
-  GC_ref(headerres) # this could get unref'd before writeheader finishes
-  GC_ref(ostream) #TODO not sure about this one, but better safe than sorry
-  defer:
-    GC_unref(headerres)
-    GC_unref(ostream)
-
-  curl.setopt(CURLOPT_HEADERDATA, headerres)
+  curl.setopt(CURLOPT_HEADERDATA, handleData)
   curl.setopt(CURLOPT_HEADERFUNCTION, curlWriteHeader)
-
-  var mime: curl_mime = nil
-
   case request.httpmethod
-  of HTTP_GET: curl.setopt(CURLOPT_HTTPGET, 1)
+  of HTTP_GET:
+    curl.setopt(CURLOPT_HTTPGET, 1)
   of HTTP_POST:
     curl.setopt(CURLOPT_POST, 1)
-    if request.multipart.issome:
-      mime = curl_mime_init(curl)
-      if mime == nil: return # fail
-      for entry in request.multipart.get.content:
-        let part = curl_mime_addpart(mime)
-        if part == nil: return # fail
-        curl_mime_name(part, cstring(entry.name))
-        if entry.isFile:
-          if entry.isStream:
-            curl_mime_filedata(part, cstring(entry.filename))
-          else:
-            let fd = readFile(entry.filename)
-            curl_mime_data(part, cstring(fd), csize_t(fd.len))
-          # may be overridden by curl_mime_filedata, so set it here
-          curl_mime_filename(part, cstring(entry.filename))
-        else:
-          curl_mime_data(part, cstring(entry.content), csize_t(entry.content.len))
-      curl.setopt(CURLOPT_MIMEPOST, mime)
-    elif request.body.issome:
-      curl.setopt(CURLOPT_POSTFIELDS, cstring(request.body.get))
-      curl.setopt(CURLOPT_POSTFIELDSIZE, request.body.get.len)
+    curl.applyPostBody(request, handleData)
   else: discard #TODO
-
-  var slist: curl_slist = nil
   for k, v in request.headers:
     let header = k & ": " & v
-    slist = curl_slist_append(slist, cstring(header))
-  if slist != nil:
-    curl.setopt(CURLOPT_HTTPHEADER, slist)
-
-  let res = curl_easy_perform(curl)
-  if res != CURLE_OK:
+    handleData.slist = curl_slist_append(handleData.slist, cstring(header))
+  if handleData.slist != nil:
+    curl.setopt(CURLOPT_HTTPHEADER, handleData.slist)
+  let res = curl_multi_add_handle(curlm, curl)
+  if res != CURLM_OK:
     ostream.swrite(int(res))
     ostream.flush()
-
-  curl_easy_cleanup(curl)
-  if mime != nil:
-    curl_mime_free(mime)
-  if slist != nil:
-    curl_slist_free_all(slist)
+    #TODO: raise here?
+    return
+  return handleData
diff --git a/src/io/loader.nim b/src/io/loader.nim
index d96d1a5a..bced0298 100644
--- a/src/io/loader.nim
+++ b/src/io/loader.nim
@@ -23,6 +23,7 @@ import bindings/curl
 import io/about
 import io/file
 import io/http
+import io/promise
 import io/request
 import io/urlfilter
 import ips/serialize
@@ -41,29 +42,114 @@ type
   LoaderCommand = enum
     LOAD, QUIT
 
+  LoaderContext = ref object
+    ssock: ServerSocket
+    alive: bool
+    curlm: CURLM
+    config: LoaderConfig
+    extra_fds: seq[curl_waitfd]
+    handleList: seq[HandleData]
+
   LoaderConfig* = object
     defaultheaders*: HeaderList
     filter*: URLFilter
     cookiejar*: CookieJar
     referrerpolicy*: ReferrerPolicy
 
-proc loadResource(request: Request, ostream: Stream) =
+proc addFd(ctx: LoaderContext, fd: int, flags: int) =
+  ctx.extra_fds.add(curl_waitfd(
+    fd: cast[cint](fd),
+    events: cast[cshort](flags)
+  ))
+
+proc loadResource(ctx: LoaderContext, request: Request, ostream: Stream) =
   case request.url.scheme
   of "file":
     loadFile(request.url, ostream)
+    ostream.close()
   of "http", "https":
-    loadHttp(request, ostream)
+    let handleData = loadHttp(ctx.curlm, request, ostream)
+    if handleData != nil:
+      ctx.handleList.add(handleData)
   of "about":
     loadAbout(request, ostream)
+    ostream.close()
   else:
     ostream.swrite(-1) # error
-    ostream.flush()
+    ostream.close()
 
-var ssock: ServerSocket
-proc runFileLoader*(fd: cint, config: LoaderConfig) =
+proc onLoad(ctx: LoaderContext, stream: Stream) =
+  var request: Request
+  stream.sread(request)
+  if not ctx.config.filter.match(request.url):
+    stream.swrite(-1) # error
+    stream.flush()
+  else:
+    for k, v in ctx.config.defaultHeaders.table:
+      if k notin request.headers.table:
+        request.headers.table[k] = v
+    if ctx.config.cookiejar != nil and ctx.config.cookiejar.cookies.len > 0:
+      if "Cookie" notin request.headers.table:
+        let cookie = ctx.config.cookiejar.serialize(request.url)
+        if cookie != "":
+          request.headers["Cookie"] = cookie
+    if request.referer != nil and "Referer" notin request.headers.table:
+      let r = getReferer(request.referer, request.url, ctx.config.referrerpolicy)
+      if r != "":
+        request.headers["Referer"] = r
+    ctx.loadResource(request, stream)
+
+proc acceptConnection(ctx: LoaderContext) =
+  #TODO TODO TODO acceptSocketStream should be non-blocking here,
+  # otherwise the client disconnecting between poll and accept could
+  # block this indefinitely.
+  let stream = ctx.ssock.acceptSocketStream()
+  try:
+    var cmd: LoaderCommand
+    stream.sread(cmd)
+    case cmd
+    of LOAD:
+      ctx.onLoad(stream)
+    of QUIT:
+      ctx.alive = false
+      stream.close()
+  except IOError:
+    # End-of-file, broken pipe, or something else. For now we just
+    # ignore it and pray nothing breaks.
+    # (TODO: this is probably not a very good idea.)
+    stream.close()
+
+proc finishCurlTransfer(ctx: LoaderContext, handleData: HandleData, res: int) =
+  if res != int(CURLE_OK):
+    handleData.ostream.swrite(int(res))
+    handleData.ostream.flush()
+  discard curl_multi_remove_handle(ctx.curlm, handleData.curl)
+  handleData.ostream.close()
+  handleData.cleanup()
+
+proc exitLoader(ctx: LoaderContext) =
+  for handleData in ctx.handleList:
+    #TODO: -1, -2, -3, ... results should be named.
+    ctx.finishCurlTransfer(handleData, -3)
+  discard curl_multi_cleanup(ctx.curlm)
+  curl_global_cleanup()
+  ctx.ssock.close()
+  quit(0)
+
+var gctx: LoaderContext
+proc initLoaderContext(fd: cint, config: LoaderConfig): LoaderContext =
   if curl_global_init(CURL_GLOBAL_ALL) != CURLE_OK:
     raise newException(Defect, "Failed to initialize libcurl.")
-  ssock = initServerSocket()
+  let curlm = curl_multi_init()
+  if curlm == nil:
+    raise newException(Defect, "Failed to initialize multi handle.")
+  var ctx = LoaderContext(
+    alive: true,
+    curlm: curlm,
+    config: config
+  )
+  gctx = ctx
+  ctx.ssock = initServerSocket()
   # The server has been initialized, so the main process can resume execution.
   var writef: File
   if not open(writef, FileHandle(fd), fmWrite):
@@ -73,45 +159,38 @@ proc runFileLoader*(fd: cint, config: LoaderConfig) =
   close(writef)
   discard close(fd)
   onSignal SIGTERM, SIGINT:
-    curl_global_cleanup()
-    ssock.close()
-    quit(1)
-  while true:
-    let stream = ssock.acceptSocketStream()
-    try:
-      var cmd: LoaderCommand
-      stream.sread(cmd)
-      case cmd
-      of LOAD:
-        var request: Request
-        stream.sread(request)
-        if not config.filter.match(request.url):
-          stream.swrite(-1) # error
-          stream.flush()
-        else:
-          for k, v in config.defaultHeaders.table:
-            if k notin request.headers.table:
-              request.headers.table[k] = v
-          if config.cookiejar != nil and config.cookiejar.cookies.len > 0:
-            if "Cookie" notin request.headers.table:
-              let cookie = config.cookiejar.serialize(request.url)
-              if cookie != "":
-                request.headers["Cookie"] = cookie
-          if request.referer != nil and "Referer" notin request.headers.table:
-            let r = getReferer(request.referer, request.url, config.referrerpolicy)
-            if r != "":
-              request.headers["Referer"] = r
-          loadResource(request, stream)
-        stream.close()
-      of QUIT:
-        stream.close()
+    gctx.exitLoader()
+  ctx.addFd(int(ctx.ssock.sock.getFd()), CURL_WAIT_POLLIN)
+  return ctx
+
+proc runFileLoader*(fd: cint, config: LoaderConfig) =
+  var ctx = initLoaderContext(fd, config)
+  while ctx.alive:
+    var numfds: cint = 0
+    #TODO do not discard
+    discard curl_multi_poll(ctx.curlm, addr ctx.extra_fds[0],
+      cuint(ctx.extra_fds.len), 30_000, addr numfds)
+    discard curl_multi_perform(ctx.curlm, addr numfds)
+    for extra_fd in ctx.extra_fds.mitems:
+      # For now, this is always ssock.sock.getFd().
+      if extra_fd.events == extra_fd.revents:
+        ctx.acceptConnection()
+        extra_fd.revents = 0
+    var msgs_left: cint = 1
+    while msgs_left > 0:
+      let msg = curl_multi_info_read(ctx.curlm, addr msgs_left)
+      if msg == nil:
         break
-    except IOError:
-      # End-of-file, broken pipe, or something.
-      stream.close()
-  curl_global_cleanup()
-  ssock.close()
-  quit(0)
+      if msg.msg == CURLMSG_DONE: # the only possible value atm
+        var idx = -1
+        for i in 0 ..< ctx.handleList.len:
+          if ctx.handleList[i].curl == msg.easy_handle:
+            idx = i
+            break
+        assert idx != -1
+        ctx.finishCurlTransfer(ctx.handleList[idx], int(msg.data.result))
+        ctx.handleList.del(idx)
+  ctx.exitLoader()
 
 #TODO async requests...
 proc doRequest*(loader: FileLoader, request: Request, blocking = true): Response =