tcp_messaging.ml.old

open Message
open Messaging
open Lwt
open Log_extra
open Lwtq

let __open_connection socket_address = 
  (* Lwt_io.open_connection socket_address *)
  let socket = Lwt_unix.socket Unix.PF_INET Unix.SOCK_STREAM 0  in
  Lwt.catch
    (fun () ->
      Lwt_unix.connect socket socket_address >>= fun () ->
      let fd_field = Obj.field (Obj.repr socket) 0 in
      let (fdi:int) = Obj.magic (fd_field) in
      Lwt_log.info_f "__open_connection SUCCEEDED (fd=%i)" fdi >>= fun () ->
      let oc = Lwt_io.of_fd ~mode:Lwt_io.output socket in
      let ic = Lwt_io.of_fd ~mode:Lwt_io.input  socket in
      Lwt.return (ic,oc))
    (fun exn -> Lwt_log.info ~exn "__open_connection failed" >>= fun () ->
      Lwt_unix.close socket;
      Lwt.fail exn)

module LockedHT = struct
  type ('a, 'b) t = { mutable entries:   ('a * 'b) list;
		      lock : Lwt_mutex.t}
  let create (n:int) = { entries = [] ; lock = Lwt_mutex.create()}

  let _remove t a = t.entries <- List.remove_assoc a t.entries
  let _mem t a = List.mem_assoc a t.entries 
  
  let _set t a b = 
    if _mem t a then _remove t a;
    t.entries <- (a,b)::t.entries
  
  let _length t = List.length t.entries
  let _find t a = List.assoc a t.entries 

  let mem t a = Lwt_mutex.with_lock t.lock  (fun () -> Lwt.return (_mem t a))
       
  let set t a b = Lwt_mutex.with_lock t.lock (fun () -> let () = _set t a b in Lwt.return ())
  let find t a = Lwt_mutex.with_lock t.lock (fun () -> let r = _find t a in Lwt.return r)
  let remove t a = Lwt_mutex.with_lock t.lock (fun () -> let () = _remove t a in Lwt.return ())
  let length t = Lwt_mutex.with_lock t.lock (fun () -> let r = _length t in Lwt.return r)

  let remove_if_present t a = Lwt_mutex.with_lock t.lock 
    (fun () -> 
      if _mem t a then _remove t a;
      Lwt.return ())

  let with_lock t f = Lwt_mutex.with_lock t.lock (fun () -> f t)
    
end

type connection = Lwt_io.input_channel * Lwt_io.output_channel
type mq = (Message.t * id) LWTQ.t

class tcp_messaging my_address =
  let never () = false in

