1 files changed, 31 insertions, 26 deletions
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())
 }