protocol/grpc/server.go

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package grpc

import (
	"fmt"
	"net"
	"sync"
	"time"
)

import (
	"github.com/dubbogo/gost/log/logger"

	"github.com/dustin/go-humanize"

	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"

	"github.com/opentracing/opentracing-go"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	grpc_health "google.golang.org/grpc/health"
	grpc_health_v1 "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/reflection"
)

import (
	"dubbo.apache.org/dubbo-go/v3/common"
	"dubbo.apache.org/dubbo-go/v3/common/constant"
	"dubbo.apache.org/dubbo-go/v3/global"
	"dubbo.apache.org/dubbo-go/v3/protocol/base"
	dubbotls "dubbo.apache.org/dubbo-go/v3/tls"
)

// DubboGrpcService is gRPC service
type DubboGrpcService interface {
	// SetProxyImpl sets proxy.
	SetProxyImpl(impl base.Invoker)
	// GetProxyImpl gets proxy.
	GetProxyImpl() base.Invoker
	// ServiceDesc gets an RPC service's specification.
	ServiceDesc() *grpc.ServiceDesc
}

// Server is a gRPC server
type Server struct {
	grpcServer   *grpc.Server
	healthServer *grpc_health.Server
	bufferSize   int
	serviceLock  sync.Mutex
	services     map[string]struct{}
}

// NewServer creates a new server
func NewServer() *Server {
	return &Server{
		healthServer: grpc_health.NewServer(),
		services:     make(map[string]struct{}),
	}
}

func (s *Server) SetBufferSize(n int) {
	s.bufferSize = n
}

// Start gRPC server with @url
func (s *Server) Start(url *common.URL) {
	var (
		addr string
		err  error
	)
	addr = url.Location
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		panic(err)
	}
	success := false
	defer func() {
		if !success {
			_ = lis.Close()
		}
	}()

	maxServerRecvMsgSize := constant.DefaultMaxServerRecvMsgSize
	if recvMsgSize, convertErr := humanize.ParseBytes(url.GetParam(constant.MaxServerRecvMsgSize, "")); convertErr == nil && recvMsgSize != 0 {
		maxServerRecvMsgSize = int(recvMsgSize)
	}
	maxServerSendMsgSize := constant.DefaultMaxServerSendMsgSize
	if sendMsgSize, convertErr := humanize.ParseBytes(url.GetParam(constant.MaxServerSendMsgSize, "")); err == convertErr && sendMsgSize != 0 {
		maxServerSendMsgSize = int(sendMsgSize)
	}

	// If global trace instance was set, then server tracer instance
	// can be get. If not, will return NoopTracer.
	tracer := opentracing.GlobalTracer()
	var serverOpts []grpc.ServerOption
	serverOpts = append(serverOpts,
		grpc.UnaryInterceptor(otgrpc.OpenTracingServerInterceptor(tracer)),
		grpc.StreamInterceptor(otgrpc.OpenTracingStreamServerInterceptor(tracer)),
		grpc.MaxRecvMsgSize(maxServerRecvMsgSize),
		grpc.MaxSendMsgSize(maxServerSendMsgSize),
	)

	var transportCreds credentials.TransportCredentials
	transportCreds = insecure.NewCredentials()
	if tlsConfRaw, ok := url.GetAttribute(constant.TLSConfigKey); ok {
		// use global TLSConfig handle tls
		tlsConf, ok := tlsConfRaw.(*global.TLSConfig)
		if !ok {
			logger.Errorf("gRPC Server initialized the TLSConfig configuration failed")
			return
		}
		if dubbotls.IsServerTLSValid(tlsConf) {
			cfg, tlsErr := dubbotls.GetServerTlSConfig(tlsConf)
			if tlsErr != nil {
				return
			}
			if cfg != nil {
				logger.Infof("gRPC Server initialized the TLSConfig configuration")
				transportCreds = credentials.NewTLS(cfg)
			}
		}
	}
	serverOpts = append(serverOpts, grpc.Creds(transportCreds))

	server := grpc.NewServer(serverOpts...)
	s.grpcServer = server
	grpc_health_v1.RegisterHealthServer(server, s.healthServer)
	success = true

	go func() {
		providerServices := getProviderServices(url)
		if providerServices == nil {
			logger.Error("no provider service found")
			return
		}
		// wait all exporter ready, then set proxy impl and grpc registerService
		waitGrpcExporter(providerServices)
		registerService(providerServices, server)
		reflection.Register(server)
		if err := server.Serve(lis); err != nil {
			logger.Errorf("server serve failed with err: %v", err)
		}
	}()
}

