rpc: RPC recording functions (#562)

* rpc: RPC recording functions

* fork.yaml: update fork.yaml with note about RPC recording functionality

* rpc: improve MsgError to return JsonError to avoid interface around typed nil

Author: protolambda

fork.yaml

@@ -327,6 +327,18 @@ def:
     - title: "Geth extras"
       description: Extend the tools available in geth to improve external testing and tooling.
       sub:
+        - title: JSON-RPC recording
+          description: |
+            Extend server and client with configurable JSON-RPC message recording.
+            This enables loggers and advanced metrics to be attached to server and client sides.
+          globs:
+            - "rpc/client.go"
+            - "rpc/client_opt.go"
+            - "rpc/handler.go"
+            - "rpc/inproc.go"
+            - "rpc/recording.go"
+            - "rpc/server.go"
+            - "rpc/subscription.go"
         - title: Simulated Backend
           globs:
             - "accounts/abi/bind/backends/simulated.go"

rpc/client.go

@@ -105,6 +105,8 @@ type Client struct {
 	reqInit     chan *requestOp  // register response IDs, takes write lock
 	reqSent     chan error       // signals write completion, releases write lock
 	reqTimeout  chan *requestOp  // removes response IDs when call timeout expires
+
+	recorder Recorder // optional, may be nil
 }
 
 type reconnectFunc func(context.Context) (ServerCodec, error)
@@ -121,6 +123,7 @@ func (c *Client) newClientConn(conn ServerCodec) *clientConn {
 	ctx = context.WithValue(ctx, clientContextKey{}, c)
 	ctx = context.WithValue(ctx, peerInfoContextKey{}, conn.peerInfo())
 	handler := newHandler(ctx, conn, c.idgen, c.services, c.batchItemLimit, c.batchResponseMaxSize)
+	handler.recorder = c.recorder
 	return &clientConn{conn, handler}
 }
 
@@ -258,6 +261,7 @@ func initClient(conn ServerCodec, services *serviceRegistry, cfg *clientConfig)
 		reqInit:              make(chan *requestOp),
 		reqSent:              make(chan error, 1),
 		reqTimeout:           make(chan *requestOp),
+		recorder:             cfg.recorder,
 	}
 
 	// Set defaults.
