Merge pull request #43 from jannson/fix-keepalive-block

fix: writeFrame may block forever in keepalive function

Author: xtaci

session.go

@@ -289,7 +289,7 @@ func (s *Session) keepalive() {
 	for {
 		select {
 		case <-tickerPing.C:
-			s.writeFrame(newFrame(cmdNOP, 0))
+			s.writeFrameInternal(newFrame(cmdNOP, 0), tickerPing.C)
 			s.notifyBucket() // force a signal to the recvLoop
 		case <-tickerTimeout.C:
 			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
@@ -335,6 +335,11 @@ func (s *Session) sendLoop() {
 // writeFrame writes the frame to the underlying connection
 // and returns the number of bytes written if successful
 func (s *Session) writeFrame(f Frame) (n int, err error) {
+	return s.writeFrameInternal(f, nil)
+}
+
+// internal writeFrame version to support deadline used in keepalive
+func (s *Session) writeFrameInternal(f Frame, deadline <-chan time.Time) (int, error) {
 	req := writeRequest{
 		frame:  f,
 		result: make(chan writeResult, 1),
@@ -343,8 +348,16 @@ func (s *Session) writeFrame(f Frame) (n int, err error) {
 	case <-s.die:
 		return 0, errors.New(errBrokenPipe)
 	case s.writes <- req:
+	case <-deadline:
+		return 0, errTimeout
 	}
 
-	result := <-req.result
-	return result.n, result.err
+	select {
+	case result := <-req.result:
+		return result.n, result.err
+	case <-deadline:
+		return 0, errTimeout
+	case <-s.die:
+		return 0, errors.New(errBrokenPipe)
+	}
 }

session_test.go

@@ -304,6 +304,44 @@ func TestKeepAliveTimeout(t *testing.T) {
 	}
 }
 
+type blockWriteConn struct {
+	net.Conn
+}
+
+func (c *blockWriteConn) Write(b []byte) (n int, err error) {
+	forever := time.Hour * 24
+	time.Sleep(forever)
+	return c.Conn.Write(b)
+}
+
+func TestKeepAliveBlockWriteTimeout(t *testing.T) {
+	ln, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	go func() {
+		ln.Accept()
+	}()
+
+	cli, err := net.Dial("tcp", ln.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cli.Close()
+	//when writeFrame block, keepalive in old version never timeout
+	blockWriteCli := &blockWriteConn{cli}
+
+	config := DefaultConfig()
+	config.KeepAliveInterval = time.Second
+	config.KeepAliveTimeout = 2 * time.Second
+	session, _ := Client(blockWriteCli, config)
+	time.Sleep(3 * time.Second)
+	if !session.IsClosed() {
+		t.Fatal("keepalive-timeout failed")
+	}
+}
+
 func TestServerEcho(t *testing.T) {
 	ln, err := net.Listen("tcp", "localhost:0")
 	if err != nil {
@@ -541,6 +579,95 @@ func TestRandomFrame(t *testing.T) {
 
 	session.conn.Write(buf)
 	cli.Close()
+
+	// writeFrame after die
+	cli, err = net.Dial("tcp", addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	session, _ = Client(cli, nil)
+	//close first
+	session.Close()
+	for i := 0; i < 100; i++ {
+		f := newFrame(byte(rand.Uint32()), rand.Uint32())
+		session.writeFrame(f)
+	}
+}
+
+func TestWriteFrameInternal(t *testing.T) {
+	addr, stop, cli, err := setupServer(t)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer stop()
+	// pure random
+	session, _ := Client(cli, nil)
+	for i := 0; i < 100; i++ {
+		rnd := make([]byte, rand.Uint32()%1024)
+		io.ReadFull(crand.Reader, rnd)
+		session.conn.Write(rnd)
+	}
+	cli.Close()
+
+	// writeFrame after die
+	cli, err = net.Dial("tcp", addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	session, _ = Client(cli, nil)
+	//close first
+	session.Close()
+	for i := 0; i < 100; i++ {
+		f := newFrame(byte(rand.Uint32()), rand.Uint32())
+		session.writeFrameInternal(f, time.After(session.config.KeepAliveTimeout))
+	}
+
+	// random cmds
+	cli, err = net.Dial("tcp", addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	allcmds := []byte{cmdSYN, cmdFIN, cmdPSH, cmdNOP}
+	session, _ = Client(cli, nil)
+	for i := 0; i < 100; i++ {
+		f := newFrame(allcmds[rand.Int()%len(allcmds)], rand.Uint32())
+		session.writeFrameInternal(f, time.After(session.config.KeepAliveTimeout))
+	}
+	//deadline occur
+	{
+		c := make(chan time.Time)
+		close(c)
+		f := newFrame(allcmds[rand.Int()%len(allcmds)], rand.Uint32())
+		_, err := session.writeFrameInternal(f, c)
+		if err != errTimeout {
+			t.Fatal("write frame with deadline failed", err)
+		}
+	}
+	cli.Close()
+
+	{
+		cli, err = net.Dial("tcp", addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		config := DefaultConfig()
+		config.KeepAliveInterval = time.Second
+		config.KeepAliveTimeout = 2 * time.Second
+		session, _ = Client(&blockWriteConn{cli}, config)
+		f := newFrame(byte(rand.Uint32()), rand.Uint32())
+		c := make(chan time.Time)
+		go func() {
+			//die first, deadline second, better for coverage
+			time.Sleep(time.Second)
+			session.Close()
+			time.Sleep(time.Second)
+			close(c)
+		}()
+		_, err = session.writeFrameInternal(f, c)
+		if err.Error() != errBrokenPipe {
+			t.Fatal("write frame with deadline failed", err)
+		}
+	}
 }
 
 func TestReadDeadline(t *testing.T) {