Add stop parameters in main-loop of server

dinosaure · dinosaure · commit 4c5cd893d37d · 2014-05-23T11:08:34.000+02:00
diff --git a/lib/lwt_unix_conduit.ml b/lib/lwt_unix_conduit.ml
@@ -41,15 +41,16 @@ ELSE
   | `TCP -> LUN.Tcp_client.connect sa
 END
 
-let serve ~mode ~sockaddr ?timeout callback =
+let serve ~mode ~sockaddr ?stop ?timeout callback =
 IFDEF HAVE_LWT_SSL THEN
   match mode with
-  | `TCP -> LUN.Tcp_server.init ~sockaddr ?timeout callback
+  | `TCP -> LUN.Tcp_server.init ~sockaddr ?stop ?timeout callback
   | `SSL (`Crt_file_path certfile, `Key_file_path keyfile) -> 
-    Lwt_unix_net_ssl.Server.init ~certfile ~keyfile ?timeout sockaddr callback
+    Lwt_unix_net_ssl.Server.init ~certfile ~keyfile
+      ?stop ?timeout sockaddr callback
 ELSE
   match mode with
-  | `TCP -> LUN.Tcp_server.init ~sockaddr ?timeout callback
+  | `TCP -> LUN.Tcp_server.init ~sockaddr ?stop ?timeout callback
   | `SSL (`Crt_file_path certfile, `Key_file_path keyfile) -> 
       fail (Failure "No SSL support compiled into Conduit")
 END
diff --git a/lib/lwt_unix_conduit.mli b/lib/lwt_unix_conduit.mli
@@ -31,6 +31,7 @@ val connect :
 val serve :
   mode:server_mode ->
   sockaddr:Lwt_unix.sockaddr ->
+  ?stop:(unit -> bool) ->
   ?timeout:int -> (ic -> oc -> unit io) -> unit io
 
 val close_in : 'a Lwt_io.channel -> unit
diff --git a/lib/lwt_unix_net.ml b/lib/lwt_unix_net.ml
@@ -58,17 +58,17 @@ module Tcp_server = struct
     Lwt_unix.setsockopt client Lwt_unix.TCP_NODELAY true;
     let ic = Lwt_io.of_fd ~mode:Lwt_io.input client in
     let oc = Lwt_io.of_fd ~mode:Lwt_io.output client in
- 
+
     let c = callback ic oc in
     let events = match timeout with
       |None -> [c]
       |Some t -> [c; (Lwt_unix.sleep (float_of_int t)) ] in
     let _ = Lwt.pick events >>= fun () -> close (ic,oc) in
     return ()
-  
-  let init ~sockaddr ?timeout callback =
+
+  let init ~sockaddr ?(stop = (fun () -> true)) ?timeout callback =
     let s = init_socket sockaddr in
-    while_lwt true do
+    while_lwt (stop ()) do
       Lwt_unix.accept s >>=
       process_accept ?timeout callback
     done
diff --git a/lib/lwt_unix_net.mli b/lib/lwt_unix_net.mli
@@ -32,6 +32,7 @@ module Tcp_server : sig
 
   val init :
     sockaddr:Lwt_unix.sockaddr ->
+    ?stop:(unit -> bool) ->
     ?timeout:int ->
     (input channel -> output channel -> unit Lwt.t) ->
     unit Lwt.t
diff --git a/lib/lwt_unix_net_ssl.ml b/lib/lwt_unix_net_ssl.ml
@@ -67,13 +67,14 @@ module Server = struct
     let _ = Lwt.pick events >>= fun () -> close (ic,oc) in
     return ()
 
-  let init ?(nconn=20) ?password ~certfile ~keyfile ?timeout sa callback =
+  let init ?(nconn=20) ?password ~certfile ~keyfile
+    ?(stop = (fun () -> true)) ?timeout sa callback =
     let s = listen ~nconn ?password ~certfile ~keyfile sa in
     let cont = ref true in
-    while_lwt !cont do
+    while_lwt !cont && (stop ()) do
       try_lwt begin
         accept s >>= process_accept ~timeout callback
-      end with 
+      end with
        | Lwt.Canceled -> cont := false; return ()
        | _ -> return ()
     done
diff --git a/lib/lwt_unix_net_ssl.mli b/lib/lwt_unix_net_ssl.mli
@@ -38,6 +38,7 @@ module Server : sig
     ?password:(bool -> string) ->
     certfile:string ->
     keyfile:string ->
+    ?stop:(unit -> bool) ->
     ?timeout:int ->
     Lwt_unix.sockaddr ->
     (Lwt_io.input_channel -> Lwt_io.output_channel -> unit Lwt.t) ->
