github.com/cloudwego/kitex@v0.9.0/pkg/remote/message.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 remote

import (
	"fmt"
	"sync"

	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/utils"
	"github.com/cloudwego/kitex/transport"
)

var (
	messagePool      sync.Pool
	transInfoPool    sync.Pool
	emptyServiceInfo serviceinfo.ServiceInfo
)

func init() {
	messagePool.New = newMessage
	transInfoPool.New = newTransInfo
}

// MessageType indicates the type of message.
type MessageType int32

// MessageTypes.
const (
	// 0-4 corresponding to thrift.TMessageType
	InvalidMessageType MessageType = 0
	Call               MessageType = 1
	Reply              MessageType = 2
	Exception          MessageType = 3
	// Oneway means there's no need to wait for the response.
	// When the actual message is transmitted, Oneway writes a Call to avoid compatibility issues
	// and to maintain consistency with the original logic.
	Oneway MessageType = 4

	Stream MessageType = 5

	Heartbeat MessageType = 6
)

const (
	// ReadFailed .
	ReadFailed string = "RFailed"

	// MeshHeader use in message.Tag to check MeshHeader
	MeshHeader string = "mHeader"
)

var emptyProtocolInfo ProtocolInfo

// ProtocolInfo is used to indicate the transport protocol and payload codec information.
type ProtocolInfo struct {
	TransProto transport.Protocol
	CodecType  serviceinfo.PayloadCodec
}

// NewProtocolInfo creates a new ProtocolInfo using the given tp and ct.
func NewProtocolInfo(tp transport.Protocol, ct serviceinfo.PayloadCodec) ProtocolInfo {
	return ProtocolInfo{
		TransProto: tp,
		CodecType:  ct,
	}
}

// Message is the core abstraction for Kitex message.
type Message interface {
	RPCInfo() rpcinfo.RPCInfo
	ServiceInfo() *serviceinfo.ServiceInfo
	SpecifyServiceInfo(svcName, methodName string) (*serviceinfo.ServiceInfo, error)
	Data() interface{}
	NewData(method string) (ok bool)
	MessageType() MessageType
	SetMessageType(MessageType)
	RPCRole() RPCRole
	PayloadLen() int
	SetPayloadLen(size int)
	TransInfo() TransInfo
	Tags() map[string]interface{}
	ProtocolInfo() ProtocolInfo
	SetProtocolInfo(ProtocolInfo)
	PayloadCodec() PayloadCodec
	SetPayloadCodec(pc PayloadCodec)
	Recycle()
}

// NewMessage creates a new Message using the given info.
func NewMessage(data interface{}, svcInfo *serviceinfo.ServiceInfo, ri rpcinfo.RPCInfo, msgType MessageType, rpcRole RPCRole) Message {
	msg := messagePool.Get().(*message)
	msg.data = data
	msg.rpcInfo = ri
	msg.targetSvcInfo = svcInfo
	msg.msgType = msgType
	msg.rpcRole = rpcRole
	msg.transInfo = transInfoPool.Get().(*transInfo)
	return msg
}

// NewMessageWithNewer creates a new Message and set data later.
func NewMessageWithNewer(targetSvcInfo *serviceinfo.ServiceInfo, svcSearchMap map[string]*serviceinfo.ServiceInfo, ri rpcinfo.RPCInfo, msgType MessageType, rpcRole RPCRole, refuseTrafficWithoutServiceName bool) Message {
	msg := messagePool.Get().(*message)
	msg.rpcInfo = ri
	msg.targetSvcInfo = targetSvcInfo
	msg.svcSearchMap = svcSearchMap
	msg.msgType = msgType
	msg.rpcRole = rpcRole
	msg.transInfo = transInfoPool.Get().(*transInfo)
	msg.refuseTrafficWithoutServiceName = refuseTrafficWithoutServiceName
	return msg
}

// RecycleMessage is used to recycle message.
func RecycleMessage(msg Message) {
	if msg != nil {
		msg.Recycle()
	}
}

func newMessage() interface{} {
	return &message{tags: make(map[string]interface{})}
}

type message struct {
	msgType                         MessageType
	data                            interface{}
	rpcInfo                         rpcinfo.RPCInfo
	targetSvcInfo                   *serviceinfo.ServiceInfo
	svcSearchMap                    map[string]*serviceinfo.ServiceInfo
	rpcRole                         RPCRole
	compressType                    CompressType
	payloadSize                     int
	transInfo                       TransInfo
	tags                            map[string]interface{}
	protocol                        ProtocolInfo
	payloadCodec                    PayloadCodec
	refuseTrafficWithoutServiceName bool
}

func (m *message) zero() {
	m.msgType = InvalidMessageType
	m.data = nil
	m.rpcInfo = nil
	m.targetSvcInfo = &emptyServiceInfo
	m.rpcRole = -1
	m.compressType = NoCompress
	m.payloadSize = 0
	if m.transInfo != nil {
		m.transInfo.Recycle()
		m.transInfo = nil
	}
	for k := range m.tags {
		delete(m.tags, k)
	}
	m.protocol = emptyProtocolInfo
}

// RPCInfo implements the Message interface.
func (m *message) RPCInfo() rpcinfo.RPCInfo {
	return m.rpcInfo
}

// ServiceInfo implements the Message interface.
func (m *message) ServiceInfo() *serviceinfo.ServiceInfo {
	return m.targetSvcInfo
}

