cleanup for lwt

Author: c-cube

src/lwt/tiny_httpd_lwt.ml

@@ -129,6 +129,8 @@ let io_backend ?addr ?port ?unix_sock ?max_connections ?max_buf_pool_size
               Lwt_unix.bind sock sockaddr |> Lwt_direct.await;
               Lwt_unix.listen sock backlog;
 
+              let cleanup () = Lwt_unix.close sock |> Lwt_direct.await in
+
               (* recover real port, if any *)
               (match Unix.getsockname (Lwt_unix.unix_file_descr sock) with
               | Unix.ADDR_INET (_, p) -> port := p
@@ -137,12 +139,6 @@ let io_backend ?addr ?port ?unix_sock ?max_connections ?max_buf_pool_size
               let handle_client client_addr fd : unit =
                 Atomic.incr active_conns;
                 Lwt_direct.run_in_the_background @@ fun () ->
-                let cleanup () =
-                  Log.debug (fun k ->
-                      k "Tiny_httpd_lwt: client handler returned");
-                  Atomic.decr active_conns
-                in
-
                 let buf_ic = Bytes.create buf_size in
                 let buf_oc = Bytes.create buf_size in
                 (*
@@ -154,6 +150,16 @@ let io_backend ?addr ?port ?unix_sock ?max_connections ?max_buf_pool_size
                 let num_open = ref 2 in
                 let ic = ic_of_fd ~num_open ~bytes:buf_ic fd in
                 let oc = oc_of_fd ~num_open ~bytes:buf_oc fd in
+
+                let cleanup () =
+                  Log.debug (fun k ->
+                      k "Tiny_httpd_lwt: client handler returned");
+                  Atomic.decr active_conns;
+                  (try Lwt_unix.shutdown fd SHUTDOWN_ALL with _ -> ());
+                  ic#close ();
+                  oc#close ()
+                in
+
                 try
                   handle.handle ~client_addr ic oc;
                   cleanup ()
@@ -167,10 +173,16 @@ let io_backend ?addr ?port ?unix_sock ?max_connections ?max_buf_pool_size
                         (Printexc.raw_backtrace_to_string bt))
               in
 
-              while Atomic.get running do
-                let fd, addr = Lwt_unix.accept sock |> Lwt_direct.await in
-                handle_client addr fd
-              done
+              try
+                while Atomic.get running do
+                  let fd, addr = Lwt_unix.accept sock |> Lwt_direct.await in
+                  handle_client addr fd
+                done;
+                cleanup ()
+              with exn ->
+                let bt = Printexc.get_raw_backtrace () in
+                cleanup ();
+                Printexc.raise_with_backtrace exn bt
             in
 
             let tcp_server : IO.TCP_server.t =