github.com/cloudwego/kitex@v0.9.0/pkg/generic/thrift/json_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  	dthrift "github.com/cloudwego/dynamicgo/thrift"
    31  	"github.com/cloudwego/dynamicgo/thrift/base"
    32  
    33  	"github.com/cloudwego/kitex/pkg/generic/descriptor"
    34  	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
    35  	cthrift "github.com/cloudwego/kitex/pkg/remote/codec/thrift"
    36  	"github.com/cloudwego/kitex/pkg/utils"
    37  )
    38  
    39  // Write write json string to out thrift.TProtocol
    40  func (m *WriteJSON) Write(ctx context.Context, out thrift.TProtocol, msg interface{}, requestBase *Base) error {
    41  	// fallback logic
    42  	if !m.dynamicgoEnabled {
    43  		return m.originalWrite(ctx, out, msg, requestBase)
    44  	}
    45  
    46  	// dynamicgo logic
    47  	var cv j2t.BinaryConv
    48  	if !m.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(m.convOptsWithThriftBase)
    55  	} else {
    56  		cv = j2t.NewBinaryConv(m.convOpts)
    57  	}
    58  
    59  	// msg is void or nil
    60  	if _, ok := msg.(descriptor.Void); ok || msg == nil {
    61  		if err := m.writeHead(out); err != nil {
    62  			return err
    63  		}
    64  		if err := m.writeFields(ctx, out, nil, nil); err != nil {
    65  			return err
    66  		}
    67  		return writeTail(out)
    68  	}
    69  
    70  	// msg is string
    71  	s, ok := msg.(string)
    72  	if !ok {
    73  		return perrors.NewProtocolErrorWithType(perrors.InvalidData, "decode msg failed, is not string")
    74  	}
    75  	transBuff := utils.StringToSliceByte(s)
    76  
    77  	if err := m.writeHead(out); err != nil {
    78  		return err
    79  	}
    80  	if err := m.writeFields(ctx, out, &cv, transBuff); err != nil {
    81  		return err
    82  	}
    83  	return writeTail(out)
    84  }
    85  
    86  type MsgType int
    87  
    88  const (
    89  	Void MsgType = iota
    90  	String
    91  )
    92  
    93  func (m *WriteJSON) writeFields(ctx context.Context, out thrift.TProtocol, cv *j2t.BinaryConv, transBuff []byte) error {
    94  	dbuf := mcache.Malloc(len(transBuff))[0:0]
    95  	defer mcache.Free(dbuf)
    96  
    97  	for _, field := range m.dynamicgoTypeDsc.Struct().Fields() {
    98  		// Exception field
    99  		if !m.isClient && field.ID() != 0 {
   100  			// generic server ignore the exception, because no description for exception
   101  			// generic handler just return error
   102  			continue
   103  		}
   104  
   105  		if err := out.WriteFieldBegin(field.Name(), field.Type().Type().ToThriftTType(), int16(field.ID())); err != nil {
   106  			return err
   107  		}
   108  		// if the field type is void, just write void and return
   109  		if field.Type().Type() == dthrift.VOID {
   110  			if err := writeFieldForVoid(field.Name(), out); err != nil {
   111  				return err
   112  			}
   113  			return nil
   114  		} else {
   115  			// encode using dynamicgo
   116  			// json []byte to thrift []byte
   117  			if err := cv.DoInto(ctx, field.Type(), transBuff, &dbuf); err != nil {
   118  				return err
   119  			}
   120  		}
   121  		// WriteFieldEnd has no content
   122  		// if err := out.WriteFieldEnd(); err != nil {
   123  		// 	return err
   124  		// }
   125  	}
   126  	tProt, ok := out.(*cthrift.BinaryProtocol)
   127  	if !ok {
   128  		return perrors.NewProtocolErrorWithMsg("TProtocol should be BinaryProtocol")
   129  	}
   130  	buf, err := tProt.ByteBuffer().Malloc(len(dbuf))
   131  	if err != nil {
   132  		return err
   133  	}
   134  	// TODO: implement MallocAck() to achieve zero copy
   135  	copy(buf, dbuf)
   136  	return nil
   137  }
   138  
   139  func (m *WriteJSON) writeHead(out thrift.TProtocol) error {
   140  	if err := out.WriteStructBegin(m.dynamicgoTypeDsc.Struct().Name()); err != nil {
   141  		return err
   142  	}
   143  	return nil
   144  }
   145  
   146  func writeTail(out thrift.TProtocol) error {
   147  	if err := out.WriteFieldStop(); err != nil {
   148  		return err
   149  	}
   150  	return out.WriteStructEnd()
   151  }
   152  
   153  func writeFieldForVoid(name string, out thrift.TProtocol) error {
   154  	if err := out.WriteStructBegin(name); err != nil {
   155  		return err
   156  	}
   157  	if err := out.WriteFieldStop(); err != nil {
   158  		return err
   159  	}
   160  	if err := out.WriteStructEnd(); err != nil {
   161  		return err
   162  	}
   163  	return nil
   164  }