rpc: add context argument to EthSubscribe

It's inconsistent not to pass it and most callers will
work with contexts anyway.

3 files changed, 38 insertions, 33 deletions

diff --git a/rpc/client.go b/rpc/client.go
index 6846e1dda..34a3b7831 100644
--- a/rpc/client.go
+++ b/rpc/client.go
@@ -354,12 +354,14 @@ func (c *Client) BatchCallContext(ctx context.Context, b []BatchElem) error {
 // sent to the given channel. The element type of the channel must match the
 // expected type of content returned by the subscription.
 //
+// The context argument cancels the RPC request that sets up the subscription but has no
+// effect on the subscription after EthSubscribe has returned.
 //
 // Slow subscribers will be dropped eventually. Client buffers up to 8000 notifications
 // before considering the subscriber dead. The subscription Err channel will receive
 // ErrSubscriptionQueueOverflow. Use a sufficiently large buffer on the channel or ensure
 // that the channel usually has at least one reader to prevent this issue.
-func (c *Client) EthSubscribe(channel interface{}, args ...interface{}) (*ClientSubscription, error) {
+func (c *Client) EthSubscribe(ctx context.Context, channel interface{}, args ...interface{}) (*ClientSubscription, error) {
 	// Check type of channel first.
 	chanVal := reflect.ValueOf(channel)
 	if chanVal.Kind() != reflect.Chan || chanVal.Type().ChanDir()&reflect.SendDir == 0 {
@@ -381,8 +383,6 @@ func (c *Client) EthSubscribe(channel interface{}, args ...interface{}) (*Client
 		resp: make(chan *jsonrpcMessage),
 		sub:  newClientSubscription(c, chanVal),
 	}
-	ctx, cancel := context.WithTimeout(context.Background(), subscribeTimeout)
-	defer cancel()
 
 	// Send the subscription request.
 	// The arrival and validity of the response is signaled on sub.quit.
diff --git a/rpc/client_example_test.go b/rpc/client_example_test.go
index 84b4b67bb..3462b3685 100644
--- a/rpc/client_example_test.go
+++ b/rpc/client_example_test.go
@@ -22,6 +22,7 @@ import (
 	"time"
 
 	"github.com/ethereum/go-ethereum/rpc"
+	"golang.org/x/net/context"
 )
 
 // In this example, our client whishes to track the latest 'block number'
@@ -41,7 +42,16 @@ func ExampleClientSubscription() {
 	// Connect the client.
 	client, _ := rpc.Dial("ws://127.0.0.1:8485")
 	subch := make(chan Block)
-	go subscribeBlocks(client, subch)
+
+	// Ensure that subch receives the latest block.
+	go func() {
+		for i := 0; ; i++ {
+			if i > 0 {
+				time.Sleep(2 * time.Second)
+			}
+			subscribeBlocks(client, subch)
+		}
+	}()
 
 	// Print events from the subscription as they arrive.
 	for block := range subch {
@@ -52,32 +62,27 @@ func ExampleClientSubscription() {
 // subscribeBlocks runs in its own goroutine and maintains
 // a subscription for new blocks.
 func subscribeBlocks(client *rpc.Client, subch chan Block) {
-	for i := 0; ; i++ {
-		if i > 0 {
-			time.Sleep(2 * time.Second)
-		}
-
-		// Subscribe to new blocks.
-		sub, err := client.EthSubscribe(subch, "newBlocks")
-		if err == rpc.ErrClientQuit {
-			return // Stop reconnecting if the client was closed.
-		} else if err != nil {
-			fmt.Println("subscribe error:", err)
-			continue
-		}
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
 
-		// The connection is established now.
-		// Update the channel with the current block.
-		var lastBlock Block
-		if err := client.Call(&lastBlock, "eth_getBlockByNumber", "latest"); err != nil {
-			fmt.Println("can't get latest block:", err)
-			continue
-		}
-		subch <- lastBlock
+	// Subscribe to new blocks.
+	sub, err := client.EthSubscribe(ctx, subch, "newBlocks")
+	if err != nil {
+		fmt.Println("subscribe error:", err)
+		return
+	}
 
-		// The subscription will deliver events to the channel. Wait for the
-		// subscription to end for any reason, then loop around to re-establish
-		// the connection.
-		fmt.Println("connection lost: ", <-sub.Err())
+	// The connection is established now.
+	// Update the channel with the current block.
+	var lastBlock Block
+	if err := client.CallContext(ctx, &lastBlock, "eth_getBlockByNumber", "latest"); err != nil {
+		fmt.Println("can't get latest block:", err)
+		return
 	}
+	subch <- lastBlock
+
+	// The subscription will deliver events to the channel. Wait for the
+	// subscription to end for any reason, then loop around to re-establish
+	// the connection.
+	fmt.Println("connection lost: ", <-sub.Err())
 }
diff --git a/rpc/client_test.go b/rpc/client_test.go
index 424d7f5bc..476c8c6f3 100644
--- a/rpc/client_test.go
+++ b/rpc/client_test.go
@@ -215,7 +215,7 @@ func TestClientSubscribeInvalidArg(t *testing.T) {
 				t.Error(string(buf))
 			}
 		}()
-		client.EthSubscribe(arg, "foo_bar")
+		client.EthSubscribe(context.Background(), arg, "foo_bar")
 	}
 	check(true, nil)
 	check(true, 1)
@@ -233,7 +233,7 @@ func TestClientSubscribe(t *testing.T) {
 
 	nc := make(chan int)
 	count := 10
-	sub, err := client.EthSubscribe(nc, "someSubscription", count, 0)
+	sub, err := client.EthSubscribe(context.Background(), nc, "someSubscription", count, 0)
 	if err != nil {
 		t.Fatal("can't subscribe:", err)
 	}
@@ -275,7 +275,7 @@ func TestClientSubscribeClose(t *testing.T) {
 		err  error
 	)
 	go func() {
-		sub, err = client.EthSubscribe(nc, "hangSubscription", 999)
+		sub, err = client.EthSubscribe(context.Background(), nc, "hangSubscription", 999)
 		errc <- err
 	}()
 
@@ -311,7 +311,7 @@ func TestClientNotificationStorm(t *testing.T) {
 		// Subscribe on the server. It will start sending many notifications
 		// very quickly.
 		nc := make(chan int)
-		sub, err := client.EthSubscribe(nc, "someSubscription", count, 0)
+		sub, err := client.EthSubscribe(ctx, nc, "someSubscription", count, 0)
 		if err != nil {
 			t.Fatal("can't subscribe:", err)
 		}