func (s *Server) SetServingStatus(service string, status grpc_health_v1.HealthCheckResponse_ServingStatus) {
	if s.healthServer == nil || service == "" {
		return
	}

	s.serviceLock.Lock()
	s.services[service] = struct{}{}
	s.serviceLock.Unlock()
	s.healthServer.SetServingStatus(service, status)
}

func (s *Server) SetAllServicesNotServing() {
	if s.healthServer == nil {
		return
	}

	s.serviceLock.Lock()
	defer s.serviceLock.Unlock()
	for service := range s.services {
		s.healthServer.SetServingStatus(service, grpc_health_v1.HealthCheckResponse_NOT_SERVING)
	}
}

// getProviderServices retrieves provider services from URL attributes.
func getProviderServices(url *common.URL) map[string]*global.ServiceConfig {
	if providerConfRaw, ok := url.GetAttribute(constant.ProviderConfigKey); ok {
		if providerConf, ok := providerConfRaw.(*global.ProviderConfig); ok && providerConf != nil {
			return providerConf.Services
		}
		logger.Error("illegal provider config")
	}
	return nil
}

// getSyncMapLen get sync map len
func getSyncMapLen(m *sync.Map) int {
	length := 0

	m.Range(func(_, _ any) bool {
		length++
		return true
	})
	return length
}

// waitGrpcExporter wait until len(providerServices) = len(ExporterMap)
func waitGrpcExporter(providerServices map[string]*global.ServiceConfig) {
	t := time.NewTicker(50 * time.Millisecond)
	defer t.Stop()
	pLen := len(providerServices)
	ta := time.NewTimer(10 * time.Second)
	defer ta.Stop()

	for {
		select {
		case <-t.C:
			mLen := getSyncMapLen(grpcProtocol.ExporterMap())
			if pLen == mLen {
				return
			}
		case <-ta.C:
			panic("wait grpc exporter timeout when start grpc server")
		}
	}
}

// registerService SetProxyImpl invoker and grpc service
func registerService(providerServices map[string]*global.ServiceConfig, server *grpc.Server) {
	for _, providerService := range providerServices {
		serviceKey := common.ServiceKey(providerService.Interface, providerService.Group, providerService.Version)
		exporter, _ := grpcProtocol.ExporterMap().Load(serviceKey)
		if exporter == nil {
			panic(fmt.Sprintf("no exporter found for servicekey: %v", serviceKey))
		}
		invoker := exporter.(base.Exporter).GetInvoker()
		if invoker == nil {
			panic(fmt.Sprintf("no invoker found for servicekey: %v", serviceKey))
		}
		service, ok := invoker.GetURL().GetAttribute(constant.RpcServiceKey)
		if !ok {
			panic(fmt.Sprintf("no rpc service found in url attribute %s for servicekey: %v", constant.RpcServiceKey, serviceKey))
		}
		ds, ok := service.(DubboGrpcService)
		if !ok {
			panic("illegal service type registered")
		}

		ds.SetProxyImpl(invoker)
		server.RegisterService(ds.ServiceDesc(), ds)
	}
}

// Stop gRPC server
func (s *Server) Stop() {
	s.grpcServer.Stop()
}

// GracefulStop gRPC server
func (s *Server) GracefulStop() {
	s.grpcServer.GracefulStop()
}