github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/nphttp2/client_handler.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 nphttp2
    18  
    19  import (
    20  	"context"
    21  	"net"
    22  
    23  	"github.com/cloudwego/kitex/pkg/kerrors"
    24  	"github.com/cloudwego/kitex/pkg/remote"
    25  	"github.com/cloudwego/kitex/pkg/remote/codec/grpc"
    26  	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/metadata"
    27  	"github.com/cloudwego/kitex/pkg/rpcinfo"
    28  )
    29  
    30  type cliTransHandlerFactory struct{}
    31  
    32  // NewCliTransHandlerFactory ...
    33  func NewCliTransHandlerFactory() remote.ClientTransHandlerFactory {
    34  	return &cliTransHandlerFactory{}
    35  }
    36  
    37  func (f *cliTransHandlerFactory) NewTransHandler(opt *remote.ClientOption) (remote.ClientTransHandler, error) {
    38  	return newCliTransHandler(opt)
    39  }
    40  
    41  func newCliTransHandler(opt *remote.ClientOption) (*cliTransHandler, error) {
    42  	return &cliTransHandler{
    43  		opt:   opt,
    44  		codec: grpc.NewGRPCCodec(grpc.WithThriftCodec(opt.PayloadCodec)),
    45  	}, nil
    46  }
    47  
    48  var _ remote.ClientTransHandler = &cliTransHandler{}
    49  
    50  type cliTransHandler struct {
    51  	opt   *remote.ClientOption
    52  	codec remote.Codec
    53  }
    54  
    55  func (h *cliTransHandler) Write(ctx context.Context, conn net.Conn, msg remote.Message) (nctx context.Context, err error) {
    56  	buf := newBuffer(conn.(*clientConn))
    57  	defer buf.Release(err)
    58  
    59  	if err = h.codec.Encode(ctx, msg, buf); err != nil {
    60  		return ctx, err
    61  	}
    62  	return ctx, buf.Flush()
    63  }
    64  
    65  func (h *cliTransHandler) Read(ctx context.Context, conn net.Conn, msg remote.Message) (nctx context.Context, err error) {
    66  	buf := newBuffer(conn.(GRPCConn))
    67  	defer buf.Release(err)
    68  
    69  	// set recv grpc compressor at client to decode the pack from server
    70  	ri := msg.RPCInfo()
    71  	remote.SetRecvCompressor(ri, conn.(hasGetRecvCompress).GetRecvCompress())
    72  	err = h.codec.Decode(ctx, msg, buf)
    73  	if bizStatusErr, isBizErr := kerrors.FromBizStatusError(err); isBizErr {
    74  		if setter, ok := ri.Invocation().(rpcinfo.InvocationSetter); ok {
    75  			setter.SetBizStatusErr(bizStatusErr)
    76  			return ctx, nil
    77  		}
    78  	}
    79  	ctx = receiveHeaderAndTrailer(ctx, conn)
    80  	return ctx, err
    81  }
    82  
    83  func (h *cliTransHandler) OnRead(ctx context.Context, conn net.Conn) (context.Context, error) {
    84  	// do nothing
    85  	hdr, err := conn.(*clientConn).Header()
    86  	if err != nil {
    87  		return nil, err
    88  	}
    89  
    90  	return metadata.NewIncomingContext(ctx, hdr), nil
    91  }
    92  
    93  func (h *cliTransHandler) OnConnect(ctx context.Context) (context.Context, error) {
    94  	return ctx, nil
    95  }
    96  
    97  func (h *cliTransHandler) OnInactive(ctx context.Context, conn net.Conn) {
    98  	panic("unimplemented")
    99  }
   100  
   101  func (h *cliTransHandler) OnError(ctx context.Context, err error, conn net.Conn) {
   102  	panic("unimplemented")
   103  }
   104  
   105  func (h *cliTransHandler) OnMessage(ctx context.Context, args, result remote.Message) (context.Context, error) {
   106  	// do nothing
   107  	return ctx, nil
   108  }
   109  
   110  func (h *cliTransHandler) SetPipeline(pipeline *remote.TransPipeline) {
   111  }