github.com/cloudwego/kitex@v0.9.0/pkg/remote/remotecli/client.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 remotecli
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"net"
    23  	"sync"
    24  	"time"
    25  
    26  	"github.com/cloudwego/kitex/pkg/kerrors"
    27  	"github.com/cloudwego/kitex/pkg/remote"
    28  	"github.com/cloudwego/kitex/pkg/rpcinfo"
    29  )
    30  
    31  // Client implementation may be different between mux and non-mux clients.
    32  type Client interface {
    33  	// RPCInfo as param just avoid to get it from ctx
    34  	Send(ctx context.Context, ri rpcinfo.RPCInfo, req remote.Message) (err error)
    35  	Recv(ctx context.Context, ri rpcinfo.RPCInfo, resp remote.Message) (err error)
    36  	Recycle()
    37  }
    38  
    39  var clientPool = &sync.Pool{
    40  	New: func() interface{} {
    41  		return new(client)
    42  	},
    43  }
    44  
    45  type client struct {
    46  	transHdlr   remote.TransHandler
    47  	connManager *ConnWrapper
    48  	conn        net.Conn
    49  }
    50  
    51  // NewClient creates a new Client using the given params.
    52  func NewClient(ctx context.Context, ri rpcinfo.RPCInfo, handler remote.TransHandler, opt *remote.ClientOption) (Client, error) {
    53  	cm := NewConnWrapper(opt.ConnPool)
    54  	var err error
    55  	// used by streaming unary
    56  	for _, shdlr := range opt.StreamingMetaHandlers {
    57  		ctx, err = shdlr.OnConnectStream(ctx)
    58  		if err != nil {
    59  			return nil, err
    60  		}
    61  	}
    62  	rawConn, err := cm.GetConn(ctx, opt.Dialer, ri)
    63  	if err != nil {
    64  		if errors.Is(err, kerrors.ErrGetConnection) {
    65  			return nil, err
    66  		}
    67  		return nil, kerrors.ErrGetConnection.WithCause(err)
    68  	}
    69  	cli := clientPool.Get().(*client)
    70  	cli.init(handler, cm, rawConn)
    71  	return cli, nil
    72  }
    73  
    74  func (c *client) Recycle() {
    75  	c.transHdlr = nil
    76  	c.connManager = nil
    77  	c.conn = nil
    78  	clientPool.Put(c)
    79  }
    80  
    81  func (c *client) init(handler remote.TransHandler, cm *ConnWrapper, conn net.Conn) {
    82  	c.transHdlr = handler
    83  	c.connManager = cm
    84  	c.conn = conn
    85  }
    86  
    87  // Send is blocked.
    88  func (c *client) Send(ctx context.Context, ri rpcinfo.RPCInfo, req remote.Message) (err error) {
    89  	_, err = c.transHdlr.Write(ctx, c.conn, req)
    90  	if err != nil {
    91  		c.connManager.ReleaseConn(err, ri)
    92  	}
    93  	return err
    94  }
    95  
    96  // Recv is blocked.
    97  func (c *client) Recv(ctx context.Context, ri rpcinfo.RPCInfo, resp remote.Message) (err error) {
    98  	// resp is nil means oneway
    99  	if resp != nil {
   100  		ctx, err = c.transHdlr.Read(ctx, c.conn, resp)
   101  		c.transHdlr.OnMessage(ctx, nil, resp)
   102  	} else {
   103  		// Wait for the request to be flushed out before closing the connection.
   104  		// 500us is an acceptable duration to keep a minimal loss rate at best effort.
   105  		time.Sleep(time.Millisecond / 2)
   106  	}
   107  
   108  	c.connManager.ReleaseConn(err, ri)
   109  	return err
   110  }