Fix context cancelled error after client close

No longer set client error from the receiver goroutine when the client
context is closed. The client context is cancelled when the client is
closed, so the receiver goroutine can stop handling messages. The client
error will be set the next dispatch.

Signed-off-by: Austin Vazquez <[email protected]>

client.go

@@ -264,6 +264,7 @@ func (c *Client) run() {
 		for {
 			select {
 			case <-c.ctx.Done():
+				// The client has been closed, so the sender goroutine can now be stopped.
 				return
 			case call := <-c.calls:
 				id := streamID
@@ -288,7 +289,7 @@ func (c *Client) run() {
 		for {
 			select {
 			case <-c.ctx.Done():
-				c.setError(c.ctx.Err())
+				// The client has been closed, so the receiver goroutine can now be stopped.
 				return
 			default:
 				mh, p, err := c.channel.recv()

client_test.go

@@ -68,3 +68,35 @@ func TestUserOnCloseWait(t *testing.T) {
 		t.Fatalf("expected error nil , but got %v", err)
 	}
 }
+
+func TestClientReturnsErrClosedAfterClose(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	server := mustServer(t)(NewServer())
+	testService := &testingServer{}
+	addr, listener := newTestListener(t)
+
+	registerTestingService(server, testService)
+	go server.Serve(ctx, listener)
+
+	t.Cleanup(func() {
+		server.Shutdown(ctx)
+		listener.Close()
+		cancel()
+	})
+
+	client, cleanup := newTestClient(t, addr)
+	testClient := newTestingClient(client)
+
+	t.Cleanup(func() {
+		cleanup()
+	})
+
+	err := client.Close()
+	if err != nil {
+		t.Errorf("Expected nil error on client close, received %v", err)
+	}
+
+	if _, err := testClient.Test(ctx, &testPayload{}); err != ErrClosed {
+		t.Errorf("Expected ErrClosed after connection has been closed, got %v", err)
+	}
+}