github.com/cloudwego/kitex@v0.9.0/server/service.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 server
    18  
    19  import (
    20  	"errors"
    21  	"fmt"
    22  
    23  	"github.com/cloudwego/kitex/pkg/remote"
    24  	"github.com/cloudwego/kitex/pkg/serviceinfo"
    25  )
    26  
    27  type service struct {
    28  	svcInfo *serviceinfo.ServiceInfo
    29  	handler interface{}
    30  }
    31  
    32  func newService(svcInfo *serviceinfo.ServiceInfo, handler interface{}) *service {
    33  	return &service{svcInfo: svcInfo, handler: handler}
    34  }
    35  
    36  type services struct {
    37  	svcSearchMap                       map[string]*service // key: "svcName.methodName" and "methodName", value: svcInfo
    38  	svcMap                             map[string]*service // key: service name, value: svcInfo
    39  	conflictingMethodHasFallbackSvcMap map[string]bool
    40  	fallbackSvc                        *service
    41  }
    42  
    43  func newServices() *services {
    44  	return &services{
    45  		svcSearchMap:                       map[string]*service{},
    46  		svcMap:                             map[string]*service{},
    47  		conflictingMethodHasFallbackSvcMap: map[string]bool{},
    48  	}
    49  }
    50  
    51  func (s *services) addService(svcInfo *serviceinfo.ServiceInfo, handler interface{}, registerOpts *RegisterOptions) error {
    52  	svc := newService(svcInfo, handler)
    53  
    54  	if err := s.checkCombineServiceWithOtherService(svcInfo); err != nil {
    55  		return err
    56  	}
    57  
    58  	if err := s.checkMultipleFallbackService(registerOpts, svc); err != nil {
    59  		return err
    60  	}
    61  
    62  	s.svcMap[svcInfo.ServiceName] = svc
    63  	s.createSearchMap(svcInfo, svc, registerOpts)
    64  	return nil
    65  }
    66  
    67  // when registering combine service, it does not allow the registration of other services
    68  func (s *services) checkCombineServiceWithOtherService(svcInfo *serviceinfo.ServiceInfo) error {
    69  	if len(s.svcMap) > 0 {
    70  		if _, ok := s.svcMap["CombineService"]; ok || svcInfo.ServiceName == "CombineService" {
    71  			return errors.New("only one service can be registered when registering combine service")
    72  		}
    73  	}
    74  	return nil
    75  }
    76  
    77  func (s *services) checkMultipleFallbackService(registerOpts *RegisterOptions, svc *service) error {
    78  	if registerOpts.IsFallbackService {
    79  		if s.fallbackSvc != nil {
    80  			return fmt.Errorf("multiple fallback services cannot be registered. [%s] is already registered as a fallback service", s.fallbackSvc.svcInfo.ServiceName)
    81  		}
    82  		s.fallbackSvc = svc
    83  	}
    84  	return nil
    85  }
    86  
    87  func (s *services) createSearchMap(svcInfo *serviceinfo.ServiceInfo, svc *service, registerOpts *RegisterOptions) {
    88  	for methodName := range svcInfo.Methods {
    89  		s.svcSearchMap[remote.BuildMultiServiceKey(svcInfo.ServiceName, methodName)] = svc
    90  		if svcFromMap, ok := s.svcSearchMap[methodName]; ok {
    91  			s.handleConflictingMethod(svcFromMap, svc, methodName, registerOpts)
    92  		} else {
    93  			s.svcSearchMap[methodName] = svc
    94  		}
    95  	}
    96  }
    97  
    98  func (s *services) handleConflictingMethod(svcFromMap, svc *service, methodName string, registerOpts *RegisterOptions) {
    99  	s.registerConflictingMethodHasFallbackSvcMap(svcFromMap, methodName)
   100  	s.updateWithFallbackSvc(registerOpts, svc, methodName)
   101  }
   102  
   103  func (s *services) registerConflictingMethodHasFallbackSvcMap(svcFromMap *service, methodName string) {
   104  	if _, ok := s.conflictingMethodHasFallbackSvcMap[methodName]; !ok {
   105  		if s.fallbackSvc != nil && svcFromMap.svcInfo.ServiceName == s.fallbackSvc.svcInfo.ServiceName {
   106  			// svc which is already registered is a fallback service
   107  			s.conflictingMethodHasFallbackSvcMap[methodName] = true
   108  		} else {
   109  			s.conflictingMethodHasFallbackSvcMap[methodName] = false
   110  		}
   111  	}
   112  }
   113  
   114  func (s *services) updateWithFallbackSvc(registerOpts *RegisterOptions, svc *service, methodName string) {
   115  	if registerOpts.IsFallbackService {
   116  		s.svcSearchMap[methodName] = svc
   117  		s.conflictingMethodHasFallbackSvcMap[methodName] = true
   118  	}
   119  }
   120  
   121  func (s *services) getSvcInfoMap() map[string]*serviceinfo.ServiceInfo {
   122  	svcInfoMap := map[string]*serviceinfo.ServiceInfo{}
   123  	for name, svc := range s.svcMap {
   124  		svcInfoMap[name] = svc.svcInfo
   125  	}
   126  	return svcInfoMap
   127  }
   128  
   129  func (s *services) getSvcInfoSearchMap() map[string]*serviceinfo.ServiceInfo {
   130  	svcInfoSearchMap := map[string]*serviceinfo.ServiceInfo{}
   131  	for name, svc := range s.svcSearchMap {
   132  		svcInfoSearchMap[name] = svc.svcInfo
   133  	}
   134  	return svcInfoSearchMap
   135  }