kratos/transport/grpc/resolver/discovery/builder.go

package discovery

import (
	"context"
	"time"

	"github.com/go-kratos/kratos/v2/log"
	"github.com/go-kratos/kratos/v2/registry"
	"google.golang.org/grpc/resolver"
)

const name = "discovery"

// Option is builder option.
type Option func(o *builder)

// WithLogger with builder logger.
func WithLogger(logger log.Logger) Option {
	return func(b *builder) {
		b.logger = logger
	}
}

// WithTimeout with timeout option.
func WithTimeout(timeout time.Duration) Option {
	return func(b *builder) {
		b.timeout = timeout
	}
}

type builder struct {
	discoverer registry.Discovery
	logger     log.Logger
	timeout    time.Duration
}

// NewBuilder creates a builder which is used to factory registry resolvers.
func NewBuilder(d registry.Discovery, opts ...Option) resolver.Builder {
	b := &builder{
		discoverer: d,
		logger:     log.DefaultLogger,
		timeout:    time.Second * 10,
	}
	for _, o := range opts {
		o(b)
	}
	return b
}

func (b *builder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
	ctx, cancel := context.WithTimeout(context.Background(), b.timeout)
	defer cancel()
	w, err := b.discoverer.Watch(ctx, target.Endpoint)
	if err != nil {
		return nil, err
	}

	r := &discoveryResolver{
		w:      w,
		cc:     cc,
		ctx:    ctx,
		cancel: cancel,
		log:    log.NewHelper(b.logger),
	}
	r.ctx, r.cancel = context.WithCancel(context.Background())
	go r.watch()

	return r, nil
}

// Scheme return scheme of discovery
func (*builder) Scheme() string {
	return name
}