Skip to content

Commit

Permalink
Add one more shs variable
Browse files Browse the repository at this point in the history
  • Loading branch information
lidizheng committed Jun 2, 2022
1 parent c71c6ac commit dec7c92
Showing 1 changed file with 7 additions and 6 deletions.
13 changes: 7 additions & 6 deletions server.go
Expand Up @@ -1430,15 +1430,16 @@ func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transp
if channelz.IsOn() {
s.incrCallsStarted()
}
shs := s.opts.statsHandlers
var statsBegin *stats.Begin
if len(s.opts.statsHandlers) != 0 {
if len(shs) != 0 {
beginTime := time.Now()
statsBegin = &stats.Begin{
BeginTime: beginTime,
IsClientStream: sd.ClientStreams,
IsServerStream: sd.ServerStreams,
}
for _, sh := range s.opts.statsHandlers {
for _, sh := range shs {
sh.HandleRPC(stream.Context(), statsBegin)
}
}
Expand All @@ -1452,10 +1453,10 @@ func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transp
maxReceiveMessageSize: s.opts.maxReceiveMessageSize,
maxSendMessageSize: s.opts.maxSendMessageSize,
trInfo: trInfo,
statsHandler: s.opts.statsHandlers,
statsHandler: shs,
}

if len(s.opts.statsHandlers) != 0 || trInfo != nil || channelz.IsOn() {
if len(shs) != 0 || trInfo != nil || channelz.IsOn() {
// See comment in processUnaryRPC on defers.
defer func() {
if trInfo != nil {
Expand All @@ -1469,15 +1470,15 @@ func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transp
ss.mu.Unlock()
}

if len(s.opts.statsHandlers) != 0 {
if len(shs) != 0 {
end := &stats.End{
BeginTime: statsBegin.BeginTime,
EndTime: time.Now(),
}
if err != nil && err != io.EOF {
end.Error = toRPCErr(err)
}
for _, sh := range s.opts.statsHandlers {
for _, sh := range shs {
sh.HandleRPC(stream.Context(), end)
}
}
Expand Down

0 comments on commit dec7c92

Please sign in to comment.