rpc: implemented Cancel and Pending methods

rpc/cancel_test.go

@@ -0,0 +1,260 @@
+// Copyright 2016 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package rpc
+
+import (
+	"encoding/json"
+	"net"
+	"testing"
+	"time"
+
+	"golang.org/x/net/context"
+)
+
+type CancelTestService struct{}
+
+func (s *CancelTestService) BlockingFunction(ctx context.Context) error {
+	select {
+	case <-time.After(time.Second):
+		return nil
+	case <-ctx.Done():
+		time.Sleep(time.Millisecond * 200)
+		return ctx.Err()
+	}
+}
+
+func TestCancel(t *testing.T) {
+	server := NewServer()
+	service := &CancelTestService{}
+
+	if err := server.RegisterName("eth", service); err != nil {
+		t.Fatalf("unable to register test service %v", err)
+	}
+
+	clientConn, serverConn := net.Pipe()
+
+	go server.ServeCodec(NewJSONCodec(serverConn), OptionMethodInvocation|OptionSubscriptions)
+
+	out := json.NewEncoder(clientConn)
+	in := json.NewDecoder(clientConn)
+
+	request := map[string]interface{}{
+		"id":      1,
+		"method":  "eth_blockingFunction",
+		"version": "2.0",
+		"params":  []interface{}{},
+	}
+
+	cancelRequest := map[string]interface{}{
+		"id":      2,
+		"method":  "rpc_cancel",
+		"version": "2.0",
+		"params":  []interface{}{1},
+	}
+
+	// test uncanceled request
+	start := time.Now()
+	// send request
+	if err := out.Encode(request); err != nil {
+		t.Fatal(err)
+	}
+
+	var response JSONErrResponse
+	if err := in.Decode(&response); err != nil {
+		t.Fatal(err)
+	}
+	elapsed := time.Since(start)
+	msg := response.Error.Message
+	// expect uncanceled request to return after 1000ms with no error
+	if elapsed < time.Millisecond*900 {
+		t.Errorf("uncanceled request returned too early (elapsed: %v, expected: 1s)", elapsed)
+	}
+	if msg != "" {
+		t.Errorf("uncanceled request returned with unexpected error: %v", msg)
+	}
+	time.Sleep(time.Millisecond * 10)
+
+	// test canceled request
+	start = time.Now()
+	// send request again
+	if err := out.Encode(request); err != nil {
+		t.Fatal(err)
+	}
+	time.Sleep(time.Millisecond * 500)
+	// send cancel request
+	if err := out.Encode(cancelRequest); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := in.Decode(&response); err != nil {
+		t.Fatal(err)
+	}
+	elapsed = time.Since(start)
+	msg = response.Error.Message
+	// expect cancel request to return after 500ms with no error
+	if elapsed > time.Millisecond*600 {
+		t.Errorf("cancel request returned too late (elapsed: %v, expected: 500ms)", elapsed)
+	}
+	if msg != "" {
+		t.Errorf("cancel request returned with unexpected error: %v", msg)
+	}
+
+	if err := in.Decode(&response); err != nil {
+		t.Fatal(err)
+	}
+	elapsed = time.Since(start)
+	msg = response.Error.Message
+	// expect canceled request to return after 700ms with "context canceled" error
+	if elapsed < time.Millisecond*600 {
+		t.Errorf("canceled request returned too early (elapsed: %v, expected: 700ms)", elapsed)
+	}
+	if elapsed > time.Millisecond*800 {
+		t.Errorf("canceled request returned too late (elapsed: %v, expected: 700ms)", elapsed)
+	}
+	if msg != "context canceled" {
+		t.Errorf("canceled request returned with unexpected or no error: \"%v\" (expected: \"context canceled\")", msg)
+	}
+
+	clientConn.Close()
+}
+
+func TestPending(t *testing.T) {
+	server := NewServer()
+	service := &CancelTestService{}
+
+	if err := server.RegisterName("eth", service); err != nil {
+		t.Fatalf("unable to register test service %v", err)
+	}
+
+	clientConn, serverConn := net.Pipe()
+
+	go server.ServeCodec(NewJSONCodec(serverConn), OptionMethodInvocation|OptionSubscriptions)
+
+	out := json.NewEncoder(clientConn)
+	in := json.NewDecoder(clientConn)
+
+	request := map[string]interface{}{
+		"id":      1,
+		"method":  "eth_blockingFunction",
+		"version": "2.0",
+		"params":  []interface{}{},
+	}
+
+	cancelRequest := map[string]interface{}{
+		"id":      123,
+		"method":  "rpc_cancel",
+		"version": "2.0",
+		"params":  []interface{}{1},
+	}
+
+	pendingRequest := map[string]interface{}{
+		"id":      456,
+		"method":  "rpc_pending",
+		"version": "2.0",
+		"params":  []interface{}{},
+	}
+
+	ids := []interface{}{1, "qwerty", 2, 42}
+	for _, id := range ids {
+		request["id"] = id
+		if err := out.Encode(request); err != nil {
+			t.Fatal(err)
+		}
+	}
+	time.Sleep(time.Millisecond * 10)
+
+	testPending := func(expRunning, expCanceled int) {
+		if err := out.Encode(pendingRequest); err != nil {
+			t.Fatal(err)
+		}
+		var response struct {
+			Version string              `json:"jsonrpc"`
+			Id      interface{}         `json:"id,omitempty"`
+			Result  []jsonPendingStatus `json:"result"`
+		}
+
+		if err := in.Decode(&response); err != nil {
+			t.Fatal(err)
+		}
+		p := response.Result
+		r := 0
+		c := 0
+		for _, status := range p {
+			if status.Status == "running" {
+				r++
+			} else {
+				if status.Status == "canceled" {
+					c++
+				} else {
+					t.Errorf("unknown pending request status: %s", status.Status)
+				}
+			}
+		}
+		if r != expRunning || c != expCanceled {
+			t.Errorf("incorrect pending status results (got %d running / %d canceled, expected %d running / %d canceled)", r, c, expRunning, expCanceled)
+		}
+		time.Sleep(time.Millisecond * 10)
+	}
+
+	cancel := func(id interface{}, expMsg string) {
+		// send cancel request
+		cancelRequest["params"] = []interface{}{id}
+		if err := out.Encode(cancelRequest); err != nil {
+			t.Fatal(err)
+		}
+		var response JSONErrResponse
+		if err := in.Decode(&response); err != nil {
+			t.Fatal(err)
+		}
+		msg := response.Error.Message
+		if msg != expMsg {
+			t.Errorf("wrong error message from cancel (got \"%s\", expected \"%s\")", msg, expMsg)
+		}
+		time.Sleep(time.Millisecond * 10)
+	}
+
+	wait := func(count int) {
+		for i := 0; i < count; i++ {
+			var response JSONErrResponse
+			if err := in.Decode(&response); err != nil {
+				t.Fatal(err)
+			}
+			msg := response.Error.Message
+			expMsg := "context canceled"
+			if msg != expMsg {
+				t.Errorf("wrong error message from canceled request (got \"%s\", expected \"%s\")", msg, expMsg)
+			}
+		}
+		time.Sleep(time.Millisecond * 10)
+	}
+
+	testPending(4, 0)
+	cancel(ids[0], "")
+	testPending(3, 1)
+	cancel(ids[1], "")
+	testPending(2, 2)
+	cancel("twrtret", "pending request not found")
+	testPending(2, 2)
+	wait(2)
+	testPending(2, 0)
+	cancel("", "")
+	testPending(0, 2)
+	wait(2)
+	testPending(0, 0)
+
+	clientConn.Close()
+}

