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

Updates grpc/grpc-go#75

server.go

@@ -264,7 +264,11 @@ func (s *Server) Serve(lis net.Listener) error {
 		}
 		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)
 	if err != nil {
 		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)
 		return
 	}
-			defer st.Close()
-			s.mu.Lock()
-			if s.conns == nil {
-				s.mu.Unlock()
+	if !s.addConn(st) {
+		c.Close()
 		return
 	}
-			s.conns[st] = true
-			s.mu.Unlock()
+	s.serveStreams(st)
+}
+
+func (s *Server) serveStreams(st transport.ServerTransport) {
+	defer s.removeConn(st)
+	defer st.Close()
 	var wg sync.WaitGroup
 	st.HandleStreams(func(stream *transport.Stream) {
-				var trInfo *traceInfo
-				if EnableTracing {
+		wg.Add(1)
+		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{
 		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 {
 		trInfo.firstLine.deadline = dl.Sub(time.Now())
 	}
-				}
-				wg.Add(1)
-				go func() {
-					s.handleStream(st, stream, trInfo)
-					wg.Done()
-				}()
-			})
-			wg.Wait()
+	return trInfo
+}
+
+func (s *Server) addConn(st transport.ServerTransport) bool {
 	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)
-			s.mu.Unlock()
-		}()
 	}
 }