func (m *message) SpecifyServiceInfo(svcName, methodName string) (*serviceinfo.ServiceInfo, error) {
	// for non-multi-service including generic server scenario
	if m.targetSvcInfo != nil {
		if mt := m.targetSvcInfo.MethodInfo(methodName); mt == nil {
			return nil, NewTransErrorWithMsg(UnknownMethod, fmt.Sprintf("unknown method %s", methodName))
		}
		return m.targetSvcInfo, nil
	}
	if svcName == "" && m.refuseTrafficWithoutServiceName {
		return nil, NewTransErrorWithMsg(NoServiceName, "no service name while the server has WithRefuseTrafficWithoutServiceName option enabled")
	}
	var key string
	if svcName == "" {
		key = methodName
	} else {
		key = BuildMultiServiceKey(svcName, methodName)
	}
	svcInfo := m.svcSearchMap[key]
	if svcInfo == nil {
		return nil, NewTransErrorWithMsg(UnknownMethod, fmt.Sprintf("unknown method %s", methodName))
	}
	m.targetSvcInfo = svcInfo
	return svcInfo, nil
}

// Data implements the Message interface.
func (m *message) Data() interface{} {
	return m.data
}

// NewData implements the Message interface.
func (m *message) NewData(method string) (ok bool) {
	if m.data != nil {
		return false
	}
	if mt := m.targetSvcInfo.MethodInfo(method); mt != nil {
		m.data = mt.NewArgs()
	}
	if m.data == nil {
		return false
	}
	return true
}

// MessageType implements the Message interface.
func (m *message) MessageType() MessageType {
	return m.msgType
}

// SetMessageType implements the Message interface.
func (m *message) SetMessageType(mt MessageType) {
	m.msgType = mt
}

// RPCRole implements the Message interface.
func (m *message) RPCRole() RPCRole {
	return m.rpcRole
}

// TransInfo implements the Message interface.
func (m *message) TransInfo() TransInfo {
	return m.transInfo
}

// Tags implements the Message interface.
func (m *message) Tags() map[string]interface{} {
	return m.tags
}

// PayloadLen implements the Message interface.
func (m *message) PayloadLen() int {
	return m.payloadSize
}

// SetPayloadLen implements the Message interface.
func (m *message) SetPayloadLen(size int) {
	m.payloadSize = size
}

// ProtocolInfo implements the Message interface.
func (m *message) ProtocolInfo() ProtocolInfo {
	return m.protocol
}

// SetProtocolInfo implements the Message interface.
func (m *message) SetProtocolInfo(pt ProtocolInfo) {
	m.protocol = pt
}

// PayloadCodec implements the Message interface.
func (m *message) PayloadCodec() PayloadCodec {
	return m.payloadCodec
}

// SetPayloadCodec implements the Message interface.
func (m *message) SetPayloadCodec(pc PayloadCodec) {
	m.payloadCodec = pc
}

// Recycle is used to recycle the message.
func (m *message) Recycle() {
	m.zero()
	messagePool.Put(m)
}

// TransInfo contains transport information.
type TransInfo interface {
	TransStrInfo() map[string]string
	TransIntInfo() map[uint16]string
	PutTransIntInfo(map[uint16]string)
	PutTransStrInfo(kvInfo map[string]string)
	Recycle()
}

func newTransInfo() interface{} {
	return &transInfo{
		intInfo: make(map[uint16]string),
		strInfo: make(map[string]string),
	}
}

type transInfo struct {
	strInfo map[string]string
	intInfo map[uint16]string
}

func (ti *transInfo) zero() {
	for k := range ti.intInfo {
		delete(ti.intInfo, k)
	}
	for k := range ti.strInfo {
		delete(ti.strInfo, k)
	}
}

// TransIntInfo implements the TransInfo interface.
func (ti *transInfo) TransIntInfo() map[uint16]string {
	return ti.intInfo
}

// PutTransIntInfo implements the TransInfo interface.
func (ti *transInfo) PutTransIntInfo(kvInfo map[uint16]string) {
	if kvInfo == nil {
		return
	}
	if len(ti.intInfo) == 0 {
		ti.intInfo = kvInfo
	} else {
		for k, v := range kvInfo {
			ti.intInfo[k] = v
		}
	}
}

// TransStrInfo implements the TransInfo interface.
func (ti *transInfo) TransStrInfo() map[string]string {
	return ti.strInfo
}

// PutTransStrInfo implements the TransInfo interface.
func (ti *transInfo) PutTransStrInfo(kvInfo map[string]string) {
	if kvInfo == nil {
		return
	}
	if len(ti.strInfo) == 0 {
		ti.strInfo = kvInfo
	} else {
		for k, v := range kvInfo {
			ti.strInfo[k] = v
		}
	}
}

// Recycle is used to recycle the transInfo.
func (ti *transInfo) Recycle() {
	ti.zero()
	transInfoPool.Put(ti)
}

// FillSendMsgFromRecvMsg is used to fill the transport information to the message to be sent.
func FillSendMsgFromRecvMsg(recvMsg, sendMsg Message) {
	sendMsg.SetProtocolInfo(recvMsg.ProtocolInfo())
	sendMsg.SetPayloadCodec(recvMsg.PayloadCodec())
}

// BuildMultiServiceKey is used to create a key to search svcInfo from svcSearchMap.
func BuildMultiServiceKey(serviceName, methodName string) string {
	var builder utils.StringBuilder
	builder.Grow(len(serviceName) + len(methodName) + 1)
	builder.WriteString(serviceName)
	builder.WriteString(".")
	builder.WriteString(methodName)
	return builder.String()
}