Repository: cloudwego/kitex
File: pkg/remote/codec/thrift/thrift.go (Go)

Code before:

	data, err := getValidData(methodName, message)
	if err != nil {
		return err
	}

	// encode with hyper codec
	// NOTE: to ensure hyperMarshalEnabled is inlined so split the check logic, or it may cause performance loss
	if c.hyperMarshalEnabled() && hyperMarshalAvailable(data) {
		return c.hyperMarshal(out, methodName, msgType, seqID, data)
	}

	// encode with FastWrite
	if c.CodecType&FastWrite != 0 {
		if msg, ok := data.(ThriftMsgFastCodec); ok {
			return encodeFastThrift(out, methodName, msgType, seqID, msg)
		}
	}

	if hyperMarshalFallback(data) {
		// fallback to frugal
		return c.hyperMarshal(out, methodName, msgType, seqID, data)
	}

	// fallback to old thrift way (slow)
	return encodeBasicThrift(out, ctx, methodName, msgType, seqID, data)
}

// encodeFastThrift encode with the FastCodec way
func encodeFastThrift(out remote.ByteBuffer, methodName string, msgType remote.MessageType, seqID int32, msg ThriftMsgFastCodec) error {
	// nocopy write is a special implementation of linked buffer, only bytebuffer implement NocopyWrite do FastWrite
	msgBeginLen := bthrift.Binary.MessageBeginLength(methodName, thrift.TMessageType(msgType), seqID)
	msgEndLen := bthrift.Binary.MessageEndLength()
	buf, err := out.Malloc(msgBeginLen + msg.BLength() + msgEndLen)
	if err != nil {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("thrift marshal, Malloc failed: %s", err.Error()))
	}
	offset := bthrift.Binary.WriteMessageBegin(buf, methodName, thrift.TMessageType(msgType), seqID)
	offset += msg.FastWriteNocopy(buf[offset:], nil)
	bthrift.Binary.WriteMessageEnd(buf[offset:])
	return nil
}

// hyperMarshalFallback indicates whether fallback to frugal is needed.
// when the generated code is using slim template, it would return true.
func hyperMarshalFallback(data interface{}) bool {
	switch data.(type) {
	case MessageWriter:
	case MessageWriterWithContext:
	default:
		if hyperMarshalAvailable(data) {

Review comment (suggestion):
fallback should be after `encodeBasicThrift` returning an error indicating it's not supported (for codegen with slim template)

Code after:
	// prepare info
	methodName := message.RPCInfo().Invocation().MethodName()
	if methodName == "" {
		return errors.New("empty methodName in thrift Marshal")
	}
	msgType := message.MessageType()
	seqID := message.RPCInfo().Invocation().SeqID()

	data, err := getValidData(methodName, message)
	if err != nil {
		return err
	}

	// encode with hyper codec
	// NOTE: to ensure hyperMarshalEnabled is inlined so split the check logic, or it may cause performance loss
	if c.hyperMarshalEnabled() && hyperMarshalAvailable(data) {
		return c.hyperMarshal(out, methodName, msgType, seqID, data)
	}

	// encode with FastWrite
	if c.CodecType&FastWrite != 0 {
		if msg, ok := data.(ThriftMsgFastCodec); ok {
			return encodeFastThrift(out, methodName, msgType, seqID, msg)
		}
	}

	// fallback to old thrift way (slow)
	if err = encodeBasicThrift(out, ctx, methodName, msgType, seqID, data); err == nil || err != errEncodeMismatchMsgType {
		return err
	}

	// Basic can be used for disabling frugal, we need to check it
	if c.CodecType != Basic && hyperMarshalAvailable(data) {
		// fallback to frugal when the generated code is using slim template
		return c.hyperMarshal(out, methodName, msgType, seqID, data)
	}

	return errEncodeMismatchMsgType
}

// encodeFastThrift encode with the FastCodec way
func encodeFastThrift(out remote.ByteBuffer, methodName string, msgType remote.MessageType, seqID int32, msg ThriftMsgFastCodec) error {
	// nocopy write is a special implementation of linked buffer, only bytebuffer implement NocopyWrite do FastWrite
	msgBeginLen := bthrift.Binary.MessageBeginLength(methodName, thrift.TMessageType(msgType), seqID)
	msgEndLen := bthrift.Binary.MessageEndLength()
	buf, err := out.Malloc(msgBeginLen + msg.BLength() + msgEndLen)
	if err != nil {
		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("thrift marshal, Malloc failed: %s", err.Error()))
	}
	offset := bthrift.Binary.WriteMessageBegin(buf, methodName, thrift.TMessageType(msgType), seqID)
	offset += msg.FastWriteNocopy(buf[offset:], nil)