curl_lwt.ml

   1 (** Lwt support for Curl *)
   2
   3 module M = Curl.Multi
   4
   5 let debug = ref false
   6 let log fmt = Printf.ksprintf prerr_endline fmt
   7
   8 let set_debug x = debug := x
   9
  10 let int_of_fd : Unix.file_descr -> int = Obj.magic
  11
  12 type multi = {
  13   mt : Curl.Multi.mt;
  14   all_events : (Unix.file_descr, Lwt_engine.event list) Hashtbl.t;
  15   wakeners : (Curl.t, Curl.curlCode Lwt.u) Hashtbl.t;
  16 }
  17
  18 let create () =
  19   let mt = M.create () in
  20   let timer_event = ref Lwt_engine.fake_event in
  21   let all_events = Hashtbl.create 32 in
  22   let wakeners = Hashtbl.create 32 in
  23   let finished s =
  24     if !debug then log "finished %s" s;
  25     let rec loop n =
  26       match M.remove_finished mt with
  27       | None -> if n > 0 && !debug then log "removed %u handles via %s" n s
  28       | Some (h,code) ->
  29         if !debug then log "wakeup";
  30         begin try
  31           let w = Hashtbl.find wakeners h in
  32           Hashtbl.remove wakeners h;
  33           Lwt.wakeup w code
  34         with Not_found ->
  35           prerr_endline "curl_lwt: orphan handle, how come?"
  36         end;
  37         loop (n+1)
  38     in
  39     loop 0
  40   in
  41   let on_readable fd _ =
  42     if !debug then log "on_readable fd %d" (int_of_fd fd);
  43     let (_:int) = M.action mt fd M.EV_IN in
  44     finished "on_readable";
  45   in
  46   let on_writable fd _ =
  47     if !debug then log "on_writable fd %d" (int_of_fd fd);
  48     let (_:int) = M.action mt fd M.EV_OUT in
  49     finished "on_writable";
  50   in
  51   let on_timer _ =
  52     if !debug then log "on_timer";
  53     Lwt_engine.stop_event !timer_event;
  54     M.action_timeout mt;
  55     finished "on_timer"
  56   in
  57   M.set_timer_function mt begin fun timeout ->
  58     if !debug then log "set timeout %d" timeout;
  59     Lwt_engine.stop_event !timer_event; (* duplicate stop_event is ok *)
  60     timer_event := Lwt_engine.on_timer (float_of_int timeout /. 1000.) false on_timer
  61   end;
  62   M.set_socket_function mt begin fun fd what ->
  63     if !debug then log "set socket fd %d %s" (int_of_fd fd)
  64       (match what with
  65       | M.POLL_NONE -> "none"
  66       | M.POLL_REMOVE -> "remove"
  67       | M.POLL_IN -> "in"
  68       | M.POLL_OUT -> "out"
  69       | M.POLL_INOUT -> "inout");
  70     begin
  71     try
  72       List.iter Lwt_engine.stop_event (Hashtbl.find all_events fd);
  73       Hashtbl.remove all_events fd;
  74       if !debug then log "removed handlers for %d" (int_of_fd fd);
  75     with
  76       Not_found -> () (* first event for the socket - no association *)
  77     end;
  78     let events = match what with
  79     | M.POLL_REMOVE | M.POLL_NONE -> []
  80     | M.POLL_IN -> [Lwt_engine.on_readable fd (on_readable fd)]
  81     | M.POLL_OUT -> [Lwt_engine.on_writable fd (on_writable fd)]
  82     | M.POLL_INOUT -> [Lwt_engine.on_readable fd (on_readable fd); Lwt_engine.on_writable fd (on_writable fd)]
  83     in
  84     match events with
  85     | [] -> ()
  86     | _ -> Hashtbl.add all_events fd events;
  87   end;
  88   { mt; all_events; wakeners; }
  89
  90 (* lwt may not run in parallel so one global is OK'ish *)
  91 let global = lazy (create ())
  92
  93 let setopt opt =
  94   let t = Lazy.force global in
  95   M.setopt t.mt opt
  96
  97 let perform h =
  98   let t = Lazy.force global in
  99   let (waiter,wakener) = Lwt.wait () in
 100   let waiter = Lwt.protected waiter in
 101   Lwt.on_cancel waiter (fun () ->
 102     Curl.Multi.remove t.mt h;
 103     Hashtbl.remove t.wakeners h;
 104   );
 105   Hashtbl.add t.wakeners h wakener;
 106   M.add t.mt h;
 107   waiter