in tools/go-agent/instrument/logger/frameworks/zap.go [105:195]
func (z *Zap) CustomizedEnhance(path string, curFile *dst.File, cursor *dstutil.Cursor, allFiles []*dst.File) (map[string]string, bool) {
switch n := cursor.Node().(type) {
case *dst.TypeSpec:
// adding the context field into entry
st, ok := n.Type.(*dst.StructType)
if !ok || st.Fields == nil || n.Name == nil {
return nil, false
}
if n.Name.Name == "CheckedEntry" {
st.Fields.List = append(st.Fields.List,
// tracing context object
&dst.Field{Names: []*dst.Ident{dst.NewIdent("SWContext")}, Type: dst.NewIdent("interface{}")},
// tracing context field
&dst.Field{Names: []*dst.Ident{dst.NewIdent("SWContextField")}, Type: &dst.StarExpr{X: dst.NewIdent("Field")}},
// existing fields needs to be added into, such as generate from log.With("key", "value")
&dst.Field{Names: []*dst.Ident{dst.NewIdent("SWFields")}, Type: &dst.ArrayType{Elt: dst.NewIdent("Field")}},
)
curFile.Decls = append(curFile.Decls, &dst.GenDecl{
Tok: token.VAR,
Specs: []dst.Spec{
&dst.ValueSpec{Names: []*dst.Ident{
dst.NewIdent("SWReporterEnable"),
dst.NewIdent("SWLogEnable"),
}, Type: dst.NewIdent("bool")},
&dst.ValueSpec{Names: []*dst.Ident{
dst.NewIdent("SWReporterLabelKeys"),
}, Type: dst.NewIdent("[]string")},
&dst.ValueSpec{Names: []*dst.Ident{
dst.NewIdent("SWLogTracingContextKey"),
}, Type: dst.NewIdent("string")},
&dst.ValueSpec{Names: []*dst.Ident{
dst.NewIdent("SWFields"),
}, Type: dst.NewIdent("[]Field")},
},
})
return nil, true
}
if n.Name.Name == "Logger" {
st.Fields.List = append(st.Fields.List,
&dst.Field{Names: []*dst.Ident{dst.NewIdent("SWFields")}, Type: dst.NewIdent("[]zapcore.Field")})
return nil, true
}
case *dst.FuncDecl:
// enhance the method which check the log and generate the entry
if n.Recv != nil && len(n.Recv.List) == 1 && tools.GenerateTypeNameByExp(n.Recv.List[0].Type) == "*Logger" &&
n.Name.Name == "check" &&
n.Type.Results != nil && len(n.Type.Results.List) > 0 &&
tools.GenerateTypeNameByExp(n.Type.Results.List[0].Type) == "*zapcore.CheckedEntry" {
entryName := tools.EnhanceParameterNames(n.Type.Results, tools.FieldListTypeResult)[0].Name
recvName := tools.EnhanceParameterNames(n.Recv, tools.FieldListTypeRecv)[0].Name
// init the zapcore variables
z.initFunction = tools.GoStringToDecls(fmt.Sprintf(`func initZapCore() {
zapcore.SWReporterEnable = %s
zapcore.SWReporterLabelKeys = %s
zapcore.SWLogEnable = %s
}`, "LogReporterEnable", "LogReporterLabelKeys", "LogTracingContextEnable"))[0].(*dst.FuncDecl)
z.initImports = []*dst.ImportSpec{
{Path: &dst.BasicLit{Kind: token.STRING, Value: `"go.uber.org/zap/zapcore"`}},
}
return z.enhanceMethod(n, fmt.Sprintf("defer func() {if %s != nil {"+
"%s.SWContext, %s.SWContextField = %s%s(%s); %s.SWFields = %s.SWFields;}}()", entryName,
entryName, entryName, rewrite.StaticMethodPrefix, "ZapTracingContextEnhance", entryName, entryName, recvName)), true
}
if n.Recv != nil && len(n.Recv.List) == 1 && n.Name.Name == "With" &&
n.Type.Params != nil && len(n.Type.Params.List) == 1 &&
tools.GenerateTypeNameByExp(n.Recv.List[0].Type) == "*Logger" && tools.GenerateTypeNameByExp(n.Type.Params.List[0].Type) == "[]Field" {
recvs := tools.EnhanceParameterNames(n.Recv, tools.FieldListTypeRecv)
parameters := tools.EnhanceParameterNames(n.Type.Params, tools.FieldListTypeParam)
results := tools.EnhanceParameterNames(n.Type.Results, tools.FieldListTypeResult)
return z.enhanceMethod(n, fmt.Sprintf(`defer func() {if %s != nil { %s.SWFields = %sZap%s(%s, %s.SWFields) }}()`,
results[0].Name, results[0].Name, rewrite.StaticMethodPrefix, "KnownFieldFilter", parameters[0].Name, recvs[0].Name)), true
}
// enhance the method which write the checked entry context
if n.Recv != nil && len(n.Recv.List) == 1 && tools.GenerateTypeNameByExp(n.Recv.List[0].Type) == "*CheckedEntry" &&
n.Name.Name == "Write" &&
n.Type.Params != nil && len(n.Type.Params.List) == 1 &&
tools.GenerateTypeNameByExp(n.Type.Params.List[0].Type) == "[]Field" {
recvs := tools.EnhanceParameterNames(n.Recv, tools.FieldListTypeRecv)
parameters := tools.EnhanceParameterNames(n.Type.Params, tools.FieldListTypeParam)
return z.enhanceMethod(n, fmt.Sprintf(`if %s != nil { %s = %sZapcore%s(%s, %s, %s.SWFields, %s.SWContext, %s.SWContextField, SWReporterEnable, SWLogEnable, SWReporterLabelKeys) }`,
recvs[0].Name, parameters[0].Name, rewrite.StaticMethodPrefix, "ReportLogFromZapEntry", recvs[0].Name,
parameters[0].Name, recvs[0].Name, recvs[0].Name, recvs[0].Name)), true
}
}
return nil, false
}