github.com/erda-project/erda-infra@v1.0.9/providers/grpcserver/provider.go (about)

     1  // Copyright (c) 2021 Terminus, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package grpcserver
    16  
    17  import (
    18  	"fmt"
    19  	"net"
    20  	"reflect"
    21  
    22  	"github.com/erda-project/erda-infra/base/logs"
    23  	"github.com/erda-project/erda-infra/base/servicehub"
    24  	grpccontext "github.com/erda-project/erda-infra/pkg/trace/inject/context/grpc"
    25  	"google.golang.org/grpc"
    26  	"google.golang.org/grpc/credentials"
    27  )
    28  
    29  // Interface .
    30  type Interface interface {
    31  	RegisterService(sd *grpc.ServiceDesc, ss interface{})
    32  }
    33  
    34  // config .
    35  type config struct {
    36  	Addr string `file:"addr" default:":7070" desc:"grpc address to listen"`
    37  	TLS  struct {
    38  		CertFile string `file:"cert_file" desc:"the TLS cert file"`
    39  		KeyFile  string `file:"key_file" desc:"the TLS key file"`
    40  	} `file:"tls"`
    41  	TraceEnable bool `file:"trace_enable" default:"true"`
    42  }
    43  
    44  type provider struct {
    45  	Cfg    *config
    46  	Log    logs.Logger
    47  	listen net.Listener
    48  	server *grpc.Server
    49  }
    50  
    51  func (p *provider) Init(ctx servicehub.Context) error {
    52  	lis, err := net.Listen("tcp", p.Cfg.Addr)
    53  	if err != nil {
    54  		return err
    55  	}
    56  	p.listen = lis
    57  
    58  	var opts []grpc.ServerOption
    59  	if len(p.Cfg.TLS.CertFile) > 0 || len(p.Cfg.TLS.KeyFile) > 0 {
    60  		creds, err := credentials.NewServerTLSFromFile(p.Cfg.TLS.CertFile, p.Cfg.TLS.KeyFile)
    61  		if err != nil {
    62  			return fmt.Errorf("fail to generate credentials %v", err)
    63  		}
    64  		opts = append(opts, grpc.Creds(creds))
    65  	}
    66  	if p.Cfg.TraceEnable {
    67  		opts = append(opts,
    68  			grpc.UnaryInterceptor(grpccontext.UnaryServerInterceptor()),
    69  			grpc.StreamInterceptor(grpccontext.StreamServerInterceptor()),
    70  		)
    71  	}
    72  	p.server = grpc.NewServer(opts...)
    73  	return nil
    74  }
    75  
    76  func (p *provider) Start() error {
    77  	p.Log.Infof("starting grpc server at %s", p.Cfg.Addr)
    78  	return p.server.Serve(p.listen)
    79  }
    80  
    81  func (p *provider) Close() error {
    82  	p.server.Stop()
    83  	return nil
    84  }
    85  
    86  func (p *provider) Provide(ctx servicehub.DependencyContext, args ...interface{}) interface{} {
    87  	return p.server
    88  }
    89  
    90  func init() {
    91  	servicehub.Register("grpc-server", &servicehub.Spec{
    92  		Services: []string{"grpc-server"},
    93  		Types: []reflect.Type{
    94  			reflect.TypeOf((*grpc.Server)(nil)),
    95  			reflect.TypeOf((*Interface)(nil)).Elem(),
    96  		},
    97  		Description: "grpc server",
    98  		ConfigFunc:  func() interface{} { return &config{} },
    99  		Creator: func() servicehub.Provider {
   100  			return &provider{}
   101  		},
   102  	})
   103  }