grpc: change type of Server.conns

Change Server.conns from a map[io.Closer]bool to a map[transport.ServerTransport]bool.
This commit is contained in:
Can Guler 2019-06-26 11:09:45 -07:00 committed by GitHub
parent 7472edcc1e
commit 915d20dcdb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 10 deletions

View File

@ -90,7 +90,7 @@ type Server struct {
mu sync.Mutex // guards following mu sync.Mutex // guards following
lis map[net.Listener]bool lis map[net.Listener]bool
conns map[io.Closer]bool conns map[transport.ServerTransport]bool
serve bool serve bool
drain bool drain bool
cv *sync.Cond // signaled when connections close for GracefulStop cv *sync.Cond // signaled when connections close for GracefulStop
@ -386,7 +386,7 @@ func NewServer(opt ...ServerOption) *Server {
s := &Server{ s := &Server{
lis: make(map[net.Listener]bool), lis: make(map[net.Listener]bool),
opts: opts, opts: opts,
conns: make(map[io.Closer]bool), conns: make(map[transport.ServerTransport]bool),
m: make(map[string]*service), m: make(map[string]*service),
quit: make(chan struct{}), quit: make(chan struct{}),
done: make(chan struct{}), done: make(chan struct{}),
@ -786,27 +786,27 @@ func (s *Server) traceInfo(st transport.ServerTransport, stream *transport.Strea
return trInfo return trInfo
} }
func (s *Server) addConn(c io.Closer) bool { func (s *Server) addConn(st transport.ServerTransport) bool {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
if s.conns == nil { if s.conns == nil {
c.Close() st.Close()
return false return false
} }
if s.drain { if s.drain {
// Transport added after we drained our existing conns: drain it // Transport added after we drained our existing conns: drain it
// immediately. // immediately.
c.(transport.ServerTransport).Drain() st.Drain()
} }
s.conns[c] = true s.conns[st] = true
return true return true
} }
func (s *Server) removeConn(c io.Closer) { func (s *Server) removeConn(st transport.ServerTransport) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
if s.conns != nil { if s.conns != nil {
delete(s.conns, c) delete(s.conns, st)
s.cv.Broadcast() s.cv.Broadcast()
} }
} }
@ -1423,8 +1423,8 @@ func (s *Server) GracefulStop() {
} }
s.lis = nil s.lis = nil
if !s.drain { if !s.drain {
for c := range s.conns { for st := range s.conns {
c.(transport.ServerTransport).Drain() st.Drain()
} }
s.drain = true s.drain = true
} }