change CallBack to handler, and move ctx out of Stats.

This commit is contained in:
Menghan Li
2016-10-31 16:42:23 -07:00
parent 46e80bf1f6
commit 7984a9c679
5 changed files with 72 additions and 97 deletions

View File

@ -559,9 +559,7 @@ func (s *Server) sendResponse(t transport.ServerTransport, stream *transport.Str
cbuf = new(bytes.Buffer)
}
if stats.On() {
outgoingPayloadStats = &stats.OutgoingPayloadStats{
Ctx: stream.Context(),
}
outgoingPayloadStats = &stats.OutgoingPayloadStats{}
}
p, err := encode(s.opts.codec, msg, cp, cbuf, outgoingPayloadStats)
if err != nil {
@ -577,7 +575,8 @@ func (s *Server) sendResponse(t transport.ServerTransport, stream *transport.Str
err = t.Write(stream, p, opts)
if outgoingPayloadStats != nil {
outgoingPayloadStats.SentTime = time.Now()
stats.CallBack()(outgoingPayloadStats)
stats.Handle(stream.Context(), outgoingPayloadStats)
}
return err
}
@ -604,7 +603,7 @@ func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.
var incomingPayloadStats *stats.IncomingPayloadStats
if stats.On() {
incomingPayloadStats = &stats.IncomingPayloadStats{
Ctx: stream.Context(),
ReceivedTime: time.Now(),
}
}
@ -675,7 +674,7 @@ func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.
if incomingPayloadStats != nil {
incomingPayloadStats.Data = req
incomingPayloadStats.Length = len(req)
stats.CallBack()(incomingPayloadStats)
stats.Handle(stream.Context(), incomingPayloadStats)
}
if trInfo != nil {
trInfo.tr.LazyLog(&payload{sent: false, msg: v}, true)

View File

@ -52,8 +52,6 @@ type Stats interface {
// InitStats indicates an RPC just started.
type InitStats struct {
// Ctx is the context associated with the RPC.
Ctx context.Context
// IsClient indicates if this stats is a client stats.
IsClient bool
// Method is the full RPC method string, i.e., /package.service/method.
@ -70,8 +68,6 @@ func (s *InitStats) isStats() {}
// IncomingPayloadStats contains the information for a incoming payload.
type IncomingPayloadStats struct {
// Ctx is the context associated with the RPC.
Ctx context.Context
// IsClient indicates if this stats is a client stats.
IsClient bool
// Data is the unencrypted message payload.
@ -88,8 +84,6 @@ func (s *IncomingPayloadStats) isStats() {}
// IncomingHeaderStats indicates a header is received.
type IncomingHeaderStats struct {
// Ctx is the context associated with the RPC.
Ctx context.Context
// IsClient indicates if this stats is a client stats.
IsClient bool
// WireLength is the wire length of header.
@ -100,8 +94,6 @@ func (s *IncomingHeaderStats) isStats() {}
// IncomingTrailerStats indicates a trailer is received.
type IncomingTrailerStats struct {
// Ctx is the context associated with the RPC.
Ctx context.Context
// IsClient indicates if this stats is a client stats.
IsClient bool
// WireLength is the wire length of header.
@ -112,8 +104,6 @@ func (s *IncomingTrailerStats) isStats() {}
// OutgoingPayloadStats contains the information for a outgoing payload.
type OutgoingPayloadStats struct {
// Ctx is the context associated with the RPC.
Ctx context.Context
// IsClient indicates if this stats is a client stats.
IsClient bool
// Data is the unencrypted message payload.
@ -130,8 +120,6 @@ func (s *OutgoingPayloadStats) isStats() {}
// OutgoingHeaderStats indicates a header is sent.
type OutgoingHeaderStats struct {
// Ctx is the context associated with the RPC.
Ctx context.Context
// IsClient indicates if this stats is a client stats.
IsClient bool
// WireLength is the wire length of header.
@ -142,8 +130,6 @@ func (s *OutgoingHeaderStats) isStats() {}
// OutgoingTrailerStats indicates a trailer is sent.
type OutgoingTrailerStats struct {
// Ctx is the context associated with the RPC.
Ctx context.Context
// IsClient indicates if this stats is a client stats.
IsClient bool
// WireLength is the wire length of header.
@ -153,8 +139,8 @@ type OutgoingTrailerStats struct {
func (s *OutgoingTrailerStats) isStats() {}
var (
on = new(int32)
f func(Stats)
on = new(int32)
handler func(context.Context, Stats)
)
// On indicates whether stats is started.
@ -162,15 +148,15 @@ func On() bool {
return atomic.LoadInt32(on) == 1
}
// CallBack returns the call back function registered by user to process the stats.
func CallBack() func(Stats) {
return f
// Handle returns the call back function registered by user to process the stats.
func Handle(ctx context.Context, s Stats) {
handler(ctx, s)
}
// RegisterCallBack registers the user callback function and starts the stats collection.
// This callback function will be called to process the stats.
func RegisterCallBack(cb func(Stats)) {
f = cb
// RegisterHandler registers the user handler function and starts the stats collection.
// This handler function will be called to process the stats.
func RegisterHandler(f func(context.Context, Stats)) {
handler = f
start()
}

View File

@ -48,7 +48,7 @@ import (
)
func TestStartStop(t *testing.T) {
stats.RegisterCallBack(nil)
stats.RegisterHandler(nil)
defer stats.Stop() // Stop stats in the case of the first Fatalf.
if stats.On() != true {
t.Fatalf("after start.RegisterCallBack(_), stats.On() = false, want true")
@ -268,17 +268,21 @@ type expectedData struct {
outgoing []*testpb.SimpleResponse
}
func checkInitStats(t *testing.T, s stats.Stats, e *expectedData) {
t.Logf(" - %T", s)
type gotData struct {
ctx context.Context
s stats.Stats
}
func checkInitStats(t *testing.T, d gotData, e *expectedData) {
var (
ok bool
st *stats.InitStats
)
if st, ok = s.(*stats.InitStats); !ok {
t.Fatalf("got %T, want InitStats", s)
if st, ok = d.s.(*stats.InitStats); !ok {
t.Fatalf("got %T, want InitStats", d.s)
}
if st.Ctx == nil {
t.Fatalf("st.Ctx = nil, want <non-nil>")
if d.ctx == nil {
t.Fatalf("d.ctx = nil, want <non-nil>")
}
if st.IsClient {
t.Fatalf("st IsClient = true, want false")
@ -294,17 +298,16 @@ func checkInitStats(t *testing.T, s stats.Stats, e *expectedData) {
}
}
func checkIncomingHeaderStats(t *testing.T, s stats.Stats, e *expectedData) {
t.Logf(" - %T", s)
func checkIncomingHeaderStats(t *testing.T, d gotData, e *expectedData) {
var (
ok bool
st *stats.IncomingHeaderStats
)
if st, ok = s.(*stats.IncomingHeaderStats); !ok {
t.Fatalf("got %T, want IncomingHeaderStats", s)
if st, ok = d.s.(*stats.IncomingHeaderStats); !ok {
t.Fatalf("got %T, want IncomingHeaderStats", d.s)
}
if st.Ctx == nil {
t.Fatalf("st.Ctx = nil, want <non-nil>")
if d.ctx == nil {
t.Fatalf("d.ctx = nil, want <non-nil>")
}
if st.IsClient {
t.Fatalf("st.IsClient = true, want false")
@ -315,17 +318,16 @@ func checkIncomingHeaderStats(t *testing.T, s stats.Stats, e *expectedData) {
}
}
func checkIncomingPayloadStats(t *testing.T, s stats.Stats, e *expectedData) {
t.Logf(" - %T", s)
func checkIncomingPayloadStats(t *testing.T, d gotData, e *expectedData) {
var (
ok bool
st *stats.IncomingPayloadStats
)
if st, ok = s.(*stats.IncomingPayloadStats); !ok {
t.Fatalf("got %T, want IncomingPayloadStats", s)
if st, ok = d.s.(*stats.IncomingPayloadStats); !ok {
t.Fatalf("got %T, want IncomingPayloadStats", d.s)
}
if st.Ctx == nil {
t.Fatalf("st.Ctx = nil, want <non-nil>")
if d.ctx == nil {
t.Fatalf("d.ctx = nil, want <non-nil>")
}
if st.IsClient {
t.Fatalf("st IsClient = true, want false")
@ -347,17 +349,16 @@ func checkIncomingPayloadStats(t *testing.T, s stats.Stats, e *expectedData) {
}
}
func checkIncomingTrailerStats(t *testing.T, s stats.Stats, e *expectedData) {
t.Logf(" - %T", s)
func checkIncomingTrailerStats(t *testing.T, d gotData, e *expectedData) {
var (
ok bool
st *stats.IncomingTrailerStats
)
if st, ok = s.(*stats.IncomingTrailerStats); !ok {
t.Fatalf("got %T, want IncomingTrailerStats", s)
if st, ok = d.s.(*stats.IncomingTrailerStats); !ok {
t.Fatalf("got %T, want IncomingTrailerStats", d.s)
}
if st.Ctx == nil {
t.Fatalf("st.Ctx = nil, want <non-nil>")
if d.ctx == nil {
t.Fatalf("d.ctx = nil, want <non-nil>")
}
if st.IsClient {
t.Fatalf("st.IsClient = true, want false")
@ -368,17 +369,16 @@ func checkIncomingTrailerStats(t *testing.T, s stats.Stats, e *expectedData) {
}
}
func checkOutgoingHeaderStats(t *testing.T, s stats.Stats, e *expectedData) {
t.Logf(" - %T", s)
func checkOutgoingHeaderStats(t *testing.T, d gotData, e *expectedData) {
var (
ok bool
st *stats.OutgoingHeaderStats
)
if st, ok = s.(*stats.OutgoingHeaderStats); !ok {
t.Fatalf("got %T, want OutgoingHeaderStats", s)
if st, ok = d.s.(*stats.OutgoingHeaderStats); !ok {
t.Fatalf("got %T, want OutgoingHeaderStats", d.s)
}
if st.Ctx == nil {
t.Fatalf("st.Ctx = nil, want <non-nil>")
if d.ctx == nil {
t.Fatalf("d.ctx = nil, want <non-nil>")
}
if st.IsClient {
t.Fatalf("st IsClient = true, want false")
@ -389,17 +389,16 @@ func checkOutgoingHeaderStats(t *testing.T, s stats.Stats, e *expectedData) {
}
}
func checkOutgoingPayloadStats(t *testing.T, s stats.Stats, e *expectedData) {
t.Logf(" - %T", s)
func checkOutgoingPayloadStats(t *testing.T, d gotData, e *expectedData) {
var (
ok bool
st *stats.OutgoingPayloadStats
)
if st, ok = s.(*stats.OutgoingPayloadStats); !ok {
t.Fatalf("got %T, want OutgoingPayloadStats", s)
if st, ok = d.s.(*stats.OutgoingPayloadStats); !ok {
t.Fatalf("got %T, want OutgoingPayloadStats", d.s)
}
if st.Ctx == nil {
t.Fatalf("st.Ctx = nil, want <non-nil>")
if d.ctx == nil {
t.Fatalf("d.ctx = nil, want <non-nil>")
}
if st.IsClient {
t.Fatalf("st IsClient = true, want false")
@ -421,17 +420,16 @@ func checkOutgoingPayloadStats(t *testing.T, s stats.Stats, e *expectedData) {
}
}
func checkOutgoingTrailerStats(t *testing.T, s stats.Stats, e *expectedData) {
t.Logf(" - %T", s)
func checkOutgoingTrailerStats(t *testing.T, d gotData, e *expectedData) {
var (
ok bool
st *stats.OutgoingTrailerStats
)
if st, ok = s.(*stats.OutgoingTrailerStats); !ok {
t.Fatalf("got %T, want OutgoingTrailerStats", s)
if st, ok = d.s.(*stats.OutgoingTrailerStats); !ok {
t.Fatalf("got %T, want OutgoingTrailerStats", d.s)
}
if st.Ctx == nil {
t.Fatalf("st.Ctx = nil, want <non-nil>")
if d.ctx == nil {
t.Fatalf("d.ctx = nil, want <non-nil>")
}
if st.IsClient {
t.Fatalf("st IsClient = true, want false")
@ -445,12 +443,12 @@ func checkOutgoingTrailerStats(t *testing.T, s stats.Stats, e *expectedData) {
func TestServerStatsUnaryRPC(t *testing.T) {
var (
mu sync.Mutex
got []stats.Stats
got []gotData
)
stats.RegisterCallBack(func(s stats.Stats) {
stats.RegisterHandler(func(ctx context.Context, s stats.Stats) {
mu.Lock()
defer mu.Unlock()
got = append(got, s)
got = append(got, gotData{ctx, s})
})
te := newTest(t, "")
@ -467,7 +465,7 @@ func TestServerStatsUnaryRPC(t *testing.T) {
outgoing: []*testpb.SimpleResponse{resp},
}
for i, f := range []func(t *testing.T, s stats.Stats, e *expectedData){
for i, f := range []func(t *testing.T, d gotData, e *expectedData){
checkInitStats,
checkIncomingHeaderStats,
checkIncomingPayloadStats,
@ -486,12 +484,12 @@ func TestServerStatsUnaryRPC(t *testing.T) {
func TestServerStatsStreamingRPC(t *testing.T) {
var (
mu sync.Mutex
got []stats.Stats
got []gotData
)
stats.RegisterCallBack(func(s stats.Stats) {
stats.RegisterHandler(func(ctx context.Context, s stats.Stats) {
mu.Lock()
defer mu.Unlock()
got = append(got, s)
got = append(got, gotData{ctx, s})
})
te := newTest(t, "gzip")
@ -510,12 +508,12 @@ func TestServerStatsStreamingRPC(t *testing.T) {
outgoing: resps,
}
checkFuncs := []func(t *testing.T, s stats.Stats, e *expectedData){
checkFuncs := []func(t *testing.T, d gotData, e *expectedData){
checkInitStats,
checkIncomingHeaderStats,
checkOutgoingHeaderStats,
}
ioPayFuncs := []func(t *testing.T, s stats.Stats, e *expectedData){
ioPayFuncs := []func(t *testing.T, d gotData, e *expectedData){
checkIncomingPayloadStats,
checkOutgoingPayloadStats,
}

View File

@ -485,9 +485,7 @@ func (ss *serverStream) SendMsg(m interface{}) (err error) {
}()
var outgoingPayloadStats *stats.OutgoingPayloadStats
if stats.On() {
outgoingPayloadStats = &stats.OutgoingPayloadStats{
Ctx: ss.s.Context(),
}
outgoingPayloadStats = &stats.OutgoingPayloadStats{}
}
out, err := encode(ss.codec, m, ss.cp, ss.cbuf, outgoingPayloadStats)
defer func() {
@ -504,7 +502,7 @@ func (ss *serverStream) SendMsg(m interface{}) (err error) {
}
if outgoingPayloadStats != nil {
outgoingPayloadStats.SentTime = time.Now()
stats.CallBack()(outgoingPayloadStats)
stats.Handle(ss.s.Context(), outgoingPayloadStats)
}
return nil
}
@ -526,9 +524,7 @@ func (ss *serverStream) RecvMsg(m interface{}) (err error) {
}()
var incomingPayloadStats *stats.IncomingPayloadStats
if stats.On() {
incomingPayloadStats = &stats.IncomingPayloadStats{
Ctx: ss.s.Context(),
}
incomingPayloadStats = &stats.IncomingPayloadStats{}
}
if err := recv(ss.p, ss.codec, ss.s, ss.dc, m, ss.maxMsgSize, incomingPayloadStats); err != nil {
if err == io.EOF {
@ -540,7 +536,7 @@ func (ss *serverStream) RecvMsg(m interface{}) (err error) {
return toRPCErr(err)
}
if incomingPayloadStats != nil {
stats.CallBack()(incomingPayloadStats)
stats.Handle(ss.s.Context(), incomingPayloadStats)
}
return nil
}

View File

@ -237,18 +237,16 @@ func (t *http2Server) operateHeaders(frame *http2.MetaHeadersFrame, handle func(
}
if stats.On() {
initStats := &stats.InitStats{
Ctx: s.ctx,
Method: s.method,
RemoteAddr: t.conn.RemoteAddr(),
LocalAddr: t.conn.LocalAddr(),
Encryption: s.recvCompress,
}
stats.CallBack()(initStats)
stats.Handle(s.ctx, initStats)
incomingHeaderStats := &stats.IncomingHeaderStats{
Ctx: s.ctx,
WireLength: int(frame.Header().Length),
}
stats.CallBack()(incomingHeaderStats)
stats.Handle(s.ctx, incomingHeaderStats)
}
handle(s)
return
@ -530,10 +528,9 @@ func (t *http2Server) WriteHeader(s *Stream, md metadata.MD) error {
}
if stats.On() {
outgoingHeaderStats := &stats.OutgoingHeaderStats{
Ctx: s.Context(),
WireLength: bufLen,
}
stats.CallBack()(outgoingHeaderStats)
stats.Handle(s.Context(), outgoingHeaderStats)
}
t.writableChan <- 0
return nil
@ -594,10 +591,9 @@ func (t *http2Server) WriteStatus(s *Stream, statusCode codes.Code, statusDesc s
}
if stats.On() {
outgoingTrailerStats := &stats.OutgoingTrailerStats{
Ctx: s.Context(),
WireLength: bufLen,
}
stats.CallBack()(outgoingTrailerStats)
stats.Handle(s.Context(), outgoingTrailerStats)
}
t.closeStream(s)
t.writableChan <- 0