micro/client.go

package micro

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"

	"git.nspix.com/golang/micro/gateway/rpc"
	"git.nspix.com/golang/micro/registry"
)

type Client struct {
	selector *registry.Selector
	clients  sync.Map
	once     sync.Once
}

func (c *Client) Do(ctx context.Context, r *Request) (resp Response, err error) {
	var (
		ok      bool
		node    *registry.ServiceNode
		value   interface{}
		client  *rpc.Client
		tryFlag int32
	)
	value, ok = c.clients.Load(r.ServiceName)
	if ok {
		client = value.(*rpc.Client)
		goto _EXEC
	}
_RETRY:
	if node, err = c.selector.Select(ctx, r.ServiceName); err != nil {
		err = fmt.Errorf("get service %s failed cause by %s", r.ServiceName, err.Error())
		return
	}
	client = rpc.NewClient()
	if err = client.DialerContext(ctx, "tcp", fmt.Sprintf("%s:%d", node.Address, node.Port)); err != nil {
		return
	}
	c.clients.Store(r.ServiceName, client)
_EXEC:
	if resp, err = client.Do(ctx, rpc.NewRequest(r.Method, r.Body)); err != nil {
		err = client.Close()
		c.clients.Delete(r.ServiceName)
		if atomic.CompareAndSwapInt32(&tryFlag, 0, 1) {
			goto _RETRY
		}
	}
	return
}

func (c *Client) Close() (err error) {
	c.clients.Range(func(key, value interface{}) bool {
		client := value.(*rpc.Client)
		_ = client.Close()
		c.clients.Delete(key)
		return true
	})
	return
}

func NewClient(reg registry.Registry) *Client {
	return &Client{
		selector: registry.NewSelector(reg),
	}
}