builder.go

package consul

import (
	"fmt"
	"strconv"
	"time"

	"github.com/go-logr/logr"
	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/api/watch"
	"google.golang.org/grpc/resolver"
)

// schemeName for the urls
// All target URLs like 'consul://.../...' will be resolved by this resolver.
const schemeName = "consul"

// builder implements resolver.Builder and use for constructing all consul resolvers.
type builder struct {
	logger     logr.Logger
	strategies []Strategy
}

// NewBuilder creates a new instance of builder with logr.Logger and strategies.
func NewBuilder(logger logr.Logger, strategies []Strategy) *builder {
	return &builder{logger: logger.WithName("grpc-consul-resolver"), strategies: strategies}
}

// Build creates a new resolver for the given target.
func (b *builder) Build(url resolver.Target, cc resolver.ClientConn, _ resolver.BuildOptions) (resolver.Resolver, error) {
	tgt, err := parseURL(url.URL.String())
	if err != nil {
		return nil, fmt.Errorf("invalid consul URL: %w", err)
	}

	cli, err := api.NewClient(tgt.consulConfig())
	if err != nil {
		return nil, fmt.Errorf("failed to connect to the ConsulAPI: %w", err)
	}

	plan, err := b.targetToPlan(tgt)
	if err != nil {
		return nil, err
	}

	r := NewResolver(cc, plan, WithLogger(b.logger.WithName(tgt.Service)), WithStrategies(b.strategies))

	go func() {
		_ = plan.RunWithClientAndHclog(cli, nil)
	}()
	go func() {
		r.resolveWithPeriod(cli, tgt)
	}()

	return r, nil
}

func (builder) targetToPlan(t target) (*watch.Plan, error) {
	const paramsNumber = 6
	params := make(map[string]interface{}, paramsNumber)
	params["type"] = "service"
	params["service"] = t.Service

	if t.Dc != "" {
		params["datacenter"] = t.Dc
	}
	if t.Token != "" {
		params["token"] = t.Token
	}
	if t.AllowStale {
		params["stale"] = t.AllowStale
	}
	if t.Tag != "" {
		params["tag"] = t.Tag
	}

	plan, err := watch.Parse(params)
	if err != nil {
		return nil, fmt.Errorf("parse watch plan: %w", err)
	}
	return plan, nil
}

// Scheme returns the scheme supported by this resolver.
// Scheme is defined at https://github.com/grpc/grpc/blob/master/doc/naming.md.
func (builder) Scheme() string {
	return schemeName
}

// resolveWithPeriod resolves service addresses with some period provided in target.
func (r *Resolver) resolveWithPeriod(cli *api.Client, t target) {
	if t.Period == 0 {
		t.Period = time.Minute
	}

	ticker := time.NewTicker(t.Period)
	defer ticker.Stop()

	for now := range ticker.C {
		logger := r.logger.WithValues("time", now, "service", t.Service)
		resp, _, err := cli.Catalog().Service(t.Service, t.Tag, nil)
		if err != nil {
			logger.Error(err, "failed to get service addresses")
			continue
		}

		if len(resp) == 0 {
			continue
		}

		addresses := make([]resolver.Address, 0, len(resp))
		for _, s := range resp {
			var addr resolver.Address

			if s.ServiceAddress != "" {
				addr.Addr = s.ServiceAddress
			} else if s.Address != "" {
				addr.Addr = s.Address
			}
			if addr.Addr == "" {
				continue
			}
			if s.ServicePort > 0 {
				addr.Addr += ":" + strconv.Itoa(s.ServicePort)
			}
			addresses = append(addresses, addr)
		}

		logger.WithValues("new_addresses", len(addresses)).Info("applying changes")
		r.mu.Lock()
		r.cc.UpdateState(resolver.State{Addresses: addresses})
		r.mu.Unlock()
	}
}