RPC tracing for Golang (ala Dapper)

gistfile1.go

diff --git a/client.go b/client.go
index 4b0c9c3..8ac98e4 100644
--- a/client.go
+++ b/client.go
@@ -32,6 +32,7 @@ type Call struct {
 	Reply         interface{} // The reply from the function (*struct).
 	Error         error       // After completion, the error status.
 	Done          chan *Call  // Strobes when call is complete.
+	Trace         uint32      // unique id used for tracing
 }
 
 // Client represents an RPC Client.
@@ -85,6 +86,7 @@ func (client *Client) send(call *Call) {
 	// Encode and send the request.
 	client.request.Seq = seq
 	client.request.ServiceMethod = call.ServiceMethod
+	client.request.Trace = call.Trace
 	err := client.codec.WriteRequest(&client.request, call.Args)
 	if err != nil {
 		client.mutex.Lock()
@@ -284,11 +286,12 @@ func (client *Client) Close() error {
 // the invocation.  The done channel will signal when the call is complete by returning
 // the same Call object.  If done is nil, Go will allocate a new channel.
 // If non-nil, done must be buffered or Go will deliberately crash.
-func (client *Client) Go(serviceMethod string, args interface{}, reply interface{}, done chan *Call) *Call {
+func (client *Client) Go(serviceMethod string, args interface{}, reply interface{}, done chan *Call, trace uint32) *Call {
 	call := new(Call)
 	call.ServiceMethod = serviceMethod
 	call.Args = args
 	call.Reply = reply
+	call.Trace = trace
 	if done == nil {
 		done = make(chan *Call, 10) // buffered.
 	} else {
@@ -306,7 +309,7 @@ func (client *Client) Go(serviceMethod string, args interface{}, reply interface
 }
 
 // Call invokes the named function, waits for it to complete, and returns its error status.
-func (client *Client) Call(serviceMethod string, args interface{}, reply interface{}) error {
-	call := <-client.Go(serviceMethod, args, reply, make(chan *Call, 1)).Done
+func (client *Client) Call(serviceMethod string, args interface{}, reply interface{}, trace uint32) error {
+	call := <-client.Go(serviceMethod, args, reply, make(chan *Call, 1), trace).Done
 	return call.Error
 }
diff --git a/server.go b/server.go
index e71b6fb..1dd04d6 100644
--- a/server.go
+++ b/server.go
@@ -167,6 +167,7 @@ type service struct {
 type Request struct {
 	ServiceMethod string   // format: "Service.Method"
 	Seq           uint64   // sequence number chosen by client
+	Trace         uint32   // unique id for tracing
 	next          *Request // for free list in Server
 }
 
@@ -275,6 +276,15 @@ func (server *Server) register(rcvr interface{}, name string, useName bool) erro
 		log.Print(str)
 		return errors.New(str)
 	}
+
+	// Special trace handler
+	if method, ok := s.typ.MethodByName("Trace"); ok {
+		s.method["Trace"] = &methodType{
+			method: method,
+			ArgType: method.Type.In(1),
+		}
+	}
+
 	server.serviceMap[s.name] = s
 	return nil
 }
@@ -287,6 +297,9 @@ func suitableMethods(typ reflect.Type, reportErr bool) map[string]*methodType {
 		method := typ.Method(m)
 		mtype := method.Type
 		mname := method.Name
+		if mname == "Trace" {
+			continue
+		}
 		// Method must be exported.
 		if method.PkgPath != "" {
 			continue
@@ -376,6 +389,11 @@ func (s *service) call(server *Server, sending *sync.Mutex, mtype *methodType, r
 	mtype.Unlock()
 	function := mtype.method.Func
 	// Invoke the method, providing a new value for the reply.
+	trace := s.method["Trace"]
+	if trace != nil {
+		traceId := reflect.ValueOf(req.Trace)
+		trace.method.Func.Call([]reflect.Value{s.rcvr, traceId})
+	}
 	returnValues := function.Call([]reflect.Value{s.rcvr, argv, replyv})
 	// The return value for the method is an error.
 	errInter := returnValues[0].Interface()
diff --git a/example/client/client.go b/example/client/client.go
new file mode 100644
index 0000000..2b3fd53
--- /dev/null
+++ b/example/client/client.go
@@ -0,0 +1,36 @@
+package main
+
+import (
+	"net/rpc"
+	"net/rpc/example"
+	"log"
+	"fmt"
+)
+
+func main() {
+
+	//client, err := rpc.Dial("tcp", "127.0.0.1:1234")
+	client, err := rpc.Dial("tcp", "127.0.0.1:1235")
+	if err != nil {
+		log.Fatal("dialing:", err)
+	}
+
+	// Synchronous call
+	/*
+	args := &example.Args{25,4}
+	var reply int
+	err = client.Call("Arith.Multiply", args, &reply, 123)
+	if err != nil {
+		log.Fatal(err)
+	}
+	fmt.Printf("Arith: %d*%d=%d\n", args.A, args.B, reply)
+	*/
+
+	var reply int
+	args := 25
+	err = client.Call("Dumb.Double", args, &reply, 12)
+	if err != nil {
+		log.Fatal(err)
+	}
+	fmt.Printf("Double: %d=%d\n", args, reply)
+}
diff --git a/example/common.go b/example/common.go
new file mode 100644
index 0000000..73da6f4
--- /dev/null
+++ b/example/common.go
@@ -0,0 +1,49 @@
+package example
+
+import (
+	"net/rpc"
+	"errors"
+	"log"
+)
+
+type Args struct {
+	A, B int
+}
+
+type Arith struct {
+	TraceId uint32
+}
+
+func (t *Arith) Trace(traceId uint32) error {
+	log.Printf("Trace: %d", traceId)
+	t.TraceId = traceId
+	return nil
+}
+
+func (t *Arith) Multiply(args *Args, reply *int) error {
+	log.Printf("Multiply (trace %d)", t.TraceId)
+	*reply = args.A * args.B
+	return nil
+}
+
+var Client *rpc.Client
+
+func SetClient(c *rpc.Client) {
+	Client = c
+}
+
+type Dumb struct {
+	TraceId uint32
+}
+
+func (d *Dumb) Trace(traceId uint32) error {
+	log.Printf("Trace: %d", traceId)
+	d.TraceId = traceId
+	return nil
+}
+
+func (d *Dumb) Double(num int, reply *int) error {
+	log.Printf("Double (trace %d)", d.TraceId)
+	args := &Args{num, 2}
+	return Client.Call("Arith.Multiply", args, &reply, d.TraceId)
+}
diff --git a/example/dumb/dumb.go b/example/dumb/dumb.go
new file mode 100644
index 0000000..cefb81b
--- /dev/null
+++ b/example/dumb/dumb.go
@@ -0,0 +1,26 @@
+package main
+
+import (
+	"net"
+	"net/rpc"
+	"net/rpc/example"
+	"log"
+)
+
+func main() {
+	dumb := new(example.Dumb)
+	rpc.Register(dumb)
+
+	client, err := rpc.Dial("tcp", "127.0.0.1:1234")
+	example.SetClient(client)
+
+	if err != nil {
+		log.Fatal("dialing:", err)
+	}
+
+	ln, err := net.Listen("tcp", ":1235")
+	if err != nil {
+		log.Fatal("listen error:", err)
+	}
+	rpc.Accept(ln)
+}
diff --git a/example/smart/smart.go b/example/smart/smart.go
new file mode 100644
index 0000000..384d863
--- /dev/null
+++ b/example/smart/smart.go
@@ -0,0 +1,19 @@
+package main
+
+import (
+	"net"
+	"net/rpc"
+	"net/rpc/example"
+	"log"
+)
+
+func main() {
+	arith := new(example.Arith)
+	rpc.Register(arith)
+
+	ln, err := net.Listen("tcp", ":1234")
+	if err != nil {
+		log.Fatal("listen error:", err)
+	}
+	rpc.Accept(ln)
+}