openbilibili/library/net/rpc/liverpc/client.go

package liverpc

import (
	"context"
	"fmt"
	"net/url"
	"sync/atomic"
	"time"

	"go-common/library/conf/env"
	"go-common/library/log"
	"go-common/library/naming"
	"go-common/library/naming/discovery"
	"go-common/library/net/metadata"
	"go-common/library/net/trace"
	"go-common/library/stat"
	xtime "go-common/library/time"

	"github.com/golang/protobuf/proto"
	"github.com/pkg/errors"
)

// Key is ContextKey
type Key int

const (
	_ Key = iota
	// KeyHeader use this in context to pass rpc header field
	// Depreated 请使用HeaderOption来传递Header
	KeyHeader
	// KeyTimeout deprecated
	// Depreated 请使用HTTPOption来传递HTTP
	KeyTimeout
)

const (
	_scheme      = "liverpc"
	_dialRetries = 3
)

// Get Implement tracer carrier interface
func (m *Header) Get(key string) string {
	if key == trace.KeyTraceID {
		return m.TraceId
	}
	return ""
}

// Set Implement tracer carrier interface
func (m *Header) Set(key string, val string) {
	if key == trace.KeyTraceID {
		m.TraceId = val
	}
}

var (
	// ErrNoClient no RPC client.
	errNoClient     = errors.New("no rpc client")
	errGroupInvalid = errors.New("invalid group")

	stats = stat.RPCClient
)

// GroupAddrs a map struct storing addrs vary groups
type GroupAddrs map[string][]string

// ClientConfig client config.
type ClientConfig struct {
	AppID       string
	Group       string
	Timeout     xtime.Duration
	ConnTimeout xtime.Duration
	Addr        string // if addr is provided, it will use add, else, use discovery
}

// Client is a RPC client.
type Client struct {
	conf     *ClientConfig
	dis      naming.Resolver
	addrs    atomic.Value // GroupAddrs
	addrsIdx int64
}

// NewClient new a RPC client with discovery.
func NewClient(c *ClientConfig) *Client {
	if c.Timeout <= 0 {
		c.Timeout = xtime.Duration(time.Second)
	}
	if c.ConnTimeout <= 0 {
		c.ConnTimeout = xtime.Duration(time.Second)
	}
	cli := &Client{
		conf: c,
	}
	if c.Addr != "" {
		groupAddrs := make(GroupAddrs)
		groupAddrs[""] = []string{c.Addr}
		cli.addrs.Store(groupAddrs)
		return cli
	}

	cli.dis = discovery.Build(c.AppID)
	// discovery watch & fetch nodes
	event := cli.dis.Watch()
	select {
	case _, ok := <-event:
		if !ok {
			panic("刚启动就从discovery拉到了关闭的event")
		}
		cli.disFetch()
		fmt.Printf("开始创建：%s 的liverpc client，等待从discovery拉取节点：%s\n", c.AppID, time.Now().Format("2006-01-02 15:04:05"))
	case <-time.After(10 * time.Second):
		fmt.Printf("失败创建：%s 的liverpc client，竟然从discovery拉取节点超时了：%s\n", c.AppID, time.Now().Format("2006-01-02 15:04:05"))
	}
	go cli.disproc(event)
	return cli
}

func (c *Client) disproc(event <-chan struct{}) {
	for {
		_, ok := <-event
		if !ok {
			return
		}
		c.disFetch()
	}
}

func (c *Client) disFetch() {
	ins, ok := c.dis.Fetch(context.Background())
	if !ok {
		return
	}
	insZone, ok := ins[env.Zone]
	if !ok {
		return
	}
	addrs := make(GroupAddrs)
	for _, svr := range insZone {
		group, ok := svr.Metadata["color"]
		if !ok {
			group = ""
		}
		for _, addr := range svr.Addrs {
			u, err := url.Parse(addr)
			if err == nil && u.Scheme == _scheme {
				addrs[group] = append(addrs[group], u.Host)
			}
		}
	}
	if len(addrs) > 0 {
		c.addrs.Store(addrs)
	}
}

// pickConn pick conn by addrs
func (c *Client) pickConn(ctx context.Context, addrs []string, dialTimeout time.Duration) (*ClientConn, error) {
	var (
		lastErr error
	)
	if len(addrs) == 0 {
		lastErr = errors.New("addrs empty")
	} else {
		for i := 0; i < _dialRetries; i++ {
			idx := atomic.AddInt64(&c.addrsIdx, 1)
			addr := addrs[int(idx)%len(addrs)]
			if dialTimeout == 0 {
				dialTimeout = time.Duration(c.conf.ConnTimeout)
			}
			cc, err := Dial(ctx, "tcp", addr, time.Duration(c.conf.Timeout), dialTimeout)
			if err != nil {
				lastErr = errors.Wrapf(err, "Dial %s error", addr)
				continue
			}
			return cc, nil
		}
	}
	if lastErr != nil {
		return nil, errors.WithMessage(errNoClient, lastErr.Error())
	}
	return nil, errors.WithStack(errNoClient)
}

