github.com/cloudwego/kitex@v0.9.0/pkg/serviceinfo/serviceinfo.go

/*
 * Copyright 2021 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 serviceinfo

import (
	"context"
)

// PayloadCodec alias type
type PayloadCodec int

// PayloadCodec supported by kitex.
const (
	Thrift PayloadCodec = iota
	Protobuf
	Hessian2
)

const (
	// GenericService name
	GenericService = "$GenericService" // private as "$"
	// GenericMethod name
	GenericMethod = "$GenericCall"
)

// ServiceInfo to record meta info of service
type ServiceInfo struct {
	// deprecated, for compatibility
	PackageName string

	// The name of the service. For generic services, it is always the constant `GenericService`.
	ServiceName string

	// HandlerType is the type value of a request handler from the generated code.
	HandlerType interface{}

	// Methods contains the meta information of methods supported by the service.
	// For generic service, there is only one method named by the constant `GenericMethod`.
	Methods map[string]MethodInfo

	// PayloadCodec is the codec of payload.
	PayloadCodec PayloadCodec

	// KiteXGenVersion is the version of command line tool 'kitex'.
	KiteXGenVersion string

	// Extra is for future feature info, to avoid compatibility issue
	// as otherwise we need to add a new field in the struct
	Extra map[string]interface{}

	// GenericMethod returns a MethodInfo for the given name.
	// It is used by generic calls only.
	GenericMethod func(name string) MethodInfo
}

// GetPackageName returns the PackageName.
// The return value is generally taken from the Extra field of ServiceInfo with a key
// `PackageName`. For some legacy generated code, it may come from the PackageName field.
func (i *ServiceInfo) GetPackageName() (pkg string) {
	if i.PackageName != "" {
		return i.PackageName
	}
	pkg, _ = i.Extra["PackageName"].(string)
	return
}

// MethodInfo gets MethodInfo.
func (i *ServiceInfo) MethodInfo(name string) MethodInfo {
	if i == nil {
		return nil
	}
	if i.ServiceName == GenericService {
		if i.GenericMethod != nil {
			return i.GenericMethod(name)
		}
		return i.Methods[GenericMethod]
	}
	return i.Methods[name]
}

type StreamingMode int

const (
	StreamingNone          StreamingMode = 0b0000 // KitexPB/KitexThrift
	StreamingUnary         StreamingMode = 0b0001 // underlying protocol for streaming, like HTTP2
	StreamingClient        StreamingMode = 0b0010
	StreamingServer        StreamingMode = 0b0100
	StreamingBidirectional StreamingMode = 0b0110
)

// MethodInfo to record meta info of unary method
type MethodInfo interface {
	Handler() MethodHandler
	NewArgs() interface{}
	NewResult() interface{}
	OneWay() bool
	IsStreaming() bool
	StreamingMode() StreamingMode
}

// MethodHandler is corresponding to the handler wrapper func that in generated code
type MethodHandler func(ctx context.Context, handler, args, result interface{}) error

// MethodInfoOption modifies the given MethodInfo
type MethodInfoOption func(*methodInfo)

// WithStreamingMode sets the streaming mode and update the streaming flag accordingly
func WithStreamingMode(mode StreamingMode) MethodInfoOption {
	return func(m *methodInfo) {
		m.streamingMode = mode
		m.isStreaming = mode != StreamingNone
	}
}

// NewMethodInfo is called in generated code to build method info
func NewMethodInfo(methodHandler MethodHandler, newArgsFunc, newResultFunc func() interface{}, oneWay bool, opts ...MethodInfoOption) MethodInfo {
	mi := methodInfo{
		handler:       methodHandler,
		newArgsFunc:   newArgsFunc,
		newResultFunc: newResultFunc,
		oneWay:        oneWay,
		isStreaming:   false,
		streamingMode: StreamingNone,
	}
	for _, opt := range opts {
		opt(&mi)
	}
	return &mi
}

type methodInfo struct {
	handler       MethodHandler
	newArgsFunc   func() interface{}
	newResultFunc func() interface{}
	oneWay        bool
	isStreaming   bool
	streamingMode StreamingMode
}

// Handler implements the MethodInfo interface.
func (m methodInfo) Handler() MethodHandler {
	return m.handler
}

// NewArgs implements the MethodInfo interface.
func (m methodInfo) NewArgs() interface{} {
	return m.newArgsFunc()
}

// NewResult implements the MethodInfo interface.
func (m methodInfo) NewResult() interface{} {
	return m.newResultFunc()
}

// OneWay implements the MethodInfo interface.
func (m methodInfo) OneWay() bool {
	return m.oneWay
}

func (m methodInfo) IsStreaming() bool {
	return m.isStreaming
}

func (m methodInfo) StreamingMode() StreamingMode {
	return m.streamingMode
}

// String prints human-readable information.
func (p PayloadCodec) String() string {
	switch p {
	case Thrift:
		return "Thrift"
	case Protobuf:
		return "Protobuf"
	case Hessian2:
		return "Hessian2"
	}
	panic("unknown payload type")
}