echo.ml

[@@@alert "-unstable"] (* my fucking god *)
open Oslo4

let port = 18624

let print_mtx = Mutex.create ()

type stats = {
  conn_total : int Atomic.t;
  conn_active : int Atomic.t;
  conn_per_domain : int Array.t;
}

let stats = { conn_total = Atomic.make 0;
              conn_active = Atomic.make 0;
              conn_per_domain = Array.make 64 0 }

let print_stats () =
  Mutex.lock print_mtx;
  for i = 0 to Odomain.num_domains () - 1 do
    Printf.printf "\tdomain %02d has %4d connections\n%!"
      i (Array.get stats.conn_per_domain i);
  done;
  Printf.printf "active: %d\ttotal: %d\n%!"
    (Atomic.get stats.conn_active)
    (Atomic.get stats.conn_total);
  Mutex.unlock print_mtx

let add_conn () =
  let num_conn = Array.get stats.conn_per_domain (Odomain.self ()) in
  Array.set stats.conn_per_domain (Odomain.self ()) (succ num_conn);
  Atomic.incr stats.conn_total;
  Atomic.incr stats.conn_active

let sub_conn () =
  let num_conn = Array.get stats.conn_per_domain (Odomain.self ()) in
  Array.set stats.conn_per_domain (Odomain.self ()) (pred num_conn);
  Atomic.decr stats.conn_active

let string_of_sockaddr = function
  | Unix.ADDR_INET (iaddr, port) ->
    Printf.sprintf "%s:%d" (Unix.string_of_inet_addr iaddr) port
  | Unix.ADDR_UNIX s -> s

let writen io buf pos len =
  let rec loop wn =
    if wn = len then
      wn
    else
      loop (wn + Io.single_write io buf (pos + wn) (len - wn))
  in
  loop 0

let handle_client fd _saddr () =
  let io = Io.of_unix fd in
  let buf_len = 1024 * 64 in
  let buf = Bytes.create buf_len in
  add_conn ();
  let rec loop total =
    let n = Io.read io buf 0 buf_len in
    let total = n + total in
    if n = 0 then (
      Unix.close fd;
      total
    ) else (
      let wn = writen io buf 0 n in
      if n <> wn then
        Printf.printf "read %d wrote %d (%b)\n%!" n wn (n = wn);
      loop total)
  in
  Fun.protect (fun () -> loop 0|>ignore) ~finally:sub_conn

let accepter ():unit =
  let unix_sock = Unix.socket ~cloexec:true PF_INET SOCK_STREAM 0 in
  Unix.setsockopt unix_sock SO_REUSEADDR true;
  Unix.setsockopt unix_sock SO_REUSEPORT true;
  let saddr = Unix.ADDR_INET (Unix.inet_addr_any, port) in
  Unix.bind unix_sock saddr;
  Unix.listen unix_sock 64;
  let accept_io = Io.of_unix unix_sock in
  let rec loop children =
    let cli_fd, cli_saddr = Io.accept accept_io in
    let cli_task = Task.async_parallel (handle_client cli_fd cli_saddr) in
    loop (cli_task :: (Task.reap children))
  in
  loop []

let () =
  Scheduler.run @@ fun () ->
  let accepters =
    List.init (Domain.recommended_domain_count ())
      (fun _ -> Task.async_parallel accepter)
  in
  while true do
    print_stats ();
    Task.sleep 1;
  done;
  List.iter Task.await accepters