server: break up the Server.Serve method into some reusable parts

Updates grpc/grpc-go#75
This commit is contained in:
Brad Fitzpatrick
2016-01-28 19:52:42 +00:00
parent 3f30c980d6
commit 9d2ecf553a

View File

@ -264,7 +264,11 @@ func (s *Server) Serve(lis net.Listener) error {
} }
s.mu.Unlock() s.mu.Unlock()
go func() { go s.serveNewHTTP2Transport(c, authInfo)
}
}
func (s *Server) serveNewHTTP2Transport(c net.Conn, authInfo credentials.AuthInfo) {
st, err := transport.NewServerTransport("http2", c, s.opts.maxConcurrentStreams, authInfo) st, err := transport.NewServerTransport("http2", c, s.opts.maxConcurrentStreams, authInfo)
if err != nil { if err != nil {
s.mu.Lock() s.mu.Lock()
@ -274,18 +278,33 @@ func (s *Server) Serve(lis net.Listener) error {
grpclog.Println("grpc: Server.Serve failed to create ServerTransport: ", err) grpclog.Println("grpc: Server.Serve failed to create ServerTransport: ", err)
return return
} }
defer st.Close() if !s.addConn(st) {
s.mu.Lock() c.Close()
if s.conns == nil {
s.mu.Unlock()
return return
} }
s.conns[st] = true s.serveStreams(st)
s.mu.Unlock() }
func (s *Server) serveStreams(st transport.ServerTransport) {
defer s.removeConn(st)
defer st.Close()
var wg sync.WaitGroup var wg sync.WaitGroup
st.HandleStreams(func(stream *transport.Stream) { st.HandleStreams(func(stream *transport.Stream) {
var trInfo *traceInfo wg.Add(1)
if EnableTracing { go func() {
defer wg.Done()
s.handleStream(st, stream, s.traceInfo(st, stream))
}()
})
wg.Wait()
}
// traceInfo returns a traceInfo and associates it with stream, if tracing is enabled.
// If tracing is not enabled, it returns nil.
func (s *Server) traceInfo(st transport.ServerTransport, stream *transport.Stream) (trInfo *traceInfo) {
if !EnableTracing {
return nil
}
trInfo = &traceInfo{ trInfo = &traceInfo{
tr: trace.New("grpc.Recv."+methodFamily(stream.Method()), stream.Method()), tr: trace.New("grpc.Recv."+methodFamily(stream.Method()), stream.Method()),
} }
@ -295,18 +314,24 @@ func (s *Server) Serve(lis net.Listener) error {
if dl, ok := stream.Context().Deadline(); ok { if dl, ok := stream.Context().Deadline(); ok {
trInfo.firstLine.deadline = dl.Sub(time.Now()) trInfo.firstLine.deadline = dl.Sub(time.Now())
} }
} return trInfo
wg.Add(1) }
go func() {
s.handleStream(st, stream, trInfo) func (s *Server) addConn(st transport.ServerTransport) bool {
wg.Done()
}()
})
wg.Wait()
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock()
if s.conns == nil {
return false
}
s.conns[st] = true
return true
}
func (s *Server) removeConn(st transport.ServerTransport) {
s.mu.Lock()
defer s.mu.Unlock()
if s.conns != nil {
delete(s.conns, st) delete(s.conns, st)
s.mu.Unlock()
}()
} }
} }