Fix deadlock caused by non-responsive ZK servers

conn.go

@@ -929,17 +929,17 @@ func (c *Conn) queueRequest(opcode int32, req interface{}, res interface{}, recv
 }
 
 func (c *Conn) request(opcode int32, req interface{}, res interface{}, recvFunc func(*request, *responseHeader, error)) (int64, error) {
-	r := <-c.queueRequest(opcode, req, res, recvFunc)
+	recv := c.queueRequest(opcode, req, res, recvFunc)
 	select {
+	case r := <-recv:
+		return r.zxid, r.err
 	case <-c.shouldQuit:
 		// queueRequest() can be racy, double-check for the race here and avoid
 		// a potential data-race. otherwise the client of this func may try to
 		// access `res` fields concurrently w/ the async response processor.
 		// NOTE: callers of this func should check for (at least) ErrConnectionClosed
 		// and avoid accessing fields of the response object if such error is present.
 		return -1, ErrConnectionClosed
-	default:
-		return r.zxid, r.err
 	}
 }
 

conn_test.go

@@ -57,6 +57,36 @@ func TestRecurringReAuthHang(t *testing.T) {
 	}
 }
 
+func TestConcurrentReadAndClose(t *testing.T) {
+	WithListenServer(t, func(server string) {
+		conn, _, err := Connect([]string{server}, 15*time.Second)
+		if err != nil {
+			t.Fatalf("Failed to create Connection %s", err)
+		}
+
+		okChan := make(chan struct{})
+		var setErr error
+		go func() {
+			_, setErr = conn.Create("/test-path", []byte("test data"), 0, WorldACL(PermAll))
+			close(okChan)
+		}()
+
+		go func() {
+			time.Sleep(1 * time.Second)
+			conn.Close()
+		}()
+
+		select {
+		case <-okChan:
+			if setErr != ErrConnectionClosed {
+				t.Fatalf("unexpected error returned from Set %v", setErr)
+			}
+		case <-time.After(3 * time.Second):
+			t.Fatal("apparent deadlock!")
+		}
+	})
+}
+
 func TestDeadlockInClose(t *testing.T) {
 	c := &Conn{
 		shouldQuit:     make(chan struct{}),

tcp_server_test.go

@@ -0,0 +1,39 @@
+package zk
+
+import (
+	"fmt"
+	"math/rand"
+	"net"
+	"testing"
+	"time"
+)
+
+func WithListenServer(t *testing.T, test func(server string)) {
+	startPort := int(rand.Int31n(6000) + 10000)
+	server := fmt.Sprintf("localhost:%d", startPort)
+	l, err := net.Listen("tcp", server)
+	if err != nil {
+		t.Fatalf("Failed to start listen server: %v", err)
+	}
+	defer l.Close()
+
+	go func() {
+		for {
+			conn, err := l.Accept()
+			if err != nil {
+				fmt.Println("Failed to accept connection: ", err.Error())
+				continue
+			}
+
+			handleRequest(conn)
+		}
+	}()
+
+	test(server)
+}
+
+// Handles incoming requests.
+func handleRequest(conn net.Conn) {
+	time.Sleep(5 * time.Second)
+	conn.Close()
+}