Two way streams for upcoming locking enhancements (#19796)

internal/grid/muxserver.go

@@ -28,21 +28,20 @@ import (
 	xioutil "github.com/minio/minio/internal/ioutil"
 )
 
-const lastPingThreshold = 4 * clientPingInterval
-
 type muxServer struct {
-	ID               uint64
-	LastPing         int64
-	SendSeq, RecvSeq uint32
-	Resp             chan []byte
-	BaseFlags        Flags
-	ctx              context.Context
-	cancel           context.CancelFunc
-	inbound          chan []byte
-	parent           *Connection
-	sendMu           sync.Mutex
-	recvMu           sync.Mutex
-	outBlock         chan struct{}
+	ID                 uint64
+	LastPing           int64
+	SendSeq, RecvSeq   uint32
+	Resp               chan []byte
+	BaseFlags          Flags
+	ctx                context.Context
+	cancel             context.CancelFunc
+	inbound            chan []byte
+	parent             *Connection
+	sendMu             sync.Mutex
+	recvMu             sync.Mutex
+	outBlock           chan struct{}
+	clientPingInterval time.Duration
 }
 
 func newMuxStateless(ctx context.Context, msg message, c *Connection, handler StatelessHandler) *muxServer {
@@ -89,16 +88,17 @@ func newMuxStream(ctx context.Context, msg message, c *Connection, handler Strea
 	}
 
 	m := muxServer{
-		ID:        msg.MuxID,
-		RecvSeq:   msg.Seq + 1,
-		SendSeq:   msg.Seq,
-		ctx:       ctx,
-		cancel:    cancel,
-		parent:    c,
-		inbound:   nil,
-		outBlock:  make(chan struct{}, outboundCap),
-		LastPing:  time.Now().Unix(),
-		BaseFlags: c.baseFlags,
+		ID:                 msg.MuxID,
+		RecvSeq:            msg.Seq + 1,
+		SendSeq:            msg.Seq,
+		ctx:                ctx,
+		cancel:             cancel,
+		parent:             c,
+		inbound:            nil,
+		outBlock:           make(chan struct{}, outboundCap),
+		LastPing:           time.Now().Unix(),
+		BaseFlags:          c.baseFlags,
+		clientPingInterval: c.clientPingInterval,
 	}
 	// Acknowledge Mux created.
 	// Send async.
@@ -153,7 +153,7 @@ func newMuxStream(ctx context.Context, msg message, c *Connection, handler Strea
 	}(m.outBlock)
 
 	// Remote aliveness check if needed.
-	if msg.DeadlineMS == 0 || msg.DeadlineMS > uint32(lastPingThreshold/time.Millisecond) {
+	if msg.DeadlineMS == 0 || msg.DeadlineMS > uint32(4*c.clientPingInterval/time.Millisecond) {
 		go func() {
 			wg.Wait()
 			m.checkRemoteAlive()
@@ -164,9 +164,21 @@ func newMuxStream(ctx context.Context, msg message, c *Connection, handler Strea
 
 // handleInbound sends unblocks when we have delivered the message to the handler.
 func (m *muxServer) handleInbound(c *Connection, inbound <-chan []byte, handlerIn chan<- []byte) {
-	for in := range inbound {
-		handlerIn <- in
-		m.send(message{Op: OpUnblockClMux, MuxID: m.ID, Flags: c.baseFlags})
+	for {
+		select {
+		case <-m.ctx.Done():
+			return
+		case in, ok := <-inbound:
+			if !ok {
+				return
+			}
+			select {
+			case <-m.ctx.Done():
+				return
+			case handlerIn <- in:
+				m.send(message{Op: OpUnblockClMux, MuxID: m.ID, Flags: c.baseFlags})
+			}
+		}
 	}
 }
 
@@ -234,7 +246,7 @@ func (m *muxServer) handleRequests(ctx context.Context, msg message, send chan<-
 
 // checkRemoteAlive will check if the remote is alive.
 func (m *muxServer) checkRemoteAlive() {
-	t := time.NewTicker(lastPingThreshold / 4)
+	t := time.NewTicker(m.clientPingInterval)
 	defer t.Stop()
 	for {
 		select {
@@ -242,7 +254,7 @@ func (m *muxServer) checkRemoteAlive() {
 			return
 		case <-t.C:
 			last := time.Since(time.Unix(atomic.LoadInt64(&m.LastPing), 0))
-			if last > lastPingThreshold {
+			if last > 4*m.clientPingInterval {
 				gridLogIf(m.ctx, fmt.Errorf("canceling remote connection %s not seen for %v", m.parent, last))
 				m.close()
 				return
@@ -257,7 +269,7 @@ func (m *muxServer) checkSeq(seq uint32) (ok bool) {
 		if debugPrint {
 			fmt.Printf("expected sequence %d, got %d\n", m.RecvSeq, seq)
 		}
-		m.disconnect(fmt.Sprintf("receive sequence number mismatch. want %d, got %d", m.RecvSeq, seq))
+		m.disconnect(fmt.Sprintf("receive sequence number mismatch. want %d, got %d", m.RecvSeq, seq), false)
 		return false
 	}
 	m.RecvSeq++
@@ -268,13 +280,13 @@ func (m *muxServer) message(msg message) {
 	if debugPrint {
 		fmt.Printf("muxServer: received message %d, length %d\n", msg.Seq, len(msg.Payload))
 	}
+	if !m.checkSeq(msg.Seq) {
+		return
+	}
 	m.recvMu.Lock()
 	defer m.recvMu.Unlock()
 	if cap(m.inbound) == 0 {
-		m.disconnect("did not expect inbound message")
-		return
-	}
-	if !m.checkSeq(msg.Seq) {
+		m.disconnect("did not expect inbound message", true)
 		return
 	}
 	// Note, on EOF no value can be sent.
@@ -296,7 +308,7 @@ func (m *muxServer) message(msg message) {
 			fmt.Printf("muxServer: Sent seq %d to handler\n", msg.Seq)
 		}
 	default:
-		m.disconnect("handler blocked")
+		m.disconnect("handler blocked", true)
 	}
 }
 
@@ -332,7 +344,9 @@ func (m *muxServer) ping(seq uint32) pongMsg {
 	}
 }
 
-func (m *muxServer) disconnect(msg string) {
+// disconnect will disconnect the mux.
+// m.recvMu must be locked when calling this function.
+func (m *muxServer) disconnect(msg string, locked bool) {
 	if debugPrint {
 		fmt.Println("Mux", m.ID, "disconnecting. Reason:", msg)
 	}
@@ -341,6 +355,11 @@ func (m *muxServer) disconnect(msg string) {
 	} else {
 		m.send(message{Op: OpDisconnectClientMux, MuxID: m.ID})
 	}
+	// Unlock, since we are calling deleteMux, which will call close - which will lock recvMu.
+	if locked {
+		m.recvMu.Unlock()
+		defer m.recvMu.Lock()
+	}
 	m.parent.deleteMux(true, m.ID)
 }