zmq4: router recv returns error when closing context

Fixes #58.

Author: speakyourcode

router.go

@@ -146,6 +146,7 @@ func (q *routerQReader) read(ctx context.Context, msg *Msg) error {
 	q.sem.lock()
 	select {
 	case <-ctx.Done():
+		return ctx.Err()
 	case *msg = <-q.c:
 	}
 	return msg.err

zmq4_routerdealer_test.go

@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"net"
 	"reflect"
 	"sync"
 	"testing"
@@ -253,3 +254,35 @@ func TestRouterWithNoDealer(t *testing.T) {
 		t.Fatalf("could not close router: %+v", err)
 	}
 }
+
+func TestRouterClose(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	router := zmq4.NewRouter(ctx)
+	err := router.Listen("tcp://*:*")
+	if err != nil {
+		t.Fatalf("router could not listen: %+v", err)
+	}
+	_, port, _ := net.SplitHostPort(router.Addr().String())
+	dealer := zmq4.NewDealer(ctx)
+	err = dealer.Dial("tcp://*:" + port)
+	if err != nil {
+		t.Fatalf("dealer could not dial: %+v", err)
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		_, err := router.Recv()
+		if err == nil {
+			t.Errorf("expected error: context canceled")
+		}
+	}()
+
+	err = router.Close()
+	if err != nil {
+		t.Fatalf("could not close router: %+v", err)
+	}
+	wg.Wait()
+}