go-pkg/cmux
2
0
Fork 0
mirror of https://github.com/soheilhy/cmux.git synced 2024-09-20 02:55:46 +08:00
Code Issues Projects Releases Wiki Activity

fix blocking select on donec

Browse Source 
Signed-off-by: Abhilash Gnan <abhilashgnan@gmail.com>
This commit is contained in:
Abhilash Gnan 2020-05-21 09:16:06 +02:00 committed by Soheil Hassas Yeganeh
parent 3c0ee784ab
commit ae889c5259
1 changed files with 27 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
cmux.go
View File

@ -15,6 +15,7 @@
package cmux package cmux
import ( import (
"errors"
"fmt" "fmt"
"io" "io"
"net" "net"
@ -61,6 +62,9 @@ func (e errListenerClosed) Timeout() bool { return false }
// listener is closed. // listener is closed.
var ErrListenerClosed = errListenerClosed("mux: listener closed") var ErrListenerClosed = errListenerClosed("mux: listener closed")
// ErrServerClosed is returned from muxListener.Accept when mux server is closed.
var ErrServerClosed = errors.New("mux: server closed")
// for readability of readTimeout // for readability of readTimeout
var noTimeout time.Duration var noTimeout time.Duration
@ -148,7 +152,7 @@ func (m *cMux) Serve() error {
var wg sync.WaitGroup var wg sync.WaitGroup
defer func() { defer func() {
m.closeDoneChan() m.closeDoneChans()
wg.Wait() wg.Wait()
for _, sl := range m.sls { for _, sl := range m.sls {
@ -161,14 +165,6 @@ func (m *cMux) Serve() error {
}() }()
for { for {
select {
case <-m.donec:
// cmux was closed with cmux.Close()
return nil
default:
// do nothing
}
c, err := m.root.Accept() c, err := m.root.Accept()
if err != nil { if err != nil {
if !m.handleErr(err) { if !m.handleErr(err) {
@ -215,16 +211,24 @@ func (m *cMux) serve(c net.Conn, donec <-chan struct{}, wg *sync.WaitGroup) {
} }
func (m *cMux) Close() { func (m *cMux) Close() {
m.closeDoneChan() m.closeDoneChans()
} }
func (m *cMux) closeDoneChan() { func (m *cMux) closeDoneChans() {
select { select {
case <-m.donec: case <-m.donec:
// Already closed. Don't close again // Already closed. Don't close again
default: default:
close(m.donec) close(m.donec)
} }
for _, sl := range m.sls {
select {
case <-sl.l.donec:
// Already closed. Don't close again
default:
close(sl.l.donec)
}
}
} }
func (m *cMux) HandleError(h ErrorHandler) { func (m *cMux) HandleError(h ErrorHandler) {
@ -246,14 +250,22 @@ func (m *cMux) handleErr(err error) bool {
type muxListener struct { type muxListener struct {
net.Listener net.Listener
connc chan net.Conn connc chan net.Conn
donec chan struct{}
} }
func (l muxListener) Accept() (net.Conn, error) { func (l muxListener) Accept() (net.Conn, error) {
c, ok := <-l.connc select {
if !ok { case c, ok := <-l.connc:
return nil, ErrListenerClosed if !ok {
return nil, ErrListenerClosed
}
return c, nil
case <-l.donec:
return nil, ErrServerClosed
default:
// do nothing
} }
return c, nil return nil, ErrServerClosed
} }
// MuxConn wraps a net.Conn and provides transparent sniffing of connection data. // MuxConn wraps a net.Conn and provides transparent sniffing of connection data.