(* Implementation of the controller process *)

open Sort2_proto_aux
open Printf

let worker_clients = Hashtbl.create 10
  (* Maps endpoint names to RPC clients *)


let rec get_worker_endpoints n_workers =
  (* Loop until at least n_workers are available *)
  let endpoints = 
    Netplex_cenv.lookup_container_sockets "sort_worker" "Worker" in
  if Array.length endpoints < n_workers then (
    Netlog.logf `Warning
      "Not enough workers found in registry (only %d)- will retry in 1 second"
      (Array.length endpoints);
    Unix.sleep 1;
    get_worker_endpoints n_workers
  )
  else (
    Netlog.logf `Info
      "Found %d endpoints, using the first %d" 
      (Array.length endpoints) n_workers;
    Array.sub endpoints 0 n_workers
  )


let get_worker_client esys endpoint =
  try
    Hashtbl.find worker_clients endpoint
  with
    | Not_found ->
        let connector =
          Netplex_sockserv.any_file_client_connector endpoint in
        let client =
          Sort2_proto_clnt.Worker.V1.create_client2
            ~esys
            (`Socket(Rpc.Tcp, connector, Rpc_client.default_socket_config)) in
        Hashtbl.replace worker_clients endpoint client;
        client


let get_subarray_info n_workers n =
  let slice_len = n / n_workers in
  let m = n mod n_workers in
  let k = ref 0 in
  Array.init
    n_workers
    (fun i ->
       let k0 = !k in
       let l =
	 if i < m then slice_len+1 else slice_len in
       k := !k + l;
       (k0, l)
    )


let shm_tbl = Hashtbl.create 10
  (* maps name to shm *)


let rec shm_open_excl pid n =
  try
    let name = sprintf "/sort2_%d_%d" pid n in
    let fd = 
      Netsys_posix.shm_open
	name
	[Netsys_posix.SHM_O_RDWR; 
	 Netsys_posix.SHM_O_CREAT;
	 Netsys_posix.SHM_O_EXCL]
	0o600 in
    (name,fd)
  with
    | Unix.Unix_error(Unix.EEXIST,_,_) ->
	shm_open_excl pid (n+1)


let shm_open_existing shm =
  let fd = 
    Netsys_posix.shm_open
      shm.shm_name
      [Netsys_posix.SHM_O_RDWR]
      0o600 in
  let mem = Netsys_mem.memory_map_file fd true (-1) in
  Unix.close fd;
  mem


let free_shm ?(immediate_unmap=false) shm =
  let (_,mem) =
    try
      Hashtbl.find shm_tbl shm.shm_name
    with
      | Not_found ->
	  failwith("Shared mem object not found: " ^ shm.shm_name) in
  Netsys_posix.shm_unlink shm.shm_name;
  if immediate_unmap then
    Netsys_mem.memory_unmap_file mem;   (* dangerous ! *)
  Hashtbl.remove shm_tbl shm.shm_name


let free_all() =
  let shm_list =
    Hashtbl.fold (fun _ (shm,_) acc -> shm::acc) shm_tbl [] in
  List.iter free_shm shm_list    


let dummy_mem =
  Netsys_mem.alloc_memory_pages 1


let sort_subarray esys worker_endpoint data (k,l) when_done when_error =
  (* The following is written as an asynchrounous sequence of RPC calls.
     On success, finally when_done is called. On error, when_error
     is called.
   *)

  let sdata = Array.sub data k l in   (* TODO: avoid this copy *)
  Netlog.logf `Info
    "sort_subarray k=%d l=%d" k l;
  
  (* First find out how large sdata is: *)
  let _,sdata_bytelen =
    Netsys_mem.init_value
      dummy_mem 0 sdata [Netsys_mem.Copy_simulate] in
  Netlog.logf `Info "first init_value done bytelen=%d" sdata_bytelen;
  
  (* Alloc the buffer in the controller: *)
  let ctrl_mem_name, ctrl_mem_fd =
    shm_open_excl (Unix.getpid()) 0 in
  let ctrl_mem =
    Netsys_mem.memory_map_file ctrl_mem_fd true sdata_bytelen in
  Unix.close ctrl_mem_fd;
  let ctrl_shm =
    { shm_name = ctrl_mem_name;
      shm_addr = Int64.of_nativeint(Netsys_mem.memory_address ctrl_mem)
    } in
  Hashtbl.replace shm_tbl ctrl_mem_name (ctrl_shm, ctrl_mem);
  Netlog.logf `Info "addr=%Lx" ctrl_shm.shm_addr;

  (* Alloc the buffer in the worker: *)
  let client = get_worker_client esys worker_endpoint in
  Sort2_proto_clnt.Worker.V1.alloc_shm'async
    client
    (Int64.of_int sdata_bytelen)
    (fun alloc_shm_reply ->
       try
	 let worker_shm = alloc_shm_reply() in
	 (* Got this buffer. Now copy the subarray to this buffer. *)
	 let worker_mem = shm_open_existing worker_shm in
	 let worker_offset,_ =
	   Netsys_mem.init_value
	     ~targetaddr:(Int64.to_nativeint worker_shm.shm_addr)
	     worker_mem 0 sdata [] in
	 Netlog.logf `Info
	   "Second init_value done";
	 (* Request the sort: *)
	 Sort2_proto_clnt.Worker.V1.sort_shm'async
	   client
	   (worker_shm,worker_offset)
	   (fun sort_shm_reply ->
	      try
		let () = sort_shm_reply() in   (* check for errors! *)
		(* Copy the sorted data to the controller buffer, so we can
                   access them
		 *)
		Sort2_proto_clnt.Worker.V1.copy_shm'async
		  client
		  (worker_shm,worker_offset,ctrl_shm)
		  (fun copy_shm_reply ->
		     try
		       Netlog.logf `Info "Got data in ctrl buffer";
		       let ctrl_offset = copy_shm_reply() in
		       (* Get the sorted subarray: *)
		       let sdata_sorted =
			 Netsys_mem.as_value ctrl_mem ctrl_offset in
		       Netsys_mem.value_area ctrl_mem;
		       when_done 
			 sdata_sorted
			 (fun () -> free_shm ~immediate_unmap:true ctrl_shm);
		       (* Do some cleanup. Note that we assume now that
                          we don't access sdata_sorted any longer! 
			*)
		       Sort2_proto_clnt.Worker.V1.free_shm'async
			 client
			 (worker_shm)
			 (fun _ -> ())
		     with
		       | error ->
			   when_error error
		  )
	      with
		| error ->
		    when_error error
	   )
       with
	 | error ->
	     when_error error
    )
  
