xframe/component/service_discovery/http/consul/builder.go

package consul

import (
	"context"
	"fmt"
	"io"
	"net/url"
	"time"

	"git.wishpal.cn/wishpal_ironfan/xframe/base/proc"
	"git.wishpal.cn/wishpal_ironfan/xframe/component/logger"
	"github.com/hashicorp/consul/api"
	"github.com/jpillora/backoff"
	"github.com/pkg/errors"
)

const (
	defaultIndex = 0
	schemeName   = "consul"
)

type Builder struct {
	target   *target
	client   *api.Client
	addrTags map[string][]string //{addr: [tags]}
}

func NewBuilder(url url.URL) (*Builder, error) {
	// parse url
	target, err := parseURL(url)
	if err != nil {
		return nil, errors.Wrap(err, "Wrong consul URL")
	}

	// new consul client
	client, err := api.NewClient(target.consulConfig())
	if err != nil {
		return nil, errors.Wrap(err, "Couldn't connect to the Consul API")
	}

	builder := &Builder{
		target:   &target,
		client:   client,
		addrTags: make(map[string][]string),
	}
	builder.init()

	return builder, nil
}

func (b *Builder) init() {
	// watch consul service
	ctx, cancel := context.WithCancel(context.Background())
	lastIndex, err := b.getConsulService(ctx, defaultIndex)
	if err != nil {
		logger.Errorf("init get consul service err:%v", err)
	}
	go b.watchConsulService(ctx, lastIndex)

	// quit
	proc.AddShutdownListener(func() {
		cancel()
	})
	return
}

func (b *Builder) watchConsulService(ctx context.Context, lastIndex uint64) {
	//退避策略
	bck := &backoff.Backoff{
		Factor: 2,
		Jitter: true,
		Min:    10 * time.Millisecond,
		Max:    b.target.MaxBackoff,
	}

	var err error
	for {
		lastIndex, err = b.getConsulService(ctx, lastIndex)
		switch {
		case err == io.EOF:
			return
		case err != nil:
			logger.Errorf("[Consul watch] Couldn't fetch endpoints. target={%s}; error={%v}",
				b.target.String(), err)
			time.Sleep(bck.Duration())
			continue
		default:
			bck.Reset()
		}
	}
}

func (b *Builder) getConsulService(ctx context.Context, lastIndex uint64) (uint64, error) {
	var (
		errChan = make(chan error, 1)
		done    = make(chan struct{})

		_lastIndex uint64
	)

	go func() {
		defer close(done)

		entries, meta, err := b.client.Health().Service(
			b.target.Service,
			b.target.Tag,
			b.target.Healthy,
			&api.QueryOptions{
				WaitIndex:         lastIndex,
				Near:              b.target.Near,
				WaitTime:          b.target.Wait,
				Datacenter:        b.target.Dc,
				AllowStale:        b.target.AllowStale,
				RequireConsistent: b.target.RequireConsistent,
			},
		)
		if err != nil {
			errChan <- err
			return
		}

		_lastIndex = meta.LastIndex

		//set addr tags
		addrTags := make(map[string][]string)
		instances := makeInstances(entries)
		for i, entry := range entries {
			tags := make([]string, 0)
			if entry.Service != nil && entry.Service.Tags != nil {
				for _, tag := range entry.Service.Tags {
					tags = append(tags, tag)
				}
			}
			addrTags[instances[i]] = tags
		}
		b.addrTags = addrTags

		logger.Infof("[Consul watch] %d endpoints fetched in(+wait) %s for target={%s}",
			len(entries), meta.RequestTime, b.target.String(),
		)
	}()

	select {
	case <-done:
		return _lastIndex, nil
	case err := <-errChan:
		return 0, err
	case <-ctx.Done():
		return 0, io.EOF
	}
}

func (b *Builder) GetAddrTags() map[string][]string {
	return b.addrTags
}

func filterEntries(entries []*api.ServiceEntry, tags ...string) []*api.ServiceEntry {
	var es []*api.ServiceEntry

ENTRIES:
	for _, entry := range entries {
		ts := make(map[string]struct{}, len(entry.Service.Tags))
		for _, tag := range entry.Service.Tags {
			ts[tag] = struct{}{}
		}

		for _, tag := range tags {
			if _, ok := ts[tag]; !ok {
				continue ENTRIES
			}
		}
		es = append(es, entry)
	}

	return es
}

func makeInstances(entries []*api.ServiceEntry) []string {
	instances := make([]string, len(entries))
	for i, entry := range entries {
		addr := entry.Node.Address
		if entry.Service.Address != "" {
			addr = entry.Service.Address
		}
		instances[i] = fmt.Sprintf("%s:%d", addr, entry.Service.Port)
	}
	return instances
}