server/service.go

/*
 * Copyright 2023 CloudWeGo Authors
 *
 * Licensed 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 server

import (
	"errors"
	"fmt"

	"github.com/cloudwego/kitex/pkg/endpoint"

	"github.com/cloudwego/kitex/pkg/serviceinfo"
)

type service struct {
	svcInfo *serviceinfo.ServiceInfo
	handler interface{}
	mw      endpoint.Middleware
}

func newService(svcInfo *serviceinfo.ServiceInfo, handler interface{}, mw endpoint.Middleware) *service {
	return &service{svcInfo: svcInfo, handler: handler, mw: mw}
}

type services struct {
	methodSvcsMap map[string][]*service // key: method name
	svcMap        map[string]*service   // key: service name
	fallbackSvc   *service

	refuseTrafficWithoutServiceName bool
}

func newServices() *services {
	return &services{
		methodSvcsMap: map[string][]*service{},
		svcMap:        map[string]*service{},
	}
}

func (s *services) addService(svcInfo *serviceinfo.ServiceInfo, handler interface{}, registerOpts *RegisterOptions) error {
	var serviceMW endpoint.Middleware
	if len(registerOpts.Middlewares) > 0 {
		serviceMW = endpoint.Chain(registerOpts.Middlewares...)
	}
	svc := newService(svcInfo, handler, serviceMW)
	if registerOpts.IsFallbackService {
		if s.fallbackSvc != nil {
			return fmt.Errorf("multiple fallback services cannot be registered. [%s] is already registered as a fallback service", s.fallbackSvc.svcInfo.ServiceName)
		}
		s.fallbackSvc = svc
	}
	s.svcMap[svcInfo.ServiceName] = svc
	// method search map
	for methodName := range svcInfo.Methods {
		svcs := s.methodSvcsMap[methodName]
		if registerOpts.IsFallbackService {
			svcs = append([]*service{svc}, svcs...)
		} else {
			svcs = append(svcs, svc)
		}
		s.methodSvcsMap[methodName] = svcs
	}
	return nil
}

func (s *services) getSvcInfoMap() map[string]*serviceinfo.ServiceInfo {
	svcInfoMap := map[string]*serviceinfo.ServiceInfo{}
	for name, svc := range s.svcMap {
		svcInfoMap[name] = svc.svcInfo
	}
	return svcInfoMap
}

func (s *services) check(refuseTrafficWithoutServiceName bool) error {
	if len(s.svcMap) == 0 {
		return errors.New("run: no service. Use RegisterService to set one")
	}
	if refuseTrafficWithoutServiceName {
		s.refuseTrafficWithoutServiceName = true
		return nil
	}
	for name, svcs := range s.methodSvcsMap {
		if len(svcs) > 1 && svcs[0] != s.fallbackSvc {
			return fmt.Errorf("method name [%s] is conflicted between services but no fallback service is specified", name)
		}
	}
	return nil
}

func (s *services) SearchService(svcName, methodName string, strict bool) *serviceinfo.ServiceInfo {
	if strict || s.refuseTrafficWithoutServiceName {
		if svc := s.svcMap[svcName]; svc != nil {
			return svc.svcInfo
		}
		return nil
	}
	var svc *service
	if svcName == "" {
		if svcs := s.methodSvcsMap[methodName]; len(svcs) > 0 {
			svc = svcs[0]
		}
	} else {
		svc = s.svcMap[svcName]
		if svc == nil {
			svcs := s.methodSvcsMap[methodName]
			// 1. no conflicting method, allow method routing
			// 2. $ generally means generic service, maybe mismatch the real service name
			if len(svcs) == 1 || len(svcs) > 1 && svcName[0] == '$' {
				svc = svcs[0]
			}
		}
	}
	if svc != nil {
		return svc.svcInfo
	}
	return nil
}