exception Merge_exit
exception Merge_done_array of int


let rec merge_into_0 out k_out in_arrays k_in =
  let l_in = Array.length in_arrays in

  assert(l_in > 0);

  if l_in = 1 then (
    let a = !(in_arrays.(0)) in
    let l = Array.length a - k_in.(0) in
    Array.blit a k_in.(0) out !k_out l;
    k_out := !k_out + l
  )
  else
    merge_into_1 out k_out in_arrays k_in
  
and merge_into_1 out k_out in_arrays k_in =
  (* at least two in_arrays *)
  let l_in = Array.length in_arrays in
  let l_out = Array.length out in
  try
    while true do
      (* Find the smallest in_array element: *)
      let smallest_v = ref !(in_arrays.(0)).(k_in.(0)) in
      let smallest_p = ref 0 in
      for i = 1 to l_in-1 do
	let v = !(in_arrays.(i)).(k_in.(i)) in
	if v < !smallest_v then (
	  smallest_v := v;
	  smallest_p := i
	)
      done;
      out.( !k_out ) <- !smallest_v;
      incr k_out;
      if !k_out = l_out then raise Merge_exit;
      let k_new = k_in.( !smallest_p ) + 1 in
      k_in.( !smallest_p ) <- k_new;
      if k_new = Array.length !(in_arrays.( !smallest_p )) then
	raise (Merge_done_array !smallest_p)
    done
  with
    | Merge_exit -> 
	()
    | Merge_done_array p ->
	(* We are done with in_array p. Remove it from our structure *)
	let in_arrays' =
	  Array.init
	    (Array.length in_arrays - 1)
	    (fun i ->
	       if i < p then in_arrays.(i) else in_arrays.(i+1)
	    ) in
	let k_in' = 
	  Array.init
	    (Array.length k_in - 1)
	    (fun i ->
	       if i < p then k_in.(i) else k_in.(i+1)
	    ) in
	merge_into_0 out k_out in_arrays' k_in'