// fetchAddrs fetch addrs by different strategies
// source_group first, come from request header if exists, currently only CallRaw supports source_group
// then env group, come from os.env
// since no invalid group found, return error
func (c *Client) fetchAddrs(ctx context.Context, request interface{}) (addrs []string, err error) {
	var (
		args        *Args
		groupAddrs  GroupAddrs
		ok          bool
		sourceGroup string
		groups      []string
	)
	defer func() {
		if err != nil {
			err = errors.WithMessage(errGroupInvalid, err.Error())
		}
	}()
	// try parse request header and fetch source group
	if args, ok = request.(*Args); ok && args.Header != nil {
		sourceGroup = args.Header.SourceGroup
		if sourceGroup != "" {
			groups = append(groups, sourceGroup)
		}
	}
	metaColor := metadata.String(ctx, metadata.Color)
	if metaColor != "" && metaColor != sourceGroup {
		groups = append(groups, metaColor)
	}

	if env.Color != "" && env.Color != metaColor {
		groups = append(groups, env.Color)
	}

	groups = append(groups, "")

	if groupAddrs, ok = c.addrs.Load().(GroupAddrs); !ok {
		err = errors.New("addrs load error")
		return
	}
	if len(groupAddrs) == 0 {
		err = errors.New("group addrs empty")
		return
	}
	for _, group := range groups {
		if addrs, ok = groupAddrs[group]; ok {
			break
		}
	}
	if len(addrs) == 0 {
		err = errors.Errorf("addrs empty source(%s), metadata(%s), env(%s), default empty, allAddrs(%+v)",
			sourceGroup, metaColor, env.Color, groupAddrs)
		return
	}
	return
}

// Call call the service method, waits for it to complete, and returns its error status.
// client: {service}
// serviceMethod: {version}|{controller.method}
// httpURL: /room/v1/Room/room_init
// httpURL: /{service}/{version}/{controller}/{method}
func (c *Client) Call(ctx context.Context, version int, serviceMethod string, in proto.Message, out proto.Message, opts ...CallOption) (err error) {
	var (
		cc    *ClientConn
		addrs []string
	)
	isPickErr := true

	defer func() {
		if cc != nil {
			cc.Close()
		}
		if err != nil && isPickErr {
			log.Error("liverpc Call pick connection error, version %d, method: %s, error: %+v", version, serviceMethod, err)
		}
	}() // for now it is non-persistent connection
	var cInfo = &callInfo{}
	for _, o := range opts {
		o.before(cInfo)
	}
	addrs, err = c.fetchAddrs(ctx, in)
	if err != nil {
		return
	}
	cc, err = c.pickConn(ctx, addrs, cInfo.DialTimeout)
	if err != nil {
		return
	}
	isPickErr = false
	cc.callInfo = cInfo
	err = cc.Call(ctx, version, serviceMethod, in, out)
	if err != nil {
		return
	}
	for _, o := range opts {
		o.after(cc.callInfo)
	}
	return
}

// CallRaw call the service method, waits for it to complete, and returns reply its error status.
// this is can be use without protobuf
// client: {service}
// serviceMethod: {version}|{controller.method}
// httpURL: /room/v1/Room/room_init
// httpURL: /{service}/{version}/{controller}/{method}
func (c *Client) CallRaw(ctx context.Context, version int, serviceMethod string, in *Args, opts ...CallOption) (out *Reply, err error) {
	var (
		cc    *ClientConn
		addrs []string
	)
	isPickErr := true

	defer func() {
		if cc != nil {
			cc.Close()
		}
		if err != nil && isPickErr {
			log.Error("liverpc CallRaw pick connection error, version %d, method: %s, error: %+v", version, serviceMethod, err)
		}
	}() // for now it is non-persistent connection
	var cInfo = &callInfo{}
	for _, o := range opts {
		o.before(cInfo)
	}
	addrs, err = c.fetchAddrs(ctx, in)
	if err != nil {
		return
	}
	cc, err = c.pickConn(ctx, addrs, cInfo.DialTimeout)
	if err != nil {
		return
	}
	isPickErr = false
	cc.callInfo = cInfo
	out, err = cc.CallRaw(ctx, version, serviceMethod, in)
	if err != nil {
		return
	}
	for _, o := range opts {
		o.after(cc.callInfo)
	}
	return
}

//Close handle client exit
func (c *Client) Close() {
	if c.dis != nil {
		c.dis.Close()
	}
}