Poor man's distributed actors

open System
open System.Net
open System.Net.Sockets
open System.IO
open System.Threading
open System.Threading.Tasks

open Microsoft.FSharp.Control

open Nessos.FsPickler

type AsyncBuilder with
    member __.Bind(f : Task<'T>, g : 'T -> Async<'S>) = __.Bind(Async.AwaitTask f, g)
    member __.Bind(f : Task, g : unit -> Async<'S>) = __.Bind(f.ContinueWith ignore, g)

type Stream with
    member s.AsyncWriteBytes (bytes : byte []) =
        async {
            do! s.WriteAsync(BitConverter.GetBytes bytes.Length, 0, 4)
            do! s.WriteAsync(bytes, 0, bytes.Length)
            do! s.FlushAsync()
        }

    member s.AsyncReadBytes(length : int) =
        let rec readSegment buf offset remaining =
            async {
                let! read = s.ReadAsync(buf, offset, remaining)
                if read < remaining then
                    return! readSegment buf (offset + read) (remaining - read)
                else
                    return ()
            }

        async {
            let bytes = Array.zeroCreate<byte> length
            do! readSegment bytes 0 length
            return bytes
        }

    member s.AsyncReadBytes() =
        async {
            let! lengthArr = s.AsyncReadBytes 4
            let length = BitConverter.ToInt32(lengthArr, 0)
            return! s.AsyncReadBytes length
        }

// existentially pack reply channels

type private IReplyChannelContainer<'T> =
    abstract PostWithReply : MailboxProcessor<'T> -> Async<obj>

and private ReplyChannelContainer<'T, 'R>(msgB : AsyncReplyChannel<'R> -> 'T) =
    interface IReplyChannelContainer<'T> with
        member __.PostWithReply (mb : MailboxProcessor<'T>) = async {
            let! r = mb.PostAndAsyncReply msgB
            return r :> obj
        }

type private ServerRequest<'T> =
    | Post of 'T
    | PostWithReply of IReplyChannelContainer<'T>

and private ServerResponse =
    | Acknowledge
    | Reply of obj
    | Fault of exn

        
type TcpActor<'T> (mailbox : MailboxProcessor<'T>, endpoint : IPEndPoint, ?pickler : FsPickler) =

    let pickler = match pickler with None -> new FsPickler() | Some p -> p
    let listener = new TcpListener(endpoint)

    let rec serverLoop () = async {
        try
            let! (client : TcpClient) = listener.AcceptTcpClientAsync()

            use client = client
            use stream = client.GetStream()

            let! (bytes : byte []) = stream.AsyncReadBytes()

            let! response = async {
                try
                    let request = pickler.UnPickle<ServerRequest<'T>> bytes

                    match request with
                    | Post msg -> 
                        do mailbox.Post msg
                        return Acknowledge
                    | PostWithReply rcc ->
                        let! reply = rcc.PostWithReply mailbox
                        return Reply reply
                with e -> return Fault e
            }

            let bytes = pickler.Pickle response
            do! stream.AsyncWriteBytes bytes

        with e -> printfn "Server error: %A" e

        return! serverLoop ()
    }

    let cts = new CancellationTokenSource()
    do 
        listener.Start()
        Async.Start(serverLoop (), cts.Token)

    member __.Stop () = cts.Cancel() ; listener.Stop()
    member __.MailboxProcessor = mailbox
    member __.GetTcpClient () = new TcpActorClient<'T>(endpoint, pickler)


and TcpActorClient<'T>(serverEndpoint : IPEndPoint, ?pickler : FsPickler) =
        
    let pickler = match pickler with None -> new FsPickler() | Some p -> p

    let sendRequest (request : ServerRequest<'T>) = async {
        use client = new TcpClient()
        do! client.ConnectAsync(serverEndpoint.Address, serverEndpoint.Port)
        use stream = client.GetStream()

        let bytes = pickler.Pickle request
        do! stream.AsyncWriteBytes bytes
        let! (reply : byte []) = stream.AsyncReadBytes()

        return pickler.UnPickle<ServerResponse> reply
    }

    member __.PostAsync (msg : 'T) = async {
        let! response = sendRequest <| Post msg
        match response with
        | Acknowledge -> return ()
        | Reply o -> return failwithf "TcpActor: invalid response '%O'." o
        | Fault e -> return raise e
    }

    member __.PostAndReplyAsync (msgB : AsyncReplyChannel<'R> -> 'T) = async {
        let wrapper = new ReplyChannelContainer<'T,'R>(msgB)
        let! response = sendRequest <| PostWithReply wrapper
        match response with
        | Reply (:? 'R as r) -> return r
        | Reply o -> return failwithf "TcpActor: invalid response '%O'." o
        | Acknowledge -> return failwith "TcpActor: invalid response."
        | Fault e -> return raise e
    }

    member __.Post (msg : 'T) = __.PostAsync msg |> Async.RunSynchronously
    member __.PostAndReply (msgB) = __.PostAndReplyAsync msgB |> Async.RunSynchronously


type TcpActor private () =

    static let parseEndpoint (endpoint : string) =
        let tokens = endpoint.Split(':')
        if tokens.Length <> 2 then raise <| new FormatException("invalid endpoint")
        let ipAddr = Dns.GetHostAddresses(tokens.[0]).[0]
        let port = int <| tokens.[1]
        new IPEndPoint(ipAddr, port)

    static member Create<'T>(behaviour : MailboxProcessor<'T> -> Async<unit>, ipEndPoint : IPEndPoint, ?pickler : FsPickler) =
        let mailbox = MailboxProcessor.Start behaviour
        new TcpActor<'T>(mailbox, ipEndPoint, ?pickler = pickler)

    static member Create<'T>(behaviour : MailboxProcessor<'T> -> Async<unit>, ipEndPoint : string, ?pickler : FsPickler) =
        TcpActor.Create(behaviour, parseEndpoint ipEndPoint, ?pickler = pickler)

    static member Connect<'T>(serverEndPoint : IPEndPoint, ?pickler : FsPickler) =
        new TcpActorClient<'T>(serverEndPoint, ?pickler = pickler)

    static member Connect<'T>(serverEndPoint : string, ?pickler : FsPickler) =
        new TcpActorClient<'T>(parseEndpoint serverEndPoint, ?pickler = pickler)