diff --git a/aws-s3-lwt/io.ml b/aws-s3-lwt/io.ml index c032373..5372f3a 100644 --- a/aws-s3-lwt/io.ml +++ b/aws-s3-lwt/io.ml @@ -199,27 +199,42 @@ module Net = struct Pipe.write input data >>= fun () -> read () in - (* We close input and output when input is closed *) + (* We close input and output when input is closed. A TLS shutdown on a + peer-closed socket can itself raise, and an exception escaping through + Lwt.async here leaves the fd in CLOSE-WAIT for good. *) + let close_quietly ch = + Lwt.catch (fun () -> Lwt_io.close ch) (fun _ -> Lwt.return ()) + in Lwt.async (fun () -> Pipe.closed reader >>= fun () -> - Lwt_io.close oc >>= fun () -> Lwt_io.close ic); + close_quietly oc >>= fun () -> close_quietly ic); Lwt.async read; let output, writer = Pipe.create () in + (* A failed write must fail the whole connection: the read side is closed so + that a pending response read errors out instead of waiting forever for a + peer that will never answer. *) let rec write () = match Queue.take output.Pipe.queue with | Flush waiter -> - Lwt_io.flush oc >>= fun () -> + catch_result (fun () -> Lwt_io.flush oc) >>= fun res -> Lwt.wakeup_later waiter (); - write () + (match res with + | Ok () -> write () + | Error _ -> fail_connection ()) | Data data -> - Lwt_io.write oc data >>= fun () -> - write () + catch_result (fun () -> Lwt_io.write oc data) >>= (function + | Ok () -> write () + | Error _ -> fail_connection ()) | exception Queue.Empty when output.Pipe.closed -> Lwt.return () | exception Queue.Empty -> Lwt_condition.wait output.Pipe.cond >>= fun () -> write () + and fail_connection () = + Pipe.close_reader output; + Pipe.close input; + Lwt.return () in Lwt.async write; Deferred.Or_error.return (reader, writer) diff --git a/aws-s3/http.ml b/aws-s3/http.ml index dac2724..19c4a75 100644 --- a/aws-s3/http.ml +++ b/aws-s3/http.ml @@ -156,13 +156,107 @@ module Make(Io : Types.Io) = struct Or_error.return (code, message, headers, error_body) - let call ?(expect=false) ?connect_timeout_ms ~(endpoint:Region.endpoint) ~path ?(query=[]) ~headers ~sink ?body (meth:meth) = - Net.connect ?connect_timeout_ms ~inet:endpoint.inet ~host:endpoint.host ~port:endpoint.port ~scheme:endpoint.scheme () >>=? fun (reader, writer) -> - (* At this point we need to make sure reader and writer are closed properly. *) - do_request ~expect ~path ~query ~headers ~sink ?body meth reader writer >>= fun result -> - (* Close the reader and writer regardless of status *) + (* Idle connections to the same (scheme, host, port) are kept open and reused, + saving a TCP + TLS handshake per request. Each entry carries the time it + was returned to the pool. *) + let pool : (string, ((string Pipe.reader * string Pipe.writer) * float) Queue.t) Hashtbl.t = + Hashtbl.create 8 + + let max_idle_per_host = 32 + + (* Servers reap idle keep-alive connections on their own schedule (Backblaze + B2 aggressively so), and handing out one the peer has already closed costs + a failed request, so how long an entry may sit here is peer-dependent. *) + let max_idle_age_s = + match Sys.getenv_opt "AWS_S3_POOL_MAX_IDLE_S" with + | Some s -> (try float_of_string s with _ -> 20.) + | None -> 20. + + let pool_key (endpoint : Region.endpoint) = + let scheme = match endpoint.scheme with `Http -> "http" | `Https -> "https" in + sprintf "%s:%s:%d" scheme endpoint.host endpoint.port + + let pool_queue key = + match Hashtbl.find_opt pool key with + | Some q -> q + | None -> let q = Queue.create () in Hashtbl.replace pool key q; q + + (* Closing the reader cascades to the socket close, see [Net.connect]. *) + let discard_conn (reader, writer) = Pipe.close writer; - Pipe.close_reader reader; - Pipe.close sink; - return result + Pipe.close_reader reader + + let rec take_idle key = + match Queue.take_opt (pool_queue key) with + | None -> None + | Some (((reader, writer) as conn), idle_since) -> + let stale = + Pipe.is_closed reader || Pipe.is_closed writer + || Unix.gettimeofday () -. idle_since > max_idle_age_s + in + match stale with + | true -> discard_conn conn; take_idle key + | false -> Some conn + + let return_idle key conn = + let q = pool_queue key in + match Queue.length q >= max_idle_per_host with + | true -> discard_conn conn + | false -> Queue.add (conn, Unix.gettimeofday ()) q + + let get_connection ?connect_timeout_ms (endpoint : Region.endpoint) = + match take_idle (pool_key endpoint) with + | Some conn -> Deferred.Or_error.return (`Reused conn) + | None -> + Net.connect ?connect_timeout_ms ~inet:endpoint.inet ~host:endpoint.host + ~port:endpoint.port ~scheme:endpoint.scheme () + >>=? fun conn -> Deferred.Or_error.return (`Fresh conn) + + (* Reuse requires that the whole response body was consumed, which only holds + when the response was framed by Content-Length or chunked transfer-encoding + (HEAD carries no body); otherwise leftover bytes corrupt the next request. *) + let response_keeps_alive ~meth ~headers = + let connection_close = + match Headers.find_opt "connection" headers with + | Some v -> + String.split_on_char ~sep:',' v + |> List.exists ~f:(fun t -> String.lowercase_ascii (String.trim t) = "close") + | None -> false + in + let framed = + meth = `HEAD + || Headers.find_opt "content-length" headers <> None + || Headers.find_opt "transfer-encoding" headers <> None + in + (not connection_close) && framed + + let call ?(expect=false) ?connect_timeout_ms ~(endpoint:Region.endpoint) ~path ?(query=[]) ~headers ~sink ?body (meth:meth) = + let key = pool_key endpoint in + let run (reader, writer) = + do_request ~expect ~path ~query ~headers ~sink ?body meth reader writer + in + let finish conn result = + (match result with + | Ok (_, _, resp_headers, _) when response_keeps_alive ~meth ~headers:resp_headers -> + return_idle key conn + | _ -> discard_conn conn); + Pipe.close sink; + return result + in + get_connection ?connect_timeout_ms endpoint >>=? fun tagged -> + let conn = match tagged with `Fresh c | `Reused c -> c in + let reused = match tagged with `Reused _ -> true | `Fresh _ -> false in + run conn >>= fun result -> + (* A pooled socket may have been dropped by the peer while idle, so a + failure on a reused connection is retried once on a fresh one -- only + without a body, which the first attempt consumed and this layer cannot + replay (PUT relies on the caller's retry, which can rebuild it). *) + match result with + | Error _ when reused && body = None -> + discard_conn conn; + Net.connect ?connect_timeout_ms ~inet:endpoint.inet ~host:endpoint.host + ~port:endpoint.port ~scheme:endpoint.scheme () >>= (function + | Ok conn -> run conn >>= fun result -> finish conn result + | Error _ as err -> Pipe.close sink; return err) + | _ -> finish conn result end diff --git a/aws-s3/s3.mli b/aws-s3/s3.mli index 051736a..fd5341f 100644 --- a/aws-s3/s3.mli +++ b/aws-s3/s3.mli @@ -7,7 +7,8 @@ it is strongly recommended that you use https. To use https, make sure to have the relevant opam packages installed: [async_ssl] for [async] and [lwt_ssl]/[tls] for [lwt]. - Please note that connections are not reused due to a limitation on the AWS endpoint. + Connections are kept alive and reused across requests to the same + (scheme, host, port) endpoint (HTTP/1.1 persistent connections). If no credentials is provided, the requests will not be signed,