diff --git a/lib/consts.ml b/lib/consts.ml index 715cf99..96414f1 100644 --- a/lib/consts.ml +++ b/lib/consts.ml @@ -21,3 +21,6 @@ let wait_for_xapi_timeout_seconds = 300.0 (** We sleep for this many seconds before the next login attempt. *) let wait_for_xapi_retry_delay_seconds = 4.0 + +(** Allow no more than this many client connections. *) +let connection_limit = 16 diff --git a/src/main.ml b/src/main.ml index 1d99308..b6acdce 100644 --- a/src/main.ml +++ b/src/main.ml @@ -125,6 +125,17 @@ let main port certfile ciphersuites = Lwt_unix.bind sock sockaddr; Lwt_unix.listen sock 5; Lwt_log.notice "Listening for incoming connections" >>= fun () -> + + let conn_count = ref 0 in + let conn_m = Lwt_mutex.create () in + let inc_conn ?(i=1) () = Lwt_mutex.with_lock conn_m (fun () -> + conn_count := !conn_count + i; + if !conn_count > Consts.connection_limit && i > 0 + then Lwt.fail_with ("Server busy: already at maximum "^(string_of_int Consts.connection_limit)^" connections.") + else Lwt.return () + ) in + let dec_conn () = inc_conn ~i:(-1) () in + let rec loop () = Lwt_unix.accept sock >>= fun (fd, _) -> @@ -135,14 +146,18 @@ let main port certfile ciphersuites = (fun () -> Lwt.finalize (fun () -> ( - xapi_says_use_tls () >>= + inc_conn () >>= + xapi_says_use_tls >>= fun tls -> ( let tls_role = if tls then tls_server_role else None in handle_connection fd tls_role) ) ) (* ignore the exception resulting from double-closing the socket *) - (ignore_exn_delayed (fun () -> Lwt_unix.close fd)) + (fun () -> + ignore_exn_delayed (fun () -> Lwt_unix.close fd) () >>= + dec_conn + ) ) in loop ()