1(*
2 * Copyright (C) 2014 Zheng Li <dev@zheng.li>
3 *
4 * This program is free software; you can redistribute it and/or modify
5 * it under the terms of the GNU Lesser General Public License as published
6 * by the Free Software Foundation; version 2.1 only. with the special
7 * exception on linking described in file LICENSE.
8 *
9 * This program is distributed in the hope that it will be useful,
10 * but WITHOUT ANY WARRANTY; without even the implied warranty of
11 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12 * GNU Lesser General Public License for more details.
13 *)
14
15
16(* The [read], [write], [except] are fields mapped to the POLLIN/OUT/PRI
17   subscription flags used by poll, which have a correspondence to the
18   readfds, writefds, exceptfds concept as in select. *)
19type event = {
20	mutable read: bool;
21	mutable write: bool;
22	mutable except: bool;
23}
24
25external select_on_poll: (Unix.file_descr * event) array -> int -> int = "stub_select_on_poll"
26external set_fd_limit: int -> unit = "stub_set_fd_limit"
27
28(* The rlim_max given to setrlimit must not go above the system level nr_open,
29   which we can read from /proc/sys. *)
30let get_sys_fs_nr_open () =
31	try
32		let ch = open_in "/proc/sys/fs/nr_open" in
33		let v = int_of_string (input_line ch) in
34		close_in_noerr ch; v
35	with _ -> 1024 * 1024
36
37let init_event () = {read = false; write = false; except = false}
38
39let poll_select in_fds out_fds exc_fds timeout =
40	let h = Hashtbl.create 57 in
41	let add_event event_set fd =
42		let e =
43			try Hashtbl.find h fd
44			with Not_found ->
45				let e = init_event () in
46				Hashtbl.add h fd e; e in
47		event_set e in
48	List.iter (add_event (fun x -> x.read <- true)) in_fds;
49	List.iter (add_event (fun x -> x.write <- true)) out_fds;
50	List.iter (add_event (fun x -> x.except <- true)) exc_fds;
51	(* Unix.stdin and init_event are dummy input as stubs, which will
52           always be overwritten later on.  *)
53	let a = Array.make (Hashtbl.length h) (Unix.stdin, init_event ()) in
54	let i = ref (-1) in
55	Hashtbl.iter (fun fd event -> incr i; Array.set a !i (fd, event)) h;
56	let n = select_on_poll a (int_of_float (timeout *. 1000.)) in
57	let r = [], [], [] in
58	if n = 0 then r else
59		Array.fold_right
60			(fun (fd, event) (r, w, x) ->
61			 (if event.read then fd :: r else r),
62			 (if event.write then fd :: w else w),
63			 (if event.except then fd :: x else x))
64			a r
65
66(* If the use_poll function is not called at all, we default to the original Unix.select behavior *)
67let select_fun = ref Unix.select
68
69let use_poll yes =
70	let sel_fun, max_fd =
71		if yes then poll_select, get_sys_fs_nr_open ()
72		else Unix.select, 1024 in
73	select_fun := sel_fun;
74	set_fd_limit max_fd
75
76let select in_fds out_fds exc_fds timeout =
77	(!select_fun) in_fds out_fds exc_fds timeout
78