github.com/cloudwego/kitex@v0.9.0/pkg/reflection/thrift/registry.go (about)

     1  /*
     2   * Copyright 2023 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package thriftreflection
    18  
    19  import (
    20  	"sync"
    21  
    22  	"github.com/cloudwego/thriftgo/reflection"
    23  )
    24  
    25  // Files is a registry for looking up or iterating over files and the
    26  // descriptors contained within them.
    27  type Files struct {
    28  	filesByPath map[string]*reflection.FileDescriptor
    29  }
    30  
    31  func NewFiles() *Files {
    32  	return &Files{filesByPath: map[string]*reflection.FileDescriptor{}}
    33  }
    34  
    35  var globalMutex sync.RWMutex
    36  
    37  // GlobalFiles is a global registry of file descriptors.
    38  var GlobalFiles = NewFiles()
    39  
    40  // RegisterIDL provides function for generated code to register their reflection info to GlobalFiles.
    41  func RegisterIDL(bytes []byte) {
    42  	desc := reflection.Decode(bytes)
    43  	GlobalFiles.Register(desc)
    44  }
    45  
    46  // Register registers the input FileDescriptor to *Files type variables.
    47  func (f *Files) Register(desc *reflection.FileDescriptor) {
    48  	if f == GlobalFiles {
    49  		globalMutex.Lock()
    50  		defer globalMutex.Unlock()
    51  	}
    52  	// TODO: check conflict
    53  	f.filesByPath[desc.Filename] = desc
    54  }
    55  
    56  // GetFileDescriptors returns the inner registered reflection FileDescriptors.
    57  func (f *Files) GetFileDescriptors() map[string]*reflection.FileDescriptor {
    58  	if f == GlobalFiles {
    59  		globalMutex.RLock()
    60  		defer globalMutex.RUnlock()
    61  		m := make(map[string]*reflection.FileDescriptor, len(f.filesByPath))
    62  		for k, v := range f.filesByPath {
    63  			m[k] = v
    64  		}
    65  		return m
    66  	}
    67  	return f.filesByPath
    68  }