@@ -343,6 +347,10 @@ func (c *Client) CallContext(ctx context.Context, result interface{}, method str
 	if err != nil {
 		return err
 	}
+	var recordDone RecordDone
+	if c.recorder != nil {
+		recordDone = c.recorder.RecordOutgoing(ctx, msg)
+	}
 	op := &requestOp{
 		ids:  []json.RawMessage{msg.ID},
 		resp: make(chan []*jsonrpcMessage, 1),
@@ -363,6 +371,9 @@ func (c *Client) CallContext(ctx context.Context, result interface{}, method str
 		return err
 	}
 	resp := batchresp[0]
+	if recordDone != nil {
+		recordDone(ctx, msg, resp)
+	}
 	switch {
 	case resp.Error != nil:
 		return resp.Error
@@ -415,7 +426,13 @@ func (c *Client) BatchCallContext(ctx context.Context, b []BatchElem) error {
 		op.ids[i] = msg.ID
 		byID[string(msg.ID)] = i
 	}
-
+	var recordDone []RecordDone
+	if c.recorder != nil {
+		recordDone = make([]RecordDone, len(b))
+		for i, msg := range msgs {
+			recordDone[i] = c.recorder.RecordOutgoing(ctx, msg)
+		}
+	}
 	var err error
 	if c.isHTTP {
 		err = c.sendBatchHTTP(ctx, op, msgs)
@@ -446,6 +463,10 @@ func (c *Client) BatchCallContext(ctx context.Context, b []BatchElem) error {
 		}
 		delete(byID, string(resp.ID))
 
+		if recordDone != nil {
+			recordDone[index](ctx, msgs[index], resp)
+		}
+
 		// Assign result and error.
 		elem := &b[index]
 		switch {
@@ -462,6 +483,9 @@ func (c *Client) BatchCallContext(ctx context.Context, b []BatchElem) error {
 	for _, index := range byID {
 		elem := &b[index]
 		elem.Error = ErrMissingBatchResponse
+		if recordDone != nil {
+			recordDone[index](ctx, msgs[index], nil)
+		}
 	}
 
 	return err

rpc/client_opt.go

@@ -41,6 +41,8 @@ type clientConfig struct {
 	idgen              func() ID
 	batchItemLimit     int
 	batchResponseLimit int
+
+	recorder Recorder
 }
 
 func (cfg *clientConfig) initHeaders() {

rpc/handler.go

@@ -68,6 +68,9 @@ type handler struct {
 
 	subLock    sync.Mutex
 	serverSubs map[ID]*Subscription
+
+	// optional, may be nil
+	recorder Recorder
 }
 
 type callProc struct {
@@ -450,6 +453,12 @@ func (h *handler) handleResponses(batch []*jsonrpcMessage, handleCall func(*json
 
 // handleSubscriptionResult processes subscription notifications.
 func (h *handler) handleSubscriptionResult(msg *jsonrpcMessage) {
+	if h.recorder != nil {
+		recordDone := h.recorder.RecordIncoming(h.rootCtx, msg)
+		if recordDone != nil {
+			defer recordDone(h.rootCtx, msg, nil)
+		}
+	}
 	var result subscriptionResult
 	if err := json.Unmarshal(msg.Params, &result); err != nil {
 		h.log.Debug("Dropping invalid subscription message")
@@ -461,7 +470,19 @@ func (h *handler) handleSubscriptionResult(msg *jsonrpcMessage) {
 }
 
 // handleCallMsg executes a call message and returns the answer.
-func (h *handler) handleCallMsg(ctx *callProc, msg *jsonrpcMessage) *jsonrpcMessage {
+func (h *handler) handleCallMsg(cp *callProc, msg *jsonrpcMessage) *jsonrpcMessage {
+	var recordDone RecordDone
+	if h.recorder != nil {
+		recordDone = h.recorder.RecordIncoming(cp.ctx, msg)
+	}
+	out := h.handleCallMsgInner(cp, msg)
+	if recordDone != nil {
+		recordDone(cp.ctx, msg, out)
+	}
+	return out
+}
+
+func (h *handler) handleCallMsgInner(ctx *callProc, msg *jsonrpcMessage) *jsonrpcMessage {
 	start := time.Now()
 	switch {
 	case msg.isNotification():

rpc/inproc.go

@@ -22,9 +22,12 @@ import (
 )
 
 // DialInProc attaches an in-process connection to the given RPC server.
-func DialInProc(handler *Server) *Client {
+func DialInProc(handler *Server, options ...ClientOption) *Client {
 	initctx := context.Background()
 	cfg := new(clientConfig)
+	for _, opt := range options {
+		opt.applyOption(cfg)
+	}
 	c, _ := newClient(initctx, cfg, func(context.Context) (ServerCodec, error) {
 		p1, p2 := net.Pipe()
 		go handler.ServeCodec(NewCodec(p1), 0)

rpc/recording.go

@@ -0,0 +1,98 @@
+package rpc
+
+import (
+	"context"
+	"encoding/json"
+)
+
+// RecordDone is called after an incoming input (request or notification) has successfully been processed,
+// and captures the output (nil or response).
+type RecordDone func(ctx context.Context, input, output RecordedMsg)
+
+// Recorder captures RPC traffic
+type Recorder interface {
+	// RecordIncoming records an incoming message (request or notification), before it has been processed.
+	// It may optionally return a function to capture the result of the processing (response or nil).
+	RecordIncoming(ctx context.Context, msg RecordedMsg) RecordDone
+
+	// RecordOutgoing records an outgoing message (request or notification), before it has been sent.
+	// It may optionally return a function to capture the result of the request (response or nil).
+	RecordOutgoing(ctx context.Context, msg RecordedMsg) RecordDone
+}
+
+// RecordedMsg wraps around the internal jsonrpcMessage type,
+// to provide a public read-only interface for recording of RPC activity.
+// Every method name is prefixed with "Msg", to avoid conflict with internal methods and future geth changes.
+type RecordedMsg interface {
+	MsgIsNotification() bool
+	MsgIsResponse() bool
+	MsgID() json.RawMessage
+	MsgMethod() string
+	MsgParams() json.RawMessage
+	MsgError() *JsonError
+	MsgResult() json.RawMessage
+}
+
+var _ RecordedMsg = (*jsonrpcMessage)(nil)
+
+func (msg *jsonrpcMessage) MsgIsNotification() bool {
+	return msg.isNotification()
+}
+
+func (msg *jsonrpcMessage) MsgIsResponse() bool {
+	return msg.isResponse()
+}
+
+func (msg *jsonrpcMessage) MsgID() json.RawMessage {
+	return msg.ID
+}
+func (msg *jsonrpcMessage) MsgMethod() string {
+	return msg.Method
+}
+func (msg *jsonrpcMessage) MsgParams() json.RawMessage {
+	return msg.Params
+}
+func (msg *jsonrpcMessage) MsgError() *JsonError {
+	return msg.Error
+}
+func (msg *jsonrpcMessage) MsgResult() json.RawMessage {
+	return msg.Result
+}
+
+var _ RecordedMsg = (*jsonrpcSubscriptionNotification)(nil)
+
+func (notif *jsonrpcSubscriptionNotification) MsgIsNotification() bool {
+	return true
+}
+
+func (notif *jsonrpcSubscriptionNotification) MsgIsResponse() bool {
+	return false
+}
+
+func (notif *jsonrpcSubscriptionNotification) MsgID() json.RawMessage {
+	return json.RawMessage{} // notifications do not have an ID
+}
+
+func (notif *jsonrpcSubscriptionNotification) MsgMethod() string {
+	return notif.Method
+}
+
+func (notif *jsonrpcSubscriptionNotification) MsgParams() json.RawMessage {
+	data, _ := json.Marshal(notif.Params)
+	return data
+}
+
+func (notif *jsonrpcSubscriptionNotification) MsgError() *JsonError {
+	return nil
+}
+
+func (notif *jsonrpcSubscriptionNotification) MsgResult() json.RawMessage {
+	return nil
+}
+
+// WithRecorder attaches a recorder to an RPC client, useful when it is serving bidirectional RPC
+func WithRecorder(r Recorder) ClientOption {
+	return optionFunc(func(cfg *clientConfig) {
+		cfg.recorder = r
+	})
+}