At Jane Street, we end up writing lots of messaging protocols, and many of these protocols end up being simple RPC-style protocols, i.e., protocols with a client and a server, where communication is done in a simple query/response style.

I’ve always found the writing of these protocols rather unsatisfying, because I could never find a clean way of writing down the types. In the following, I’d like to describe some nice tricks I’ve learned recently for specifying these protocols more cleanly.

A Simple Example

I’ll start with a concrete example: a set of RPCs for accessing a remote filesystem. Here are the signatures for a set of functions that we want to make available via RPC.

type path = Path of string list with sexp 
type 'a result = Ok of 'a | Error of string with sexp

val listdir : path -> string list result
val read_file : path -> string result
val move : path * path -> unit result
val put_file : path * string -> unit result
val file_size : path -> int result
val file_exists : path -> bool 

The with sexp appended to the end of the type definitions comes from the Jane Street’s publicly available sexplib macros. These macros generate functions for converting values to and from s-expressions. This is fantastically helpful for writing messaging protocols, since it gives you a simple hassle-free mechanism for serializing values over the wire. (Unfortunately, s-expression generation is not the fastest thing in the world, which is why we’ve written a set of binary serialization macros for high-performance messaging applications, which we intend to release.)

The usual next step would be to build up two types, one for the queries and one for the responses. Here’s what you might write down to support the functions shown above.

module Request = struct 
  type t = | Listdir of path 
      | Read_file of path 
      | Move of path * path 
      | Put_file of path * string 
      | File_size of path 
      | File_exists of path 
  with sexp 
end

module Response = struct 
  type t = | Ok 
  | Error of string 
  | File_size of int 
  | Contents of string list 
  | File_exists of bool 
  with sexp  
end 

In some ways, this is great. The types are simple to write down and understand, and you get the wire protocol virtually for free from the s-expression converters. And both the server and the client code are pretty easy to write. Let’s look at how that code might look.

First, let’s assume we have functions for sending and receiving s-expressions over some connection object, with the following signature:

val send : conn -> Sexp.t -> unit 
val recv : conn -> Sexp.t 

Then the server code should look something like this:

let handle_query conn = 
  let module Q = Query in 
  let module R = Response in 
  let msg = Q.t_of_sexp (recv conn) in 
  let resp = 
    match query with 
    | Q.Listdir path -> 
      begin match listdir path with 
      | Ok x -> R.Contents x 
      | Error s -> R.Error s 
      end 
    | Q.Read_file path -> 
    . 
    . 
    . 
  in 
  send (R.sexp_of_t resp) 

And the client code could look like something this:

let rpc_listdir conn path = 
  let module Q = Query in 
  let module R = Response in 
  send conn (Q.sexp_of_t (Q.Listdir path)); 
  match R.t_of_sexp (recv conn) with 
  | R.Contents x -> Ok x 
  | R.Error s -> Error s 
  | _ -> assert false 

Unfortunately, to make this all work, you’ve been forced to turn your type definitions sideways: rather than specifying for each RPC a pair of a request type and a response type, as you do in the specification of ordinary function type, you have to specify all the requests and all the responses at once. And there’s nothing in the types tying the two sides together. This means that there is no consistency check between the server code and the client code. In particular, the server code could receive a File_size query and return Contents, or Ok, when really it should only be returning either a File_size or Error, and you would only catch it at runtime.

Specifying RPCs with Embeddings

But all is not lost! With just a little bit of infrastructure, we can specify our protocol in a way that ties together the client and server pieces. The first thing we need is something that we’re going to call an embedding, but which you might see referred to elsewhere as an embedding-projection pair. An embedding is basically a pair of functions, one for converting values of a given type into some universal type, and the other for converting back from the universal type. (For another take on universal types, take a look at this post from Steven). The universal type we’ll use is S-expressions:

type 'a embedding = { inj: 'a -> Sexp.t; 
prj: Sexp.t -> 'a; } 

It’s worth noting that the projection function is always going to be partial, meaning it will fail on some inputs. In this case, we’ll encode that partiality with exceptions, since our s-expression macro library generates conversion functions that throw exceptions when a value doesn’t parse. But it’s often better to explicitly encode the partiality in the return type of the projection function.

