func setup(s *testerator.Setup) error {
if sniff, ok := s.Context.Value(&ctxKey).(*searchSniffer); ok {
sniff.searchIndexDocumentRequests = nil
return nil
}
sniff := &searchSniffer{}
s.Context = context.WithValue(s.Context, &ctxKey, sniff)
s.Context = appengine.WithAPICallFunc(s.Context, func(ctx context.Context, service, method string, in, out proto.Message) error {
if service == "search" && method == "IndexDocument" {
b, err := proto.Marshal(in)
if err != nil {
return err
}
req := &searchpb.IndexDocumentRequest{}
err = proto.Unmarshal(b, req)
if err != nil {
return err
}
sniff.searchIndexDocumentRequests = append(sniff.searchIndexDocumentRequests, req)
}
return appengine.APICall(ctx, service, method, in, out)
})
return nil
}
func override(ctx context.Context, service, method string, in, out proto.Message) error {
stats := ctx.Value(statsKey).(*requestStats)
stats.wg.Add(1)
defer stats.wg.Done()
if service == "__go__" {
return appengine.APICall(ctx, service, method, in, out)
}
// NOTE: This limits size of stack traces to 65KiB.
b := make([]byte, 65536)
i := runtime.Stack(b, false)
stat := rpcStat{
Service: service,
Method: method,
Start: time.Now(),
Offset: time.Since(stats.Start),
StackData: string(b[0:i]),
}
err := appengine.APICall(ctx, service, method, in, out)
stat.Duration = time.Since(stat.Start)
stat.In = in.String()
stat.Out = out.String()
stat.Cost = getCost(out)
if len(stat.In) > ProtoMaxBytes {
stat.In = stat.In[:ProtoMaxBytes] + "..."
}
if len(stat.Out) > ProtoMaxBytes {
stat.Out = stat.Out[:ProtoMaxBytes] + "..."
}
stats.lock.Lock()
stats.RPCStats = append(stats.RPCStats, stat)
stats.Cost += stat.Cost
stats.lock.Unlock()
return err
}
func override(ctx context.Context, service, method string, in, out proto.Message) error {
stats := stats(ctx)
stats.wg.Add(1)
defer stats.wg.Done()
if service == "__go__" {
return appengine.APICall(ctx, service, method, in, out)
}
stat := rpcStat{
Service: service,
Method: method,
Start: time.Now(),
Offset: time.Since(stats.Start),
StackData: string(debug.Stack()),
}
err := appengine.APICall(ctx, service, method, in, out)
stat.Duration = time.Since(stat.Start)
stat.In = in.String()
stat.Out = out.String()
stat.Cost = getCost(out)
if len(stat.In) > ProtoMaxBytes {
stat.In = stat.In[:ProtoMaxBytes] + "..."
}
if len(stat.Out) > ProtoMaxBytes {
stat.Out = stat.Out[:ProtoMaxBytes] + "..."
}
stats.lock.Lock()
stats.RPCStats = append(stats.RPCStats, stat)
stats.Cost += stat.Cost
stats.lock.Unlock()
return err
}