rpc/json.go

@@ -79,6 +79,12 @@ type jsonNotification struct {
 	Params  jsonSubscription `json:"params"`
 }
 
+// JSON-RPC pending request status
+type jsonPendingStatus struct {
+	Id     interface{} `json:"id"`
+	Status string      `json:"status"`
+}
+
 // jsonCodec reads and writes JSON-RPC messages to the underlying connection. It
 // also has support for parsing arguments and serializing (result) objects.
 type jsonCodec struct {
@@ -129,20 +135,21 @@ func (c *jsonCodec) ReadRequestHeaders() ([]rpcRequest, bool, RPCError) {
 	return parseRequest(incomingMsg)
 }
 
-// checkReqId returns an error when the given reqId isn't valid for RPC method calls.
+// parseReqId converts reqId to the appropriate format or returns an error
+// when the given reqId isn't valid for RPC method calls.
 // valid id's are strings, numbers or null
-func checkReqId(reqId json.RawMessage) error {
+func parseReqId(reqId json.RawMessage) (interface{}, error) {
 	if len(reqId) == 0 {
-		return fmt.Errorf("missing request id")
+		return nil, fmt.Errorf("missing request id")
 	}
-	if _, err := strconv.ParseFloat(string(reqId), 64); err == nil {
-		return nil
+	if f, err := strconv.ParseFloat(string(reqId), 64); err == nil {
+		return f, nil
 	}
 	var str string
 	if err := json.Unmarshal(reqId, &str); err == nil {
-		return nil
+		return str, nil
 	}
-	return fmt.Errorf("invalid request id")
+	return nil, fmt.Errorf("invalid request id")
 }
 
 // parseRequest will parse a single request from the given RawMessage. It will return
@@ -154,13 +161,14 @@ func parseRequest(incomingMsg json.RawMessage) ([]rpcRequest, bool, RPCError) {
 		return nil, false, &invalidMessageError{err.Error()}
 	}
 
-	if err := checkReqId(in.Id); err != nil {
+	id, err := parseReqId(in.Id)
+	if err != nil {
 		return nil, false, &invalidMessageError{err.Error()}
 	}
 
 	// subscribe are special, they will always use `subscribeMethod` as first param in the payload
 	if in.Method == subscribeMethod {
-		reqs := []rpcRequest{rpcRequest{id: &in.Id, isPubSub: true}}
+		reqs := []rpcRequest{rpcRequest{id: id, isPubSub: true}}
 		if len(in.Payload) > 0 {
 			// first param must be subscription name
 			var subscribeMethod [1]string
@@ -178,7 +186,7 @@ func parseRequest(incomingMsg json.RawMessage) ([]rpcRequest, bool, RPCError) {
 	}
 
 	if in.Method == unsubscribeMethod {
-		return []rpcRequest{rpcRequest{id: &in.Id, isPubSub: true,
+		return []rpcRequest{rpcRequest{id: id, isPubSub: true,
 			method: unsubscribeMethod, params: in.Payload}}, false, nil
 	}
 
@@ -189,10 +197,10 @@ func parseRequest(incomingMsg json.RawMessage) ([]rpcRequest, bool, RPCError) {
 	}
 
 	if len(in.Payload) == 0 {
-		return []rpcRequest{rpcRequest{service: elems[0], method: elems[1], id: &in.Id}}, false, nil
+		return []rpcRequest{rpcRequest{service: elems[0], method: elems[1], id: id}}, false, nil
 	}
 
-	return []rpcRequest{rpcRequest{service: elems[0], method: elems[1], id: &in.Id, params: in.Payload}}, false, nil
+	return []rpcRequest{rpcRequest{service: elems[0], method: elems[1], id: id, params: in.Payload}}, false, nil
 }
 
 // parseBatchRequest will parse a batch request into a collection of requests from the given RawMessage, an indication
@@ -205,12 +213,11 @@ func parseBatchRequest(incomingMsg json.RawMessage) ([]rpcRequest, bool, RPCErro
 
 	requests := make([]rpcRequest, len(in))
 	for i, r := range in {
-		if err := checkReqId(r.Id); err != nil {
+		id, err := parseReqId(r.Id)
+		if err != nil {
 			return nil, false, &invalidMessageError{err.Error()}
 		}
 
-		id := &in[i].Id
-
 		// subscribe are special, they will always use `subscribeMethod` as first param in the payload
 		if r.Method == subscribeMethod {
 			requests[i] = rpcRequest{id: id, isPubSub: true}

rpc/json_test.go

@@ -21,7 +21,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"reflect"
-	"strconv"
 	"testing"
 )
 
@@ -69,16 +68,12 @@ func TestJSONRequestParsing(t *testing.T) {
 		t.Fatalf("Expected method 'Add' but got '%s'", requests[0].method)
 	}
 
-	if rawId, ok := requests[0].id.(*json.RawMessage); ok {
-		id, e := strconv.ParseInt(string(*rawId), 0, 64)
-		if e != nil {
-			t.Fatalf("%v", e)
-		}
+	if id, ok := requests[0].id.(float64); ok {
 		if id != 1234 {
-			t.Fatalf("Expected id 1234 but got %d", id)
+			t.Fatalf("Expected id 1234 but got %v", id)
 		}
 	} else {
-		t.Fatalf("invalid request, expected *json.RawMesage got %T", requests[0].id)
+		t.Fatalf("invalid request, expected float64 got %T", requests[0].id)
 	}
 
 	var arg int

rpc/server.go

@@ -17,6 +17,7 @@
 package rpc
 
 import (
+	"errors"
 	"fmt"
 	"reflect"
 	"runtime"
@@ -57,6 +58,7 @@ func NewServer() *Server {
 		subscriptions: make(subscriptionRegistry),
 		codecs:        set.New(),
 		run:           1,
+		pending:       make(serverPendingRequests),
 	}
 
 	// register a default service which will provide meta information about the RPC service such as the services and
@@ -82,6 +84,91 @@ func (s *RPCService) Modules() map[string]string {
 	return modules
 }
 
+var (
+	// ErrPendingUnsupported is returned when the connection doesn't support access to pending requests
+	ErrPendingUnsupported = errors.New("access to pending requests not supported")
+
+	// ErrPendingNotFound is returned when no pending request with the given id is found
+	ErrPendingNotFound = errors.New("pending request not found")
+)
+
+type (
+	codecKey struct{}
+	idKey    struct{}
+)
+
+// Cancel cancels a pending request initiated through the same codec as this one.
+// If id == nil, all such requests are canceled.
+func (s *RPCService) Cancel(ctx context.Context, id interface{}) error {
+	_, fl := id.(float64)
+	_, st := id.(string)
+	if !fl && !st {
+		return ErrPendingNotFound
+	}
+
+	s.server.pendingMu.Lock()
+	defer s.server.pendingMu.Unlock()
+
+	codec, _ := ctx.Value(codecKey{}).(ServerCodec)
+	if codec == nil {
+		return ErrPendingUnsupported
+	}
+	codecPending := s.server.pending[codec]
+	if codecPending == nil {
+		return ErrPendingUnsupported
+	}
+
+	if id == "" {
+		for _, req := range codecPending {
+			if !req.canceled {
+				req.cancelFn()
+				req.canceled = true
+			}
+		}
+		return nil
+	}
+
+	if req, ok := codecPending[id]; ok {
+		if !req.canceled {
+			req.cancelFn()
+			req.canceled = true
+		}
+		return nil
+	}
+	return ErrPendingNotFound
+}
+
+// Pending returns the status of all pending requests initiated through the same
+// codec as this one. The status is either "running" or "canceled".
+func (s *RPCService) Pending(ctx context.Context) ([]jsonPendingStatus, error) {
+	s.server.pendingMu.Lock()
+	defer s.server.pendingMu.Unlock()
+
+	codec, _ := ctx.Value(codecKey{}).(ServerCodec)
+	if codec == nil {
+		return nil, ErrPendingUnsupported
+	}
+	codecPending := s.server.pending[codec]
+	if codecPending == nil {
+		return nil, ErrPendingUnsupported
+	}
+
+	selfId := ctx.Value(idKey{})
+	var res []jsonPendingStatus
+	for id, req := range codecPending {
+		if id != selfId {
+			s := jsonPendingStatus{Id: id}
+			if req.canceled {
+				s.Status = "canceled"
+			} else {
+				s.Status = "running"
+			}
+			res = append(res, s)
+		}
+	}
+	return res, nil
+}
+
 // RegisterName will create an service for the given rcvr type under the given name. When no methods on the given rcvr
 // match the criteria to be either a RPC method or a subscription an error is returned. Otherwise a new service is
 // created and added to the service collection this server instance serves.
@@ -162,7 +249,7 @@ func (s *Server) serveRequest(codec ServerCodec, singleShot bool, options CodecO
 		return
 	}()
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(context.WithValue(context.Background(), codecKey{}, codec))
 	defer cancel()
 
 	// if the codec supports notification include a notifier that callbacks can use
@@ -224,7 +311,16 @@ func (s *Server) serveRequest(codec ServerCodec, singleShot bool, options CodecO
 // response back using the given codec. It will block until the codec is closed or the server is
 // stopped. In either case the codec is closed.
 func (s *Server) ServeCodec(codec ServerCodec, options CodecOption) {
-	defer codec.Close()
+	defer func() {
+		s.pendingMu.Lock()
+		delete(s.pending, codec)
+		s.pendingMu.Unlock()
+		codec.Close()
+	}()
+
+	s.pendingMu.Lock()
+	s.pending[codec] = make(codecPendingRequests)
+	s.pendingMu.Unlock()
 	s.serveRequest(codec, false, options)
 }
 
@@ -337,12 +433,35 @@ func (s *Server) handle(ctx context.Context, codec ServerCodec, req *serverReque
 
 // exec executes the given request and writes the result back using the codec.
 func (s *Server) exec(ctx context.Context, codec ServerCodec, req *serverRequest) {
+	lctx := ctx
+	s.pendingMu.Lock()
+	codecPending := s.pending[codec]
+	var preq *pendingRequest
+	if codecPending != nil {
+		var cancelFn func()
+		lctx, cancelFn = context.WithCancel(context.WithValue(ctx, idKey{}, req.id))
+		preq = &pendingRequest{cancelFn, false}
+		codecPending[req.id] = preq
+	}
+	s.pendingMu.Unlock()
+
+	if codecPending != nil {
+		defer func() {
+			s.pendingMu.Lock()
+			if !preq.canceled {
+				preq.cancelFn() // always cancel contexts in order to avoid leaks
+			}
+			delete(codecPending, req.id)
+			s.pendingMu.Unlock()
+		}()
+	}
+
 	var response interface{}
 	var callback func()
 	if req.err != nil {
 		response = codec.CreateErrorResponse(&req.id, req.err)
 	} else {
-		response, callback = s.handle(ctx, codec, req)
+		response, callback = s.handle(lctx, codec, req)
 	}
 
 	if err := codec.Write(response); err != nil {
@@ -359,6 +478,33 @@ func (s *Server) exec(ctx context.Context, codec ServerCodec, req *serverRequest
 // execBatch executes the given requests and writes the result back using the codec.
 // It will only write the response back when the last request is processed.
 func (s *Server) execBatch(ctx context.Context, codec ServerCodec, requests []*serverRequest) {
+	lctx := ctx
+	s.pendingMu.Lock()
+	codecPending := s.pending[codec]
+	var preq *pendingRequest
+	if codecPending != nil {
+		var cancelFn func()
+		lctx, cancelFn = context.WithCancel(ctx)
+		preq = &pendingRequest{cancelFn, false}
+		for _, req := range requests {
+			codecPending[req.id] = preq
+		}
+	}
+	s.pendingMu.Unlock()
+
+	if codecPending != nil {
+		defer func() {
+			s.pendingMu.Lock()
+			if !preq.canceled {
+				preq.cancelFn() // always cancel contexts in order to avoid leaks
+			}
+			for _, req := range requests {
+				delete(codecPending, req.id)
+			}
+			s.pendingMu.Unlock()
+		}()
+	}
+
 	responses := make([]interface{}, len(requests))
 	var callbacks []func()
 	for i, req := range requests {
@@ -366,7 +512,8 @@ func (s *Server) execBatch(ctx context.Context, codec ServerCodec, requests []*s
 			responses[i] = codec.CreateErrorResponse(&req.id, req.err)
 		} else {
 			var callback func()
-			if responses[i], callback = s.handle(ctx, codec, req); callback != nil {
+			cctx := context.WithValue(lctx, idKey{}, req.id)
+			if responses[i], callback = s.handle(cctx, codec, req); callback != nil {
 				callbacks = append(callbacks, callback)
 			}
 		}

rpc/types.go

@@ -70,6 +70,19 @@ type callbacks map[string]*callback            // collection of RPC callbacks
 type subscriptions map[string]*callback        // collection of subscription callbacks
 type subscriptionRegistry map[string]*callback // collection of subscription callbacks
 
+// pendingRequest contains status and cancel info about currently running requests
+type pendingRequest struct {
+	cancelFn func()
+	canceled bool
+}
+
+// codecPendingRequests stores all pending requests initiated from the same codec
+// keys are request IDs (either float64 or string)
+type codecPendingRequests map[interface{}]*pendingRequest
+
+// serverPendingRequests stores all pending requests handled by a server
+type serverPendingRequests map[ServerCodec]codecPendingRequests
+
 // Server represents a RPC server
 type Server struct {
 	services       serviceRegistry
@@ -79,6 +92,9 @@ type Server struct {
 	run      int32
 	codecsMu sync.Mutex
 	codecs   *set.Set
+
+	pending   serverPendingRequests
+	pendingMu sync.Mutex
 }
 
 // rpcRequest represents a raw incoming RPC request