github.com/cloudwego/kitex@v0.9.0/pkg/generic/thrift/http_go116plus_amd64.go (about)

     1  //go:build amd64 && go1.16
     2  // +build amd64,go1.16
     3  
     4  /*
     5   * Copyright 2023 CloudWeGo Authors
     6   *
     7   * Licensed under the Apache License, Version 2.0 (the "License");
     8   * you may not use this file except in compliance with the License.
     9   * You may obtain a copy of the License at
    10   *
    11   *     http://www.apache.org/licenses/LICENSE-2.0
    12   *
    13   * Unless required by applicable law or agreed to in writing, software
    14   * distributed under the License is distributed on an "AS IS" BASIS,
    15   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    16   * See the License for the specific language governing permissions and
    17   * limitations under the License.
    18   */
    19  
    20  package thrift
    21  
    22  import (
    23  	"context"
    24  	"unsafe"
    25  
    26  	"github.com/apache/thrift/lib/go/thrift"
    27  	"github.com/bytedance/gopkg/lang/mcache"
    28  	"github.com/cloudwego/dynamicgo/conv"
    29  	"github.com/cloudwego/dynamicgo/conv/j2t"
    30  	"github.com/cloudwego/dynamicgo/thrift/base"
    31  
    32  	"github.com/cloudwego/kitex/pkg/generic/descriptor"
    33  	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
    34  	cthrift "github.com/cloudwego/kitex/pkg/remote/codec/thrift"
    35  )
    36  
    37  // Write ...
    38  func (w *WriteHTTPRequest) Write(ctx context.Context, out thrift.TProtocol, msg interface{}, requestBase *Base) error {
    39  	// fallback logic
    40  	if !w.dynamicgoEnabled {
    41  		return w.originalWrite(ctx, out, msg, requestBase)
    42  	}
    43  
    44  	// dynamicgo logic
    45  	req := msg.(*descriptor.HTTPRequest)
    46  
    47  	var cv j2t.BinaryConv
    48  	if !w.hasRequestBase {
    49  		requestBase = nil
    50  	}
    51  	if requestBase != nil {
    52  		base := (*base.Base)(unsafe.Pointer(requestBase))
    53  		ctx = context.WithValue(ctx, conv.CtxKeyThriftReqBase, base)
    54  		cv = j2t.NewBinaryConv(w.convOptsWithThriftBase)
    55  	} else {
    56  		cv = j2t.NewBinaryConv(w.convOpts)
    57  	}
    58  
    59  	if err := out.WriteStructBegin(w.dynamicgoTypeDsc.Struct().Name()); err != nil {
    60  		return err
    61  	}
    62  
    63  	ctx = context.WithValue(ctx, conv.CtxKeyHTTPRequest, req)
    64  	body := req.GetBody()
    65  	dbuf := mcache.Malloc(len(body))[0:0]
    66  	defer mcache.Free(dbuf)
    67  
    68  	for _, field := range w.dynamicgoTypeDsc.Struct().Fields() {
    69  		if err := out.WriteFieldBegin(field.Name(), field.Type().Type().ToThriftTType(), int16(field.ID())); err != nil {
    70  			return err
    71  		}
    72  
    73  		// json []byte to thrift []byte
    74  		if err := cv.DoInto(ctx, field.Type(), body, &dbuf); err != nil {
    75  			return err
    76  		}
    77  
    78  		// WriteFieldEnd has no content
    79  		// if err := out.WriteFieldEnd(); err != nil {
    80  		// 	return err
    81  		// }
    82  	}
    83  
    84  	tProt, ok := out.(*cthrift.BinaryProtocol)
    85  	if !ok {
    86  		return perrors.NewProtocolErrorWithMsg("TProtocol should be BinaryProtocol")
    87  	}
    88  	buf, err := tProt.ByteBuffer().Malloc(len(dbuf))
    89  	if err != nil {
    90  		return err
    91  	}
    92  	// TODO: implement MallocAck() to achieve zero copy
    93  	copy(buf, dbuf)
    94  
    95  	if err := out.WriteFieldStop(); err != nil {
    96  		return err
    97  	}
    98  	if err := out.WriteStructEnd(); err != nil {
    99  		return err
   100  	}
   101  	return nil
   102  }