Merge pull request #48 from moteus/master

Update multi-uv example.

Author: moteus

examples/cURLv3/multi-uv.lua

@@ -1,12 +1,15 @@
 local curl = require "cURL"
 local uv   = require "lluv"
+local ut   = require "lluv.utils"
 
 local fprintf = function(f, ...) f:write((string.format(...))) end
 local printf  = function(...) fprintf(io.stdout, ...) end
 
 local stderr = io.stderr
 
-local timeout, curl_handle
+local trace = false
+
+trace = trace and print or function() end
 
 local ACTION_NAMES = {
   [curl.POLL_IN     ] = "POLL_IN";
@@ -15,141 +18,240 @@ local ACTION_NAMES = {
   [curl.POLL_NONE   ] = "POLL_NONE";
   [curl.POLL_REMOVE ] = "POLL_REMOVE";
 }
+
 local EVENT_NAMES = {
   [ uv.READABLE ] = "READABLE";
   [ uv.WRITABLE ] = "WRITABLE";
   [ uv.READABLE + uv.WRITABLE ] = "READABLE + WRITABLE";
 }
+
 local FLAGS = {
-  [ uv.READABLE ]               = curl.CSELECT_IN;
-  [ uv.WRITABLE ]               = curl.CSELECT_OUT;
+  [ uv.READABLE               ] = curl.CSELECT_IN;
+  [ uv.WRITABLE               ] = curl.CSELECT_OUT;
   [ uv.READABLE + uv.WRITABLE ] = curl.CSELECT_IN + curl.CSELECT_OUT;
+}
 
+local POLL_IO_FLAGS = {
+  [ curl.POLL_IN    ] = uv.READABLE;
+  [ curl.POLL_OUT   ] = uv.WRITABLE;
+  [ curl.POLL_INOUT ] = uv.READABLE + uv.WRITABLE;
 }
 
-local trace = true
+local Context = ut.class() do
 
-trace = trace and print or function() end
+function Context:__init(fd)
+  self._fd   = assert(fd)
+  self._poll = uv.poll_socket(fd)
+  self._poll.data = {context = self}
+
+  assert(self._poll:fileno() == fd)
+
+  return self
+end
 
-local CONTEXT = {}
+function Context:close()
+  if not self._poll then return end
+  self._poll.data = nil
+  self._poll:close()
+  self._poll, self._fd = nil
+end
+
+function Context:poll(...)
+  self._poll:start(...)
+end
+
+function Context:fileno()
+  return self._fd
+end
 
-function create_curl_context(sockfd)
-  local context = {
-    sockfd      = sockfd;
-    poll_handle = uv.poll_socket(sockfd);
-  }
-  context.poll_handle.data = context
-  
-  return context
 end
 
-function destroy_curl_context(context)
-  context.poll_handle:close()
-end 
+-- Number of parallel request
+local MAX_REQUESTS
+local timer, multi
+local qtask = ut.Queue.new() -- wait tasks
+local qfree = ut.Queue.new() -- avaliable easy handles
+local qeasy = {} -- all easy handles
 
-function add_download(url, num)
+local function on_begin(handle, url, num)
   local filename = tostring(num) ..  ".download"
   local file = io.open(filename, "w")
   if not file then
     fprintf(stderr, "Error opening %s\n", filename)
     return
   end
+  handle.data.file = file
+  handle:setopt_writefunction(file)
 
-  local handle = curl.easy{
-    url           = url;
-    writefunction = file;
-  }
+  fprintf(stderr, "Added download %s -> %s\n", url, filename);
+  return true
+end
 
-  handle.data = file
+local function on_end(handle, err, url)
+  handle.data.file:close()
+  handle.data.file = nil
 
-  curl_handle:add_handle(handle)
-  fprintf(stderr, "Added download %s -> %s\n", url, filename);
+  if err then
+    printf("%s ERROR - %s\n", url, tostring(err));
+  else
+    printf("%s DONE\n", url);
+  end
 end
 
-function check_multi_info()
-  while true do
-    local easy, ok, err = curl_handle:info_read(true)
-    if not easy then curl_handle:close() error(err) end
-    if easy == 0 then break end
+local function cleanup()
+  timer:close()
 
-    local context = CONTEXT[e]
-    if context then destroy_curl_context(context) end
-    local file = assert(easy.data)
-    file:close()
-    local done_url = easy:getinfo_effective_url()
+  for i, easy in ipairs(qeasy) do
+    multi:remove_handle(easy)
     easy:close()
-    if ok then
-      printf("%s DONE\n", done_url);
-    elseif data == "error" then
-      printf("%s ERROR - %s\n", done_url, tostring(err));
-    end
   end
+
+  multi:close()
 end
 
-function curl_perform(handle, err, events)
-  -- calls by libuv --
-  trace("UV::POLL", handle, err, EVENT_NAMES[events] or events)
+local proceed_queue, add_download do
 
-  local flags = assert(FLAGS[events], ("unknown event:" .. events))
+proceed_queue = function()
+  while true do
+    if qtask:empty() then return end
+
+    if qfree:empty() then
+      if #qeasy < MAX_REQUESTS then
+        local easy = assert(curl.easy())
+        qeasy[#qeasy + 1] = easy
+        qfree:push(easy)
+      else
+        return
+      end
+    end
 
-  context = handle.data
+    local task = assert(qtask:pop())
+    local url, num = task[1], task[2]
 
-  curl_handle:socket_action(context.sockfd, flags)
+    local handle = assert(qfree:pop())
 
-  check_multi_info()
+    handle:setopt{
+      url           = url;
+      fresh_connect = true;
+      forbid_reuse  = true;
+    }
+
+    handle.data = {}
+
+    if on_begin(handle, url, num) then
+      multi:add_handle(handle)
+    else
+      handle:reset().data = nil
+      qfree:push(handle)
+    end
+  end
 end
 
-function on_timeout(timer)
-  -- calls by libuv --
-  trace("UV::TIMEOUT", timer)
+add_download = function(url, num)
+  qtask:push{url, num}
 
-  local running_handles, err = curl_handle:socket_action()
+  proceed_queue()
+end
 
-  check_multi_info()
 end
 
-function start_timeout(timeout_ms)
+local on_libuv_poll, on_libuv_timeout
+
+local on_curl_timeout, on_curl_action do
+
+on_curl_timeout = function(ms)
   -- calls by curl --
-  trace("CURL::TIMEOUT", timeout_ms)
+  trace("CURL::TIMEOUT", ms)
 
-  -- 0 means directly call socket_action, but we'll do it in a bit
-  if timeout_ms <= 0 then timeout_ms = 1 end
+  if ms <= 0 then ms = 1 end
 
-  timeout:stop():start(timeout_ms, 0, on_timeout)
+  timer:start(ms, 0, on_libuv_timeout)
 end
 
-function handle_socket(easy, s, action)
+on_curl_action = function(easy, fd, action)
   local ok, err = pcall(function()
-    -- calls by curl --
     trace("CURL::SOCKET", easy, s, ACTION_NAMES[action] or action)
 
-    local curl_context = CONTEXT[easy] or create_curl_context(s)
-    CONTEXT[easy] = curl_context
-
-    assert(curl_context.sockfd == s)
+    local context = easy.data.context
 
-    if action == curl.POLL_IN then
-      curl_context.poll_handle:start(uv.READABLE, curl_perform)
-    elseif action == curl.POLL_OUT then
-      curl_context.poll_handle:start(uv.WRITABLE, curl_perform)
+    local flag = POLL_IO_FLAGS[action]
+    if flag then
+      if not context then
+        context = Context.new(fd)
+        easy.data.context = context
+      end
+      context:poll(flag, on_libuv_poll)
     elseif action == curl.POLL_REMOVE then
-      CONTEXT[easy] = nil
-      destroy_curl_context(curl_context)
+      if context then
+        easy.data.context = nil
+        context:close()
+      end
     end
   end)
+
   if not ok then uv.defer(function() error(err) end) end
 end
 
-timeout = uv.timer()
+end
 
-curl_handle = curl.multi{
-  socketfunction = handle_socket;
-  timerfunction  = start_timeout;
-}
+-- on_libuv_poll, on_libuv_timeout 
+do
+
+local curl_check_multi_info = function()
+  while true do
+    local easy, ok, err = multi:info_read(true)
+
+    if not easy then
+      multi:close()
+      error(err)
+    end
+
+    if easy == 0 then break end
 
-curl_handle = curl.multi{
-  socketfunction = handle_socket;
-  timerfunction  = start_timeout;
+    local done_url = easy:getinfo_effective_url()
+
+    local context = easy.data.context
+    if context then context:close() end
+    easy.data.context = nil
+
+    if ok then on_end(easy, nil, done_url) else on_end(easy, err, done_url) end
+
+    easy:reset().data = nil
+    qfree:push(easy)
+  end
+
+  proceed_queue()
+end
+
+on_libuv_poll = function(handle, err, events)
+  trace("UV::POLL", handle, err, EVENT_NAMES[events] or events)
+
+  local flags = assert(FLAGS[events], ("unknown event:" .. events))
+
+  context = handle.data.context
+
+  multi:socket_action(context:fileno(), flags)
+
+  curl_check_multi_info()
+end
+
+on_libuv_timeout = function(timer)
+  trace("UV::TIMEOUT", timer)
+
+  local running_handles, err = multi:socket_action()
+
+  curl_check_multi_info()
+end
+
+end
+
+MAX_REQUESTS = 64
+
+timer = uv.timer()
+
+multi = curl.multi{
+  timerfunction = on_curl_timeout;
+  socketfunction = on_curl_action;
 }
 
 for i = 1, math.huge do
@@ -158,4 +260,6 @@ for i = 1, math.huge do
   add_download(url, i)
 end
 
-uv.run(loop, UV_RUN_DEFAULT)
+uv.run()
+
+cleanup()