feat(rpc2): transport keepalive with ping pong

. #22478859

Signed-off-by: slasher <shenjie1@oppo.com>

blobstore/common/rpc2/transport/frame.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"sync"
+	"sync/atomic"
 )
 
 const ( // cmds
@@ -12,7 +13,8 @@ const ( // cmds
 	cmdSYN byte = iota // stream open
 	cmdFIN             // stream close, a.k.a EOF mark
 	cmdPSH             // data push
-	cmdNOP             // no operation
+	cmdPIN             // client send  keepalive
+	cmdPON             // server reply keepalive
 
 	// protocol version 2 extra commands
 	// notify bytes consumed by remote peer-end
@@ -48,15 +50,15 @@ const (
 type rawHeader [headerSize]byte
 
 func (h rawHeader) Version() byte {
-	return h[0] >> 4
+	return h[3] >> 4 // little endian
 }
 
 func (h rawHeader) Cmd() byte {
-	return h[0] & 0x0f
+	return h[3] & 0x0f // little endian
 }
 
 func (h rawHeader) Length() uint32 {
-	return binary.LittleEndian.Uint32(h[1:]) & 0xffffff
+	return binary.LittleEndian.Uint32(h[:]) & 0xffffff
 }
 
 func (h rawHeader) StreamID() uint32 {
@@ -87,12 +89,21 @@ type FrameWrite struct {
 	off  int
 	data []byte // with frame header
 
+	done   uint32 // 0 = new, 1 = locked, 2 == closed
 	once   sync.Once
 	closer interface {
 		Free([]byte) error
 	}
 }
 
+func (f *FrameWrite) tryLock() bool {
+	return atomic.CompareAndSwapUint32(&f.done, 0, 1)
+}
+
+func (f *FrameWrite) unlock() {
+	atomic.CompareAndSwapUint32(&f.done, 1, 0)
+}
+
 func (f *FrameWrite) Write(p []byte) (int, error) {
 	if f.off >= len(f.data) {
 		return 0, nil
@@ -126,6 +137,8 @@ func (f *FrameWrite) Len() int {
 
 func (f *FrameWrite) Close() (err error) {
 	f.once.Do(func() {
+		for !atomic.CompareAndSwapUint32(&f.done, 0, 2) {
+		}
 		data := f.data
 		f.data = nil
 		f.off = 0

blobstore/common/rpc2/transport/session.go

@@ -78,6 +78,7 @@ type Session struct {
 	chAccepts chan *Stream
 
 	dataReady int32 // flag data has arrived
+	pingpong  chan struct{}
 
 	goAway int32 // flag id exhausted
 
@@ -100,6 +101,7 @@ func newSession(config *Config, conn io.ReadWriteCloser, client bool) *Session {
 	s.chAccepts = make(chan *Stream, defaultAcceptBacklog)
 	s.bucket = int32(config.MaxReceiveBuffer)
 	s.bucketNotify = make(chan struct{}, 1)
+	s.pingpong = make(chan struct{}, 1)
 	s.ctrl = make(chan writeRequest, 1)
 	s.writes = make(chan writeRequest)
 	s.resultChPool = sync.Pool{New: func() interface{} {
@@ -359,7 +361,13 @@ func (s *Session) recvLoop() {
 			}
 			sid := hdr.StreamID()
 			switch hdr.Cmd() {
-			case cmdNOP:
+			case cmdPIN:
+				s.pong()
+			case cmdPON:
+				select {
+				case s.pingpong <- struct{}{}:
+				default:
+				}
 			case cmdSYN:
 				s.streamLock.Lock()
 				if _, ok := s.streams[sid]; !ok {
@@ -426,6 +434,25 @@ func (s *Session) recvLoop() {
 	}
 }
 
+func (s *Session) ping(ticker *time.Ticker) {
+	deadline := writeDealine{
+		time: time.Now().Add(s.config.KeepAliveInterval),
+		wait: ticker.C,
+	}
+	frame, err := s.newFrameWrite(cmdPIN, 0, 0)
+	if err == nil {
+		s.writeFrameInternal(frame, deadline, CLSCTRL)
+		frame.Close()
+	}
+}
+
+func (s *Session) pong() {
+	frame, err := s.newFrameWrite(cmdPON, 0, 0)
+	if err == nil {
+		s.writeFrame(frame)
+	}
+}
+
 func (s *Session) keepalive(client bool) {
 	tickerTimeout := time.NewTicker(s.config.KeepAliveTimeout)
 	defer tickerTimeout.Stop()
@@ -453,17 +480,7 @@ func (s *Session) keepalive(client bool) {
 	for {
 		select {
 		case <-tickerPing.C:
-			deadline := writeDealine{
-				time: time.Now().Add(s.config.KeepAliveInterval),
-				wait: tickerPing.C,
-			}
-
-			frame, err := s.newFrameWrite(cmdNOP, 0, 0)
-			if err == nil {
-				if _, err = s.writeFrameInternal(frame, deadline, CLSCTRL); err == nil {
-					alive = true
-				}
-			}
+			s.ping(tickerPing)
 			s.notifyBucket() // force a signal to the recvLoop
 		case <-tickerTimeout.C:
 			if !atomic.CompareAndSwapInt32(&s.dataReady, 1, 0) {
@@ -475,6 +492,8 @@ func (s *Session) keepalive(client bool) {
 				}
 			}
 			alive = false
+		case <-s.pingpong:
+			alive = true
 		case <-s.die:
 			return
 		}
@@ -483,6 +502,7 @@ func (s *Session) keepalive(client bool) {
 
 func (s *Session) sendLoop() {
 	var buf []byte
+	var first uint32
 	var n, nn int
 	var err error
 	var request writeRequest
@@ -497,6 +517,8 @@ func (s *Session) sendLoop() {
 	for {
 		select {
 		case request = <-s.ctrl:
+		case <-s.die:
+			return
 		default:
 			select {
 			case <-s.die:
@@ -506,9 +528,15 @@ func (s *Session) sendLoop() {
 			}
 		}
 
+		if !request.frame.tryLock() { // closed
+			continue
+		}
+
 		buf = request.frame.data[:request.frame.off]
-		buf[0] = (request.frame.ver << 4) | (request.frame.cmd & 0x0f)
-		binary.LittleEndian.PutUint32(buf[1:], uint32(request.frame.Len()))
+		first = (uint32(request.frame.ver) << 28) |
+			(uint32(request.frame.cmd&0x0f) << 24) |
+			(uint32(request.frame.Len()) & 0xffffff)
+		binary.LittleEndian.PutUint32(buf[:], first)
 		binary.LittleEndian.PutUint32(buf[4:], request.frame.sid)
 
 		setWriteDeadline(request.deadline)
@@ -525,6 +553,7 @@ func (s *Session) sendLoop() {
 		if n -= headerSize; n < 0 {
 			n = 0
 		}
+		request.frame.unlock()
 
 		result := writeResult{
 			n:   n,
@@ -546,6 +575,7 @@ func (s *Session) sendLoop() {
 func (s *Session) writeFrame(f *FrameWrite) (n int, err error) {
 	timer := time.NewTimer(openCloseTimeout)
 	defer timer.Stop()
+	defer f.Close()
 	deadline := writeDealine{
 		time: time.Now().Add(openCloseTimeout),
 		wait: timer.C,

blobstore/common/rpc2/transport/session_test.go

@@ -825,7 +825,7 @@ func TestRandomFrame(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	allcmds := []byte{cmdSYN, cmdFIN, cmdPSH, cmdNOP}
+	allcmds := []byte{cmdSYN, cmdFIN, cmdPSH, cmdPIN, cmdPON}
 	session, _ = Client(cli, nil)
 	for i := 0; i < 100; i++ {
 		f, _ := session.newFrameWrite(allcmds[rand.Int()%len(allcmds)], rand.Uint32(), 0)
@@ -871,8 +871,10 @@ func TestRandomFrame(t *testing.T) {
 	f.Write(rnd)
 
 	buf := make([]byte, headerSize+len(f.data))
-	buf[0] = (f.ver << 4) | (f.cmd & 0x0f)
-	binary.LittleEndian.PutUint32(buf[1:], uint32(len(rnd)+1)) // incorrect size
+
+	first := (uint32(f.ver) << 28) | (uint32(f.cmd&0x0f) << 24) |
+		(uint32(f.Len()+1) & 0xffffff)
+	binary.LittleEndian.PutUint32(buf[:], first) // incorrect size
 	binary.LittleEndian.PutUint32(buf[4:], f.sid)
 	copy(buf[headerSize:], rnd)
 
@@ -930,7 +932,7 @@ func TestWriteFrameInternal(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	allcmds := []byte{cmdSYN, cmdFIN, cmdPSH, cmdNOP}
+	allcmds := []byte{cmdSYN, cmdFIN, cmdPSH, cmdPIN, cmdPON}
 	session, _ = Client(cli, nil)
 	for i := 0; i < 100; i++ {
 		f, _ := session.newFrameWrite(allcmds[rand.Int()%len(allcmds)], rand.Uint32(), 0)

blobstore/common/rpc2/transport/stream.go

@@ -280,7 +280,7 @@ func (s *Stream) writeFrameV2(frame *FrameWrite, deadline writeDealine) (n int,
 		}
 
 		win := int32(atomic.LoadUint32(&s.peerWindow)) - inflight
-		if win >= int32(frame.Len()) {
+		if win >= int32(frame.Len()) || s.numWritten == 0 {
 			sent, err := s.sess.writeFrameInternal(frame, deadline, CLSDATA)
 			s.numWritten += uint32(sent)
 			return sent, err