sig-pub/pkg/grpc/discovery/consul/resolver.go

package consul

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

	"github.com/hashicorp/consul/api"
	"google.golang.org/grpc/attributes"
	"google.golang.org/grpc/resolver"
)

const (
	ConsulScheme = "consul"
)

// ConsulBuilder implements resolver.Builder
type ConsulBuilder struct {
	client *api.Client
}

func NewConsulBuilder(client *api.Client) *ConsulBuilder {
	return &ConsulBuilder{
		client: client,
	}
}

// Build creates a new resolver for the given target
func (b *ConsulBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
	// Parse target: consul://consul-host:port/service-name
	targetUrl := target.URL.String()
	u, err := url.Parse(targetUrl)
	if err != nil {
		return nil, fmt.Errorf("invalid target URL: %v", err)
	}
	if u.Scheme != ConsulScheme {
		return nil, fmt.Errorf("invalid scheme, expected %s, got %s", ConsulScheme, u.Scheme)
	}

	serviceName := u.Path[1:] // Remove leading '/'
	if serviceName == "" {
		return nil, fmt.Errorf("service name is empty")
	}

	r := &consulResolver{
		client:      b.client,
		cc:          cc,
		serviceName: serviceName,
		ctx:         context.Background(),
		cancel:      nil,
	}
	r.ctx, r.cancel = context.WithCancel(r.ctx)
	r.wg.Add(1)
	go r.watch()
	return r, nil
}

// Scheme returns the resolver scheme
func (b *ConsulBuilder) Scheme() string {
	return ConsulScheme
}

// consulResolver implements resolver.Resolver
type consulResolver struct {
	client      *api.Client
	cc          resolver.ClientConn
	serviceName string
	ctx         context.Context
	cancel      context.CancelFunc
	wg          sync.WaitGroup
}

// ResolveNow triggers a resolution immediately
func (r *consulResolver) ResolveNow(_ resolver.ResolveNowOptions) {
	// Fetch service instances from Consul
	services, _, err := r.client.Health().Service(r.serviceName, "", true, nil)
	if err != nil {
		log.Printf("ResolveNow error: %v", err)
		// serviceconfig.ParseResult
		r.cc.UpdateState(resolver.State{ServiceConfig: r.cc.ParseServiceConfig(`{"loadBalancingPolicy": "round_robin"}`)})
		return
	}

	// Convert Consul services to gRPC addresses
	var addrs []resolver.Address
	for _, s := range services {
		addr := resolver.Address{
			Addr:       fmt.Sprintf("%s:%d", s.Service.Address, s.Service.Port),
			Attributes: attributes.New("service_id", s.Service.ID),
		}
		addrs = append(addrs, addr)
	}

	// Update gRPC client connection state
	state := resolver.State{
		Addresses:     addrs,
		ServiceConfig: r.cc.ParseServiceConfig(`{"loadBalancingPolicy": "round_robin"}`),
	}
	if err := r.cc.UpdateState(state); err != nil {
		log.Printf("Failed to update state: %v", err)
	}
}

// Close stops the resolver and cleans up
func (r *consulResolver) Close() {
	r.cancel()
	r.wg.Wait()
}

// watch polls Consul for service changes
func (r *consulResolver) watch() {
	defer r.wg.Done()
	var lastIndex uint64
	for {
		select {
		case <-r.ctx.Done():
			return
		default:
			services, meta, err := r.client.Health().Service(r.serviceName, "", true, &api.QueryOptions{WaitIndex: lastIndex})
			if err != nil {
				log.Printf("Watch error: %v", err)
				time.Sleep(1 * time.Second)
				continue
			}
			lastIndex = meta.LastIndex

			var addrs []resolver.Address
			for _, s := range services {
				addr := resolver.Address{
					Addr:       fmt.Sprintf("%s:%d", s.Service.Address, s.Service.Port),
					Attributes: attributes.New("service_id", s.Service.ID),
				}
				addrs = append(addrs, addr)
			}

			state := resolver.State{
				Addresses:     addrs,
				ServiceConfig: r.cc.ParseServiceConfig(`{"loadBalancingPolicy": "round_robin"}`),
			}
			if err := r.cc.UpdateState(state); err != nil {
				log.Printf("Failed to update state: %v", err)
			}
			time.Sleep(1 * time.Second) // Avoid tight loop
		}
	}
}

// func init() {
// 	resolver.Register(&consulBuilder{})
// }

// func main() {
// 	// Dial gRPC service using Consul resolver
// 	conn, err := grpc.NewClient(
// 		"consul://localhost:8500/greeter",
// 		grpc.WithInsecure(), // Use WithTransportCredentials for TLS in production
// 		grpc.WithDefaultServiceConfig(`{"loadBalancingPolicy": "round_robin"}`),
// 	)
// 	if err != nil {
// 		log.Fatalf("Failed to dial: %v", err)
// 	}
// 	defer conn.Close()

// 	client := pb.NewGreeterClient(conn)
// 	resp, err := client.SayHello(context.Background(), &pb.HelloRequest{Name: "World"})
// 	if err != nil {
// 		log.Fatalf("SayHello failed: %v", err)
// 	}
// 	log.Printf("Response: %s", resp.Message)
// }