We can now write up a type that specifies the type of the RPC from which we can derive both the client and server code.

module RPC = struct 
  type ('a,'b) t = { 
    tag: string; 
    query: 'a embedding;
    resp: 'b embedding;
  } 
end 

Here’s a how you could write the RPC.t corresponding to the listdir function:

module RPC_specs = struct 
  type listdir_resp = string list result with sexp 
  let listdir = { RPC. 
    tag = "listdir"; 
    query = { inj = sexp_of_path; 
          prj = path_of_sexp; }; 
    resp = { inj = sexp_of_listdir_resp; 
          prj = listdir_resp_of_sexp; }; 
  }

  .
  .
  .

end 

One slightly annoying aspect of the above code is that we had to define the type listdir_resp purely for the purpose of getting the corresponding s-expression converters. At some point, we should do a post on type-indexed values to explain how one could get around the need for such a declaration.

Note that the above specifies the interface, but not actually the function used to implement the RPC on the server side. The embeddings basically specify the types of the requests and responses, and the tag is used to distinguish different RPCs on the wire.

As you may have noticed, an ('a,'b) RPC.t corresponds to a function of type 'a -> 'b. We can put this correspondence to work by writing a function that takes an ('a,'b) RPC.t and an ordinary function of type

'a -> 'b

and produces an RPC handler. We’ll write down a simple implementation below.

type full_query = string * Sexp.t with sexp 
(* The first part is the tag, the second half is the s-expression for the arguments to the query. We only declare this type to get the s-expression converters *)

module Handler : sig 
  type t 
  val implement : ('a,'b) RPC.t -> ('a -> 'b) -> t 
  val handle : t list -> Sexp.t -> Sexp.t 
end
 = 
 struct 
  type t = { tag: string; 
        handle: Sexp.t -> Sexp.t; }

let implement rpc f = 
  { tag = rpc.RPC.tag; 
  handle = (fun sexp -> 
    let query = rpc.RPC.query.prj sexp in 
    rpc.RPC.resp.inj (f query)); }

let handle handlers sexp = 
  let (tag,query_sexp) = full_query_of_sexp sexp in 
  let handler = List.find ~f:(fun x -> x.tag = tag) handlers in 
  handler.handle query_sexp 
end 

Using the RPC.t’s we started writing as part of the RPC_specs module, we can now write the server as follows:

let handle_query conn = 
  let query = recv conn in 
  let resp = 
    Handler.handle [ Handler.implement RPC_specs.listdir listdir; 
            Handler.implement RPC_specs.read_file read_file; 
            Handler.implement RPC_specs.move move; 
            Handler.implement RPC_specs.put_file put_file; 
            Handler.implement RPC_specs.file_size file_size;] 
    query 
  in 
  send conn resp 

And we can implement the client side just as easily.

let query rpc conn x = 
  let query_sexp = rpc.RPC.query.inj x in 
  send (sexp_of_full_query (rpc.RPC.tag,query_sexp)); 
  rpc.RPC.resp.prj (recv conn)

module Client : sig 
  val listdir : path -> string list result 
  val read_file : path -> string result 
  val move : path * path -> unit result 
  val put_file : path * string -> unit result 
  val file_size : path -> int result 
  val file_exists : path -> bool 
end
 = 
struct 
  let listdir = query RPC_specs.listdir 
  let read_file = query RPC_specs.read_file 
  let move = query RPC_specs.move 
  let put_file = query RPC_specs.put_file 
  let file_size = query RPC_specs.file_size 
  let file_exists = query RPC_specs.file_exists 
end 

Pleasantly, the signature of the Client module is exactly the same as the signature of the underlying functions we’re exposing via RPC.

To be clear, this is far from a complete implementation – particularly notable is the weak error handling, and we haven’t said anything about how to deal with versioning of the protocol. But even though the implementation we’ve sketched out is a toy, we think this approach scales well to a full implementation.

There are still some problems. Although we’ve added static checks for some errors, we’ve eliminated some others. For instance, it’s now possible for the user to specify multiple RPC.t’s with the same tag, and there’s no guarantee that the server has exhaustively implemented all of expected RPC.t’s. I’m not aware of a clean way of getting all of these static checks working cleanly together in the same implementation.