diff --git a/thrift/annotation.go b/thrift/annotation.go index 4c84271c..4c20d761 100644 --- a/thrift/annotation.go +++ b/thrift/annotation.go @@ -20,6 +20,7 @@ import ( "context" "fmt" "strings" + "sync" "github.com/cloudwego/dynamicgo/http" "github.com/cloudwego/dynamicgo/meta" @@ -124,6 +125,15 @@ type OptionMapping interface { Map(ctx context.Context, opts Options) Options } +func applyOptionMapping(ctx context.Context, mapper OptionMapping, opts *Options) { + annotationMappers := opts.annotationMappers + mapped := mapper.Map(ctx, *opts) + if mapped.annotationMappers == nil { + mapped.annotationMappers = annotationMappers + } + *opts = mapped +} + // ValueMapping is used to convert thrift value while running convertion. // See also: thrift/annotation/value_mapping.go type ValueMapping interface { @@ -196,6 +206,53 @@ func makeAnnotation(ctx context.Context, anns []parser.Annotation, scope AnnoSco } } +type annotationMapperRegistry map[string]map[AnnoScope]AnnotationMapper + +func (r annotationMapperRegistry) find(key string, scope AnnoScope) AnnotationMapper { + m := r[key] + if m == nil { + return nil + } + return m[scope] +} + +func (r annotationMapperRegistry) clone() annotationMapperRegistry { + if r == nil { + return nil + } + ret := make(annotationMapperRegistry, len(r)) + for key, scopes := range r { + if scopes == nil { + continue + } + copied := make(map[AnnoScope]AnnotationMapper, len(scopes)) + for scope, mapper := range scopes { + copied[scope] = mapper + } + ret[key] = copied + } + return ret +} + +func (r annotationMapperRegistry) register(scope AnnoScope, mapper AnnotationMapper, keys ...string) { + for _, key := range keys { + m := r[key] + if m == nil { + m = make(map[AnnoScope]AnnotationMapper) + r[key] = m + } + m[scope] = mapper + } +} + +func (r annotationMapperRegistry) remove(scope AnnoScope, keys ...string) { + for _, key := range keys { + if m := r[key]; m != nil { + delete(m, scope) + } + } +} + // AnnotationMapper is used to convert a annotation to equivalent annotations // desc is specific to its registered AnnoScope: // @@ -208,37 +265,76 @@ type AnnotationMapper interface { Map(ctx context.Context, ann []parser.Annotation, desc interface{}, opt Options) (cur []parser.Annotation, next []parser.Annotation, err error) } -var annotationMapper = map[string]map[AnnoScope]AnnotationMapper{} +var ( + defaultAnnotationMapperMu sync.RWMutex + defaultAnnotationMapper = annotationMapperRegistry{} +) // RegisterAnnotationMapper register a annotation mapper on specific scope func RegisterAnnotationMapper(scope AnnoScope, mapper AnnotationMapper, keys ...string) { - for _, key := range keys { - m := annotationMapper[key] - if m == nil { - m = make(map[AnnoScope]AnnotationMapper) - annotationMapper[key] = m - } - m[scope] = mapper - } + defaultAnnotationMapperMu.Lock() + defer defaultAnnotationMapperMu.Unlock() + defaultAnnotationMapper.register(scope, mapper, keys...) } func FindAnnotationMapper(key string, scope AnnoScope) AnnotationMapper { - m := annotationMapper[key] - if m == nil { - return nil - } - return m[scope] + defaultAnnotationMapperMu.RLock() + defer defaultAnnotationMapperMu.RUnlock() + return defaultAnnotationMapper.find(key, scope) } func RemoveAnnotationMapper(scope AnnoScope, keys ...string) { - for _, key := range keys { - m := annotationMapper[key] - if m != nil { - if _, ok := m[scope]; ok { - delete(m, scope) - } - } + defaultAnnotationMapperMu.Lock() + defer defaultAnnotationMapperMu.Unlock() + defaultAnnotationMapper.remove(scope, keys...) +} + +func cloneDefaultAnnotationMappers() annotationMapperRegistry { + defaultAnnotationMapperMu.RLock() + defer defaultAnnotationMapperMu.RUnlock() + return defaultAnnotationMapper.clone() +} + +func (opts Options) cloneAnnotationMappersForUpdate() annotationMapperRegistry { + if opts.annotationMappers == nil { + return cloneDefaultAnnotationMappers() + } + return (*opts.annotationMappers).clone() +} + +// RegisterAnnotationMapper registers a mapper on this Options. It uses +// copy-on-write so modifying a copied Options value does not affect the source. +// It must not be called concurrently on the same *Options. +func (opts *Options) RegisterAnnotationMapper(scope AnnoScope, mapper AnnotationMapper, keys ...string) { + registry := opts.cloneAnnotationMappersForUpdate() + registry.register(scope, mapper, keys...) + opts.annotationMappers = ®istry +} + +// RemoveAnnotationMapper removes mappers from this Options. It uses +// copy-on-write so modifying a copied Options value does not affect the source. +// It must not be called concurrently on the same *Options. +func (opts *Options) RemoveAnnotationMapper(scope AnnoScope, keys ...string) { + registry := opts.cloneAnnotationMappersForUpdate() + registry.remove(scope, keys...) + opts.annotationMappers = ®istry +} + +// FindAnnotationMapper finds a mapper from this Options or the global defaults +// if this Options has not taken a snapshot yet. +func (opts Options) FindAnnotationMapper(key string, scope AnnoScope) AnnotationMapper { + if opts.annotationMappers != nil { + return (*opts.annotationMappers).find(key, scope) + } + return FindAnnotationMapper(key, scope) +} + +func (opts Options) prepareAnnotationMappersForParse() Options { + if opts.annotationMappers == nil { + registry := cloneDefaultAnnotationMappers() + opts.annotationMappers = ®istry } + return opts } //------------------------------- IDL processing logic ------------------------------- @@ -275,7 +371,7 @@ func mapAnnotations(ctx context.Context, as parser.Annotations, scope AnnoScope, cur := make([]parser.Annotation, 0, len(as)) // try find mapper for _, a := range as { - if mapper := FindAnnotationMapper(a.Key, scope); mapper != nil { + if mapper := opt.FindAnnotationMapper(a.Key, scope); mapper != nil { con.Add(*a, mapper) } else { // no mapper found, just append it to the result @@ -340,7 +436,7 @@ func handleAnnotation(ctx context.Context, scope AnnoScope, ann Annotation, valu if !ok { return fmt.Errorf("annotation %#v for %d is not OptionMaker", handle, ann.ID()) } - *opts = om.Map(ctx, *opts) + applyOptionMapping(ctx, om, opts) return nil default: //NOTICE: ignore unsupported annotations @@ -362,7 +458,7 @@ func handleFieldAnnotation(ctx context.Context, ann Annotation, values []parser. if !ok { return fmt.Errorf("annotation %#v for %d is not OptionMaker", handle, ann.ID()) } - *opts = om.Map(ctx, *opts) + applyOptionMapping(ctx, om, opts) return nil case AnnoKindHttpMappping: hm, ok := handle.(HttpMapping) diff --git a/thrift/annotation/anno_mapping_test.go b/thrift/annotation/anno_mapping_test.go index b67b2ded..4edf360b 100644 --- a/thrift/annotation/anno_mapping_test.go +++ b/thrift/annotation/anno_mapping_test.go @@ -17,13 +17,44 @@ package annotation import ( + "context" + "fmt" + "sync" "testing" "github.com/cloudwego/dynamicgo/thrift" + "github.com/cloudwego/thriftgo/parser" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) +const resetOptionsAnnotationKey = "dynamicgo.test.reset_options" + +type resetOptionsAnnotation struct{} + +func (resetOptionsAnnotation) ID() thrift.AnnoID { + return thrift.MakeAnnoID(thrift.AnnoKindOptionMapping, thrift.AnnoScopeService, 65000) +} + +func (resetOptionsAnnotation) Make(context.Context, []parser.Annotation, interface{}) (interface{}, error) { + return resetOptionsMapping{}, nil +} + +type resetOptionsMapping struct{} + +func (resetOptionsMapping) Map(context.Context, thrift.Options) thrift.Options { + return thrift.Options{} +} + +type countingAnnotationMapper struct { + calls *int +} + +func (m countingAnnotationMapper) Map(context.Context, []parser.Annotation, interface{}, thrift.Options) ([]parser.Annotation, []parser.Annotation, error) { + (*m.calls)++ + return nil, nil, nil +} + func TestMain(m *testing.M) { InitAGWAnnos() m.Run() @@ -80,6 +111,176 @@ func TestGoTagJSON(t *testing.T) { require.NoError(t, err) } +func TestGoTagMapperOptionsIsolation(t *testing.T) { + content := ` + namespace go kitex.test.server + struct Base { + 1: string Msg (go.tag = "json:\"message\"") + } + + service InboxService { + string ExampleMethod(1: Base req) + } + ` + + opts := thrift.NewDefaultOptions() + opts.RemoveAnnotationMapper(thrift.AnnoScopeField, "go.tag") + p, err := GetDescFromContentWithOptions(content, "ExampleMethod", opts) + require.NoError(t, err) + req := p.Request().Struct().Fields()[0].Type() + require.Equal(t, "Msg", req.Struct().FieldById(1).Alias()) + + p, err = GetDescFromContent(content, "ExampleMethod") + require.NoError(t, err) + req = p.Request().Struct().Fields()[0].Type() + require.Equal(t, "message", req.Struct().FieldById(1).Alias()) + + opts = thrift.NewDefaultOptions() + opts.RemoveAnnotationMapper(thrift.AnnoScopeField, "go.tag") + opts.RegisterAnnotationMapper(thrift.AnnoScopeField, goTagMapper{}, "go.tag") + p, err = GetDescFromContentWithOptions(content, "ExampleMethod", opts) + require.NoError(t, err) + req = p.Request().Struct().Fields()[0].Type() + require.Equal(t, "message", req.Struct().FieldById(1).Alias()) +} + +func TestOptionMappingPreservesAnnotationMapperSnapshot(t *testing.T) { + const mapperKey = "dynamicgo.test.count_mapper" + + thrift.RegisterAnnotation(resetOptionsAnnotation{}, resetOptionsAnnotationKey) + calls := 0 + thrift.RegisterAnnotationMapper(thrift.AnnoScopeField, countingAnnotationMapper{calls: &calls}, mapperKey) + t.Cleanup(func() { + thrift.RemoveAnnotationMapper(thrift.AnnoScopeField, mapperKey) + }) + + opts := thrift.NewDefaultOptions() + opts.RemoveAnnotationMapper(thrift.AnnoScopeField, mapperKey) + _, err := GetDescFromContentWithOptions(` + namespace go kitex.test.server + struct Base { + 1: string Msg (`+mapperKey+` = "true") + } + + service InboxService { + string ExampleMethod(1: Base req) + } (`+resetOptionsAnnotationKey+` = "true") + `, "ExampleMethod", opts) + require.NoError(t, err) + require.Zero(t, calls, "removed mapper was restored after OptionMapping returned fresh Options") +} + +func TestAnnotationMapperOptionsCopyOnWrite(t *testing.T) { + content := ` + namespace go kitex.test.server + struct Base { + 1: string Msg (go.tag = "json:\"message\"") + } + + service InboxService { + string ExampleMethod(1: Base req) + } + ` + + base := thrift.NewDefaultOptions() + base.RegisterAnnotationMapper(thrift.AnnoScopeField, goTagMapper{}, "go.tag") + copied := base + require.True(t, base == copied) + + const workers = 8 + var wg sync.WaitGroup + errCh := make(chan error, workers) + for i := 0; i < workers; i++ { + wg.Add(1) + go func(disableGoTag bool) { + defer wg.Done() + for j := 0; j < 50; j++ { + opts := base + want := "message" + if disableGoTag { + opts.RemoveAnnotationMapper(thrift.AnnoScopeField, "go.tag") + want = "Msg" + } else { + opts.RegisterAnnotationMapper(thrift.AnnoScopeField, goTagMapper{}, "go.tag") + } + + p, err := GetDescFromContentWithOptions(content, "ExampleMethod", opts) + if err != nil { + errCh <- err + return + } + req := p.Request().Struct().Fields()[0].Type() + if got := req.Struct().FieldById(1).Alias(); got != want { + errCh <- fmt.Errorf("unexpected field alias: got %q, want %q", got, want) + return + } + } + }(i%2 == 0) + } + + wg.Wait() + close(errCh) + for err := range errCh { + require.NoError(t, err) + } + + p, err := GetDescFromContentWithOptions(content, "ExampleMethod", base) + require.NoError(t, err) + req := p.Request().Struct().Fields()[0].Type() + require.Equal(t, "message", req.Struct().FieldById(1).Alias()) +} + +func TestAnnotationMapperConcurrentDefaultMutationAndParse(t *testing.T) { + t.Cleanup(func() { + thrift.RegisterAnnotationMapper(thrift.AnnoScopeField, goTagMapper{}, "go.tag") + }) + content := ` + namespace go kitex.test.server + struct Base { + 1: string Msg (go.tag = "json:\"message\"") + } + + service InboxService { + string ExampleMethod(1: Base req) + } + ` + + var wg sync.WaitGroup + errCh := make(chan error, 512) + for i := 0; i < 8; i++ { + wg.Add(1) + go func(disableGoTag bool) { + defer wg.Done() + for j := 0; j < 50; j++ { + opts := thrift.NewDefaultOptions() + if disableGoTag { + opts.RemoveAnnotationMapper(thrift.AnnoScopeField, "go.tag") + } + _, err := GetDescFromContentWithOptions(content, "ExampleMethod", opts) + if err != nil { + errCh <- err + return + } + } + }(i%2 == 0) + } + + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 100; i++ { + thrift.RemoveAnnotationMapper(thrift.AnnoScopeField, "go.tag") + thrift.RegisterAnnotationMapper(thrift.AnnoScopeField, goTagMapper{}, "go.tag") + } + }() + + wg.Wait() + close(errCh) + for err := range errCh { + require.NoError(t, err) + } +} + func TestApiKey(t *testing.T) { p, err := GetDescFromContent(` namespace go kitex.test.server diff --git a/thrift/idl.go b/thrift/idl.go index 3df59b99..2cd3da1f 100644 --- a/thrift/idl.go +++ b/thrift/idl.go @@ -104,6 +104,8 @@ type Options struct { // ForceHashMapAsFieldNameMap indicates to use hash map as underlying field name map. // By default we try to use trie tree as field name map, which is usually faster than go map but consume more memory. ForceHashMapAsFieldNameMap bool + + annotationMappers *annotationMapperRegistry } // NewDefaultOptions creates a default Options. @@ -129,6 +131,7 @@ func NewDescritorFromPath(ctx context.Context, path string, includeDirs ...strin // NewDescritorFromContent creates a ServiceDescriptor from a thrift path and its includes, which uses the given options. // The includeDirs is used to find the include files. func (opts Options) NewDescritorFromPath(ctx context.Context, path string, includeDirs ...string) (*ServiceDescriptor, error) { + opts = opts.prepareAnnotationMappersForParse() tree, err := parser.ParseFile(path, includeDirs, true) if err != nil { return nil, err @@ -144,6 +147,7 @@ func (opts Options) NewDescritorFromPath(ctx context.Context, path string, inclu // If methods is empty, all methods will be parsed. // The includeDirs is used to find the include files. func (opts Options) NewDescriptorFromPathWithMethod(ctx context.Context, path string, includeDirs []string, methods ...string) (*ServiceDescriptor, error) { + opts = opts.prepareAnnotationMappersForParse() tree, err := parser.ParseFile(path, includeDirs, true) if err != nil { return nil, err @@ -165,6 +169,7 @@ func NewDescritorFromContent(ctx context.Context, path, content string, includes // includes is the thrift file content map, and its keys are specific including thrift file path. // isAbsIncludePath argument has become obsolete. Regardless of whether its value is true or false, both absolute path and relative path will be searched. func (opts Options) NewDescritorFromContent(ctx context.Context, path, content string, includes map[string]string, isAbsIncludePath bool) (*ServiceDescriptor, error) { + opts = opts.prepareAnnotationMappersForParse() tree, err := parseIDLContent(path, content, includes) if err != nil { return nil, err @@ -178,6 +183,7 @@ func (opts Options) NewDescritorFromContent(ctx context.Context, path, content s // NewDescritorFromContentWithMethod creates a ServiceDescriptor from a thrift content and its includes, but only parse specific methods. func (opts Options) NewDescriptorFromContentWithMethod(ctx context.Context, path, content string, includes map[string]string, isAbsIncludePath bool, methods ...string) (*ServiceDescriptor, error) { + opts = opts.prepareAnnotationMappersForParse() tree, err := parseIDLContent(path, content, includes) if err != nil { return nil, err @@ -955,6 +961,7 @@ func makeDefaultValue(typ *TypeDescriptor, val *parser.ConstValue, tree *parser. // file is the main thrift file path, name is the type name to parse (supports format like "package.TypeName" for cross-file references). // Returns a complete TypeDescriptor with all referenced types resolved. func (opts Options) NewDescriptorByName(ctx context.Context, file string, name string, includes map[string]string) (*TypeDescriptor, error) { + opts = opts.prepareAnnotationMappersForParse() // Parse the main IDL file and all includes tree, err := parseIDLContent(file, includes[file], includes) if err != nil {