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

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package generic
    18  
    19  import (
    20  	"context"
    21  	"sync/atomic"
    22  
    23  	"github.com/cloudwego/dynamicgo/conv"
    24  
    25  	"github.com/cloudwego/kitex/pkg/generic/descriptor"
    26  	"github.com/cloudwego/kitex/pkg/generic/thrift"
    27  	"github.com/cloudwego/kitex/pkg/remote"
    28  	"github.com/cloudwego/kitex/pkg/remote/codec"
    29  	"github.com/cloudwego/kitex/pkg/remote/codec/perrors"
    30  	"github.com/cloudwego/kitex/pkg/serviceinfo"
    31  )
    32  
    33  var (
    34  	_ remote.PayloadCodec = &jsonThriftCodec{}
    35  	_ Closer              = &jsonThriftCodec{}
    36  )
    37  
    38  // JSONRequest alias of string
    39  type JSONRequest = string
    40  
    41  type jsonThriftCodec struct {
    42  	svcDsc                 atomic.Value // *idl
    43  	provider               DescriptorProvider
    44  	codec                  remote.PayloadCodec
    45  	binaryWithBase64       bool
    46  	opts                   *Options
    47  	convOpts               conv.Options // used for dynamicgo conversion
    48  	convOptsWithThriftBase conv.Options // used for dynamicgo conversion with EnableThriftBase turned on
    49  	convOptsWithException  conv.Options // used for dynamicgo conversion with ConvertException turned on
    50  	dynamicgoEnabled       bool
    51  }
    52  
    53  func newJsonThriftCodec(p DescriptorProvider, codec remote.PayloadCodec, opts *Options) (*jsonThriftCodec, error) {
    54  	svc := <-p.Provide()
    55  	c := &jsonThriftCodec{codec: codec, provider: p, binaryWithBase64: true, opts: opts, dynamicgoEnabled: false}
    56  	if dp, ok := p.(GetProviderOption); ok && dp.Option().DynamicGoEnabled {
    57  		c.dynamicgoEnabled = true
    58  
    59  		convOpts := opts.dynamicgoConvOpts
    60  		c.convOpts = convOpts
    61  
    62  		convOptsWithThriftBase := convOpts
    63  		convOptsWithThriftBase.EnableThriftBase = true
    64  		c.convOptsWithThriftBase = convOptsWithThriftBase
    65  
    66  		convOptsWithException := convOpts
    67  		convOptsWithException.ConvertException = true
    68  		c.convOptsWithException = convOptsWithException
    69  	}
    70  	c.svcDsc.Store(svc)
    71  	go c.update()
    72  	return c, nil
    73  }
    74  
    75  func (c *jsonThriftCodec) update() {
    76  	for {
    77  		svc, ok := <-c.provider.Provide()
    78  		if !ok {
    79  			return
    80  		}
    81  		c.svcDsc.Store(svc)
    82  	}
    83  }
    84  
    85  func (c *jsonThriftCodec) Marshal(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
    86  	method := msg.RPCInfo().Invocation().MethodName()
    87  	if method == "" {
    88  		return perrors.NewProtocolErrorWithMsg("empty methodName in thrift Marshal")
    89  	}
    90  	if msg.MessageType() == remote.Exception {
    91  		return c.codec.Marshal(ctx, msg, out)
    92  	}
    93  	svcDsc, ok := c.svcDsc.Load().(*descriptor.ServiceDescriptor)
    94  	if !ok {
    95  		return perrors.NewProtocolErrorWithMsg("get parser ServiceDescriptor failed")
    96  	}
    97  
    98  	wm, err := thrift.NewWriteJSON(svcDsc, method, msg.RPCRole() == remote.Client)
    99  	if err != nil {
   100  		return err
   101  	}
   102  	wm.SetBase64Binary(c.binaryWithBase64)
   103  	if c.dynamicgoEnabled {
   104  		if err = wm.SetDynamicGo(svcDsc, method, &c.convOpts, &c.convOptsWithThriftBase); err != nil {
   105  			return err
   106  		}
   107  	}
   108  
   109  	msg.Data().(WithCodec).SetCodec(wm)
   110  	return c.codec.Marshal(ctx, msg, out)
   111  }
   112  
   113  func (c *jsonThriftCodec) Unmarshal(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
   114  	if err := codec.NewDataIfNeeded(serviceinfo.GenericMethod, msg); err != nil {
   115  		return err
   116  	}
   117  	svcDsc, ok := c.svcDsc.Load().(*descriptor.ServiceDescriptor)
   118  	if !ok {
   119  		return perrors.NewProtocolErrorWithMsg("get parser ServiceDescriptor failed")
   120  	}
   121  
   122  	rm := thrift.NewReadJSON(svcDsc, msg.RPCRole() == remote.Client)
   123  	rm.SetBinaryWithBase64(c.binaryWithBase64)
   124  	// Transport protocol should be TTHeader, Framed, or TTHeaderFramed to enable dynamicgo
   125  	if c.dynamicgoEnabled && msg.PayloadLen() != 0 {
   126  		rm.SetDynamicGo(&c.convOpts, &c.convOptsWithException, msg)
   127  	}
   128  
   129  	msg.Data().(WithCodec).SetCodec(rm)
   130  	return c.codec.Unmarshal(ctx, msg, in)
   131  }
   132  
   133  func (c *jsonThriftCodec) getMethod(req interface{}, method string) (*Method, error) {
   134  	fnSvc, err := c.svcDsc.Load().(*descriptor.ServiceDescriptor).LookupFunctionByMethod(method)
   135  	if err != nil {
   136  		return nil, err
   137  	}
   138  	return &Method{method, fnSvc.Oneway}, nil
   139  }
   140  
   141  func (c *jsonThriftCodec) Name() string {
   142  	return "JSONThrift"
   143  }
   144  
   145  func (c *jsonThriftCodec) Close() error {
   146  	return c.provider.Close()
   147  }