rpc: fix goroutine leaks in subscription tests

This improves the fixes sent previously by refactoring waitForMessages
to exit when an error happens. This avoids the need for the error
channel.

rpc/subscription_test.go

@@ -60,11 +60,8 @@ func TestSubscriptions(t *testing.T) {
 		successes              = make(chan subConfirmation)
 		notifications          = make(chan subscriptionResult)
 		errors                 = make(chan error, subCount*notificationCount+1)
-		stop                   = make(chan struct{})
 	)
 
-	defer close(stop)
-
 	// setup and start server
 	for _, namespace := range namespaces {
 		if err := server.RegisterName(namespace, service); err != nil {
@@ -75,7 +72,7 @@ func TestSubscriptions(t *testing.T) {
 	defer server.Stop()
 
 	// wait for message and write them to the given channels
-	go waitForMessages(in, successes, notifications, errors, stop)
+	go waitForMessages(in, successes, notifications, errors)
 
 	// create subscriptions one by one
 	for i, namespace := range namespaces {
@@ -128,25 +125,26 @@ func TestSubscriptions(t *testing.T) {
 
 // This test checks that unsubscribing works.
 func TestServerUnsubscribe(t *testing.T) {
+	p1, p2 := net.Pipe()
+	defer p2.Close()
+
 	// Start the server.
 	server := newTestServer()
-	service := &notificationTestService{unsubscribed: make(chan string)}
+	service := &notificationTestService{unsubscribed: make(chan string, 1)}
 	server.RegisterName("nftest2", service)
-	p1, p2 := net.Pipe()
 	go server.ServeCodec(NewCodec(p1), 0)
 
-	p2.SetDeadline(time.Now().Add(10 * time.Second))
-
 	// Subscribe.
+	p2.SetDeadline(time.Now().Add(10 * time.Second))
 	p2.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"nftest2_subscribe","params":["someSubscription",0,10]}`))
 
 	// Handle received messages.
-	resps := make(chan subConfirmation)
-	notifications := make(chan subscriptionResult)
-	errors := make(chan error)
-	stop := make(chan struct{})
-	defer close(stop)
-	go waitForMessages(json.NewDecoder(p2), resps, notifications, errors, stop)
+	var (
+		resps         = make(chan subConfirmation)
+		notifications = make(chan subscriptionResult)
+		errors        = make(chan error, 1)
+	)
+	go waitForMessages(json.NewDecoder(p2), resps, notifications, errors)
 
 	// Receive the subscription ID.
 	var sub subConfirmation
@@ -178,38 +176,45 @@ type subConfirmation struct {
 	subid ID
 }
 
-func waitForMessages(in *json.Decoder, successes chan subConfirmation, notifications chan subscriptionResult, errors chan error, stop chan struct{}) {
+// waitForMessages reads RPC messages from 'in' and dispatches them into the given channels.
+// It stops if there is an error.
+func waitForMessages(in *json.Decoder, successes chan subConfirmation, notifications chan subscriptionResult, errors chan error) {
 	for {
+		resp, notification, err := readAndValidateMessage(in)
+		if err != nil {
+			errors <- err
+			return
+		} else if resp != nil {
+			successes <- *resp
+		} else {
+			notifications <- *notification
+		}
+	}
+}
+
+func readAndValidateMessage(in *json.Decoder) (*subConfirmation, *subscriptionResult, error) {
 	var msg jsonrpcMessage
 	if err := in.Decode(&msg); err != nil {
-			errors <- fmt.Errorf("decode error: %v", err)
-			return
+		return nil, nil, fmt.Errorf("decode error: %v", err)
 	}
 	switch {
 	case msg.isNotification():
 		var res subscriptionResult
 		if err := json.Unmarshal(msg.Params, &res); err != nil {
-				errors <- fmt.Errorf("invalid subscription result: %v", err)
-			} else {
-				notifications <- res
+			return nil, nil, fmt.Errorf("invalid subscription result: %v", err)
 		}
+		return nil, &res, nil
 	case msg.isResponse():
 		var c subConfirmation
 		if msg.Error != nil {
-				errors <- msg.Error
+			return nil, nil, msg.Error
 		} else if err := json.Unmarshal(msg.Result, &c.subid); err != nil {
-				errors <- fmt.Errorf("invalid response: %v", err)
+			return nil, nil, fmt.Errorf("invalid response: %v", err)
 		} else {
 			json.Unmarshal(msg.ID, &c.reqid)
-				select {
-				case successes <- c:
-				case <-stop:
-					return
-				}
+			return &c, nil, nil
 		}
 	default:
-			errors <- fmt.Errorf("unrecognized message: %v", msg)
-			return
-		}
+		return nil, nil, fmt.Errorf("unrecognized message: %v", msg)
 	}
 }