object(self : # messaging )
  val _id2address = Hashtbl.create 10
  val _connections = LockedHT.create 10
  val _connections_lock = Lwt_mutex.create ()
  val _qs = Hashtbl.create 10
  val _outgoing = LWTQ.create ()
  

  method register_receivers mapping =
    List.iter
      (fun (id,address) -> Hashtbl.add _id2address id address) mapping

  method private _get_target_address ~target =
    Hashtbl.find _id2address target


  method send_message m ~source ~target =
    LWTQ.add (source, target, m) _outgoing


  method private _get_q ~(target:id) =
    try Hashtbl.find _qs target
    with Not_found ->
      begin
  let tq = LWTQ.create () in
    Hashtbl.add _qs target tq;
          tq
      end

  method recv_message ~target =
    let q = self # _get_q ~target in LWTQ.take q

  method peek_message ~target =
    let q = self # _get_q ~target in LWTQ.peek q

  method private _establish_connection address =
    let host_ip, port = address in
    let socket_address = Network.make_address host_ip port in
    Lwt_log.debug_f "establishing connection to (%s,%i)" host_ip port
    >>= fun () ->
    (*Backoff.backoff (fun () -> *) __open_connection socket_address (* ) *)
    >>= fun (ic,oc) ->
    Lwt_log.debug_f "got connection to (%s,%i)" host_ip port >>= fun () ->
    let my_ip, my_port = my_address in
    Llio.output_string oc my_ip >>= fun () ->
    Llio.output_int oc my_port  >>= fun () ->
    Lwt.return (ic,oc)
      (* open_connection can also fail with Unix.Unix_error (63, "connect",_)
	 on local host *)

  method private _get_connection address =
    LockedHT.with_lock _connections 
      (fun t ->
	if LockedHT._mem t address 
	then 
	  Lwt.return (LockedHT._find t address )
	else
	  begin
	    self # _establish_connection address >>= fun conn ->
	    LockedHT._set t address conn;
	    Lwt_log.debug_f "connection added (%i in total)" (LockedHT._length t) >>= fun () ->
	    Lwt.return conn
	  end
      )

  method private _drop_connection address =

    Lwt_log.debug "_drop_connection" >>= fun () ->
    LockedHT.with_lock _connections
      (fun t ->
	if LockedHT._mem t address then
	  begin
	    let conn = LockedHT._find t address in
	    Lwt_log.debug "found connection, closing it" >>= fun () ->
	    let ic,oc = conn in
	    (* something with conn *)
	    Lwt.catch
	      (fun () ->
		Lwt_io.close ic >>= fun () ->
		Lwt_io.close oc >>= fun () ->
		Lwt_log.debug "closed connection"
	      )
	      (fun exn -> Lwt_log.warning ~exn "exception while closing, too little too late" )
	    >>= fun () ->
	    let () = LockedHT._remove t address in
	    Lwt.return ()
	  end
	else
	  begin
	    let h,p = address in
	    Lwt_log.debug_f "connection to (%s,%i) not found. we never had one..." h p
	  end
      )


  method private _pickle source target msg =
    let buffer = Buffer.create 40 in
    let () = Llio.string_to buffer source in
    let () = Llio.string_to buffer target in
    let () = Message.to_buffer msg buffer in
    Buffer.contents buffer


  method private _maybe_insert_connection address =
    let host,port = address in
    LockedHT.mem _connections address >>= function
      | true  -> Lwt_log.debug_f "XXX already have connection with (%S,%i)" host port
      | false -> Lwt_log.debug_f "XXX first connection with (%S,%i)" host port
	
  method run ?(stop=never) () =
    let conditionally f =
      let b = stop () in
      if b then
	begin
	  Lwt_log.info "ending loop" >>= fun () ->
	  Lwt.return ()
	end
      else f ()
    in
    let protocol (ic,oc) =
      Llio.input_string ic >>= fun ip ->
      Llio.input_int ic    >>= fun port ->
      self # _maybe_insert_connection (ip,port) >>= fun () ->
      begin
	let rec loop () =
	  begin
	    Llio.input_int ic >>= fun msg_size ->
	    let buffer = String.create msg_size in
	    Lwt_io.read_into_exactly ic buffer 0 msg_size >>= fun () ->
	    let (source:id), pos1 = Llio.string_from buffer 0 in
	    let target, pos2 = Llio.string_from buffer pos1 in
	    let msg, _   = Message.from_buffer buffer pos2 in
	    let q = self # _get_q target in
	    LWTQ.add (msg, source) q >>=  fun () ->
	    conditionally loop
	  end
	in
	catch
	  (fun () -> loop ())
	  (fun exn ->
	    Lwt_log.info ~exn "going to drop outgoing connection as well" >>= fun () ->
	    let address = (ip,port) in
	    self # _drop_connection address >>= fun () ->
	    Lwt.fail exn)

      end
    in
    let rec sender_loop () =
      LWTQ.take _outgoing >>= fun (source, target, msg) ->
      (* Lwt_log.debug_f "sender_loop got %S" (Message.string_of msg) >>= fun () -> *)
      let address = self # _get_target_address ~target in
      let try_send () =
	self # _get_connection address >>= fun connection ->
	let ic,oc = connection in
	let pickled = self # _pickle source target msg in
        Llio.output_string oc pickled >>= fun () ->
	Lwt_io.flush oc
      in
      Lwt.catch
	(fun () -> try_send ())
	(function
	  | Unix.Unix_error(Unix.EPIPE,_,_) -> (* stale connection *)
	    begin
	      Lwt_log.debug_f "stale connection" >>= fun () ->
	      self # _drop_connection address >>= fun () ->
	      Lwt.catch
		(fun () -> try_send ())
		(fun exn -> Lwt_log.info_f ~exn "dropped message")
	    end
	  | exn ->
	    begin
	      Lwt_log.info_f ~exn
		"dropping message %s with destination '%s' because of"
		(Message.string_of msg) target >>= fun () ->
	      self # _drop_connection address >>= fun () ->
	      Lwt_log.debug "end of connection epilogue"
	    end
	)
      >>= fun () -> conditionally sender_loop
    in
    let _, my_port = my_address in
    let server_t = Server.make_server_thread my_port protocol
    in
    Lwt.pick [server_t ();sender_loop ();] >>= fun () ->
    Lwt_log.info "end of tcp_messaging"    >>= fun () ->
    Lwt.return ()


end