func()

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
}