github.com/cloudwego/kitex@v0.9.0/pkg/generic/pbidl_provider.go (about)

     1  /*
     2   * Copyright 2021 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 generic
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"sync"
    23  
    24  	dproto "github.com/cloudwego/dynamicgo/proto"
    25  	"github.com/jhump/protoreflect/desc/protoparse"
    26  
    27  	"github.com/cloudwego/kitex/pkg/generic/proto"
    28  )
    29  
    30  type PbContentProvider struct {
    31  	closeOnce sync.Once
    32  	svcs      chan proto.ServiceDescriptor
    33  }
    34  
    35  type PbFileProviderWithDynamicGo struct {
    36  	closeOnce sync.Once
    37  	svcs      chan *dproto.ServiceDescriptor
    38  }
    39  
    40  var (
    41  	_ PbDescriptorProvider          = (*PbContentProvider)(nil)
    42  	_ PbDescriptorProviderDynamicGo = (*PbFileProviderWithDynamicGo)(nil)
    43  )
    44  
    45  func NewPbContentProvider(main string, includes map[string]string) (PbDescriptorProvider, error) {
    46  	p := &PbContentProvider{
    47  		svcs: make(chan proto.ServiceDescriptor, 1),
    48  	}
    49  
    50  	sd, err := parseProto(main, includes)
    51  	if err != nil {
    52  		return nil, err
    53  	}
    54  	p.svcs <- sd
    55  
    56  	return p, nil
    57  }
    58  
    59  func (p *PbContentProvider) UpdateIDL(main string, includes map[string]string) error {
    60  	sd, err := parseProto(main, includes)
    61  	if err != nil {
    62  		return err
    63  	}
    64  
    65  	select {
    66  	case <-p.svcs:
    67  	default:
    68  	}
    69  
    70  	select {
    71  	case p.svcs <- sd:
    72  	default:
    73  	}
    74  
    75  	return nil
    76  }
    77  
    78  func parseProto(main string, includes map[string]string) (proto.ServiceDescriptor, error) {
    79  	var pbParser protoparse.Parser
    80  	pbParser.Accessor = protoparse.FileContentsFromMap(includes)
    81  	pbParser.IncludeSourceCodeInfo = true
    82  	pbParser.ValidateUnlinkedFiles = true
    83  	pbParser.InterpretOptionsInUnlinkedFiles = true
    84  	fds, err := pbParser.ParseFiles(main)
    85  	if err != nil {
    86  		return nil, err
    87  	}
    88  	if len(fds) < 1 {
    89  		return nil, errors.New("no file descriptor found")
    90  	}
    91  
    92  	services := fds[0].GetServices()
    93  	if len(services) < 1 {
    94  		return nil, errors.New("no service descriptor found")
    95  	}
    96  
    97  	return services[0], nil
    98  }
    99  
   100  func (p *PbContentProvider) Provide() <-chan proto.ServiceDescriptor {
   101  	return p.svcs
   102  }
   103  
   104  func (p *PbContentProvider) Close() error {
   105  	p.closeOnce.Do(func() {
   106  		close(p.svcs)
   107  	})
   108  	return nil
   109  }
   110  
   111  // PbFileProviderWithDynamicGo
   112  func NewPbFileProviderWithDynamicGo(main string, ctx context.Context, options dproto.Options) (PbDescriptorProviderDynamicGo, error) {
   113  	p := &PbFileProviderWithDynamicGo{
   114  		svcs: make(chan *dproto.ServiceDescriptor, 1),
   115  	}
   116  
   117  	svc, err := options.NewDescriptorFromPath(ctx, main)
   118  	if err != nil {
   119  		return nil, err
   120  	}
   121  	p.svcs <- svc
   122  
   123  	return p, nil
   124  }
   125  
   126  func (p *PbFileProviderWithDynamicGo) Provide() <-chan *dproto.ServiceDescriptor {
   127  	return p.svcs
   128  }
   129  
   130  func (p *PbFileProviderWithDynamicGo) Close() error {
   131  	p.closeOnce.Do(func() {
   132  		close(p.svcs)
   133  	})
   134  	return nil
   135  }