let merge_into out in_arrays =
  let t0 = Unix.gettimeofday() in
  Netlog.logf `Info "merge_into";
  let l_in = Array.length in_arrays in
  assert(l_in > 0);
  let k_out = ref 0 in
  let k_in = Array.make l_in 0 in
  merge_into_0 out k_out in_arrays k_in;
  assert (!k_out = Array.length out);
  let t1 = Unix.gettimeofday() in
  Netlog.logf `Info "time for merge_into: %f" (t1-.t0)


let sort esys data worker_endpoints when_done =
  (* sorts asynchronously. Calls [when_done] with [Some sorted_data] on
     success, or [when_done None] on error
   *)
  let n_workers = Array.length worker_endpoints in
  let subarray_info = 
    get_subarray_info n_workers (Array.length data) in
  let to_merge =
    Array.make n_workers (ref [| |]) in
    
  let runcounter = ref 0 in
  let errorflag = ref false in
    
  Array.iteri
    (fun p (k,l) ->
       let worker_endpoint = worker_endpoints.(p) in
       sort_subarray
	 esys worker_endpoint data (k,l) 
	 (fun sdata_sorted cleanup ->  (* when_done *)
	    decr runcounter;
	    if not !errorflag then (
	      (* Note that sdata_sorted is only valid memory until [cleanup]
                 is called. Be careful with that! We cannot attach a finaliser
                 to sdata_sorted directly (it is living outside of the heap)
                 but with some care the ref cell around it will work.

                 Note that in there is a ref to the bigarray in the cleanup
                 closure, and this ref prevents that the shm is unmapped.
	       *)
	      let sdata_sorted_cell = ref sdata_sorted in
	      to_merge.(p) <- sdata_sorted_cell;
	      Gc.finalise 
		(fun _ -> Netlog.logf `Info "Cleanup"; cleanup()) 
		sdata_sorted_cell;
	      (* If all subarrays have arrived, start the merge: *)
	      if !runcounter = 0 then (
		merge_into data to_merge;
		(* and reply: *)
		when_done (Some data)
	      )
	    );
	    if !errorflag && !runcounter = 0 then (
	      when_done None
	    )
	 )
	 (fun error ->         (* when_error *)
	    decr runcounter;
	    Netlog.logf `Err
	      "Got exception: %s" (Netexn.to_string error);
	    errorflag := true;
	    if !runcounter = 0 then (
	      when_done None
	    )
	 );
       incr runcounter
    )
    subarray_info


let emit_error session =
  Rpc_server.reply_error session Rpc.System_err


let proc_sort n_workers session (data,sort_flag) emit =
  if sort_flag then (
    let esys = (Netplex_cenv.self_cont()) # event_system in
    let worker_endpoints =
      get_worker_endpoints n_workers in
    sort 
      esys data worker_endpoints
      (function
	 | Some sorted_data ->
	     emit sorted_data
	 | None ->
	     emit_error session
      )
  )
  else
    emit data


let proc_get_workers n_workers session () emit =
  let worker_endpoints =
    get_worker_endpoints n_workers in
  emit worker_endpoints



let configure cf addr =
  let n_workers =
    try cf # int_param(cf # resolve_parameter addr "n_workers") 
    with Not_found -> 1 in
  Netlog.logf `Info
    "Using param n_workers=%d" n_workers;
  (n_workers)



let setup srv (n_workers) =
  Sort2_proto_srv.Interface.V1.bind_async
    ~proc_null:(fun _ _ emit -> emit ())
    ~proc_sort:(proc_sort n_workers)
    ~proc_get_workers:(proc_get_workers n_workers)
    srv

let controller_factory() =
  Rpc_netplex.rpc_factory
    ~name:"sort"
    ~configure
    ~setup
    ~hooks:(fun _ ->
              object(self)
                inherit Netplex_kit.empty_processor_hooks() 
                method post_start_hook _ =
                  let _t =
                    Netplex_cenv.create_timer
                      (fun _ -> 
                         Gc.major();
                         true
                      )
                      1.0 in
		  List.iter
		    (fun signo ->
		       Netsys_signal.register_handler
			 ~name:"Sort2_controller"
			 ~signal:signo
			 ~callback:(fun _ -> free_all())
			 ~keep_default:true
			 ()
		    )
		    [ Sys.sigint; Sys.sigterm ];
                  ()
		method pre_finish_hook _ =
		  free_all()
              end
           )
    ()