diff --git a/rpc/cancel_test.go b/rpc/cancel_test.go new file mode 100644 index 0000000000..29bd280004 --- /dev/null +++ b/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 . + +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() +} diff --git a/rpc/json.go b/rpc/json.go index 151ed546e7..88528405cb 100644 --- a/rpc/json.go +++ b/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} diff --git a/rpc/json_test.go b/rpc/json_test.go index 5048d2f7a0..1762218648 100644 --- a/rpc/json_test.go +++ b/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 diff --git a/rpc/server.go b/rpc/server.go index 69f3271e8e..b6ef9d03e9 100644 --- a/rpc/server.go +++ b/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) } } diff --git a/rpc/types.go b/rpc/types.go index a1f36fbd26..f56c2e0dd5 100644 --- a/rpc/types.go +++ b/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