github.com/cloudwego/kitex@v0.9.0/pkg/remote/codec/thrift/thrift_frugal.go (about)

     1  //go:build (amd64 || arm64) && !windows && go1.16 && !go1.23 && !disablefrugal
     2  // +build amd64 arm64
     3  // +build !windows
     4  // +build go1.16
     5  // +build !go1.23
     6  // +build !disablefrugal
     7  
     8  /*
     9   * Copyright 2021 CloudWeGo Authors
    10   *
    11   * Licensed under the Apache License, Version 2.0 (the "License");
    12   * you may not use this file except in compliance with the License.
    13   * You may obtain a copy of the License at
    14   *
    15   *     http://www.apache.org/licenses/LICENSE-2.0
    16   *
    17   * Unless required by applicable law or agreed to in writing, software
    18   * distributed under the License is distributed on an "AS IS" BASIS,
    19   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    20   * See the License for the specific language governing permissions and
    21   * limitations under the License.
    22   */
    23  
    24  package thrift
    25  
    26  import (
    27  	"fmt"
    28  	"reflect"
    29  
    30  	"github.com/apache/thrift/lib/go/thrift"
    31  	"github.com/bytedance/gopkg/lang/mcache"
    32  	"github.com/cloudwego/frugal"
    33  
    34  	"github.com/cloudwego/kitex/pkg/protocol/bthrift"
    35  	"github.com/cloudwego/kitex/pkg/remote"
    36  	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
    37  )
    38  
    39  const (
    40  	// 0b0001 and 0b0010 are used for FastWrite and FastRead, so Frugal starts from 0b0100
    41  	FrugalWrite CodecType = 0b0100
    42  	FrugalRead  CodecType = 0b1000
    43  
    44  	FrugalReadWrite = FrugalWrite | FrugalRead
    45  )
    46  
    47  // hyperMarshalEnabled indicates that if there are high priority message codec for current platform.
    48  func (c thriftCodec) hyperMarshalEnabled() bool {
    49  	return c.CodecType&FrugalWrite != 0
    50  }
    51  
    52  // hyperMarshalAvailable indicates that if high priority message codec is available.
    53  func hyperMarshalAvailable(data interface{}) bool {
    54  	dt := reflect.TypeOf(data).Elem()
    55  	if dt.NumField() > 0 && dt.Field(0).Tag.Get("frugal") == "" {
    56  		return false
    57  	}
    58  	return true
    59  }
    60  
    61  // hyperMessageUnmarshalEnabled indicates that if there are high priority message codec for current platform.
    62  func (c thriftCodec) hyperMessageUnmarshalEnabled() bool {
    63  	return c.CodecType&FrugalRead != 0
    64  }
    65  
    66  // hyperMessageUnmarshalAvailable indicates that if high priority message codec is available.
    67  func hyperMessageUnmarshalAvailable(data interface{}, payloadLen int) bool {
    68  	if payloadLen == 0 {
    69  		return false
    70  	}
    71  	dt := reflect.TypeOf(data).Elem()
    72  	if dt.NumField() > 0 && dt.Field(0).Tag.Get("frugal") == "" {
    73  		return false
    74  	}
    75  	return true
    76  }
    77  
    78  func (c thriftCodec) hyperMarshal(out remote.ByteBuffer, methodName string, msgType remote.MessageType,
    79  	seqID int32, data interface{},
    80  ) error {
    81  	// calculate and malloc message buffer
    82  	msgBeginLen := bthrift.Binary.MessageBeginLength(methodName, thrift.TMessageType(msgType), seqID)
    83  	msgEndLen := bthrift.Binary.MessageEndLength()
    84  	objectLen := frugal.EncodedSize(data)
    85  	buf, err := out.Malloc(msgBeginLen + objectLen + msgEndLen)
    86  	if err != nil {
    87  		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("thrift marshal, Malloc failed: %s", err.Error()))
    88  	}
    89  
    90  	// encode message
    91  	offset := bthrift.Binary.WriteMessageBegin(buf, methodName, thrift.TMessageType(msgType), seqID)
    92  	var writeLen int
    93  	writeLen, err = frugal.EncodeObject(buf[offset:], nil, data)
    94  	if err != nil {
    95  		return perrors.NewProtocolErrorWithMsg(fmt.Sprintf("thrift marshal, Encode failed: %s", err.Error()))
    96  	}
    97  	offset += writeLen
    98  	bthrift.Binary.WriteMessageEnd(buf[offset:])
    99  	return nil
   100  }
   101  
   102  func (c thriftCodec) hyperMarshalBody(data interface{}) (buf []byte, err error) {
   103  	objectLen := frugal.EncodedSize(data)
   104  	buf = mcache.Malloc(objectLen)
   105  	_, err = frugal.EncodeObject(buf, nil, data)
   106  	return buf, err
   107  }
   108  
   109  func (c thriftCodec) hyperMessageUnmarshal(buf []byte, data interface{}) error {
   110  	_, err := frugal.DecodeObject(buf, data)
   111  	if err != nil {
   112  		return remote.NewTransError(remote.ProtocolError, err)
   113  	}
   114  	return nil
   115  }