lone-services/pkg/discovery/nacos.go

package discovery

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Config struct {
	Hosts       []string
	NamespaceId string
	Group       string
}

type Instance struct {
	ServiceName string
	IP          string
	Port        uint64
	Group       string
}

type ServiceInstance struct {
	IP          string
	Port        uint64
	ServiceName string
	Healthy     bool
	Weight      float64
	Enabled     bool
	Ephemeral   bool
	ClusterName string
	Metadata    map[string]string
}

var (
	mu         sync.Mutex
	baseURL    string
	namespace  string
	group      string
	httpClient = &http.Client{Timeout: 5 * time.Second}
	current    Instance
	inited     bool
	registered bool
	stopBeat   chan struct{}

	pickMu sync.Mutex
	pickRR = map[string]uint64{}
)

func Init(cfg Config) error {
	mu.Lock()
	defer mu.Unlock()

	if inited {
		return nil
	}
	if len(cfg.Hosts) == 0 {
		return fmt.Errorf("discovery: Hosts is empty")
	}

	host, port, err := parseHostPort(cfg.Hosts[0])
	if err != nil {
		return err
	}

	group = cfg.Group
	if group == "" {
		group = "default"
	}
	namespace = cfg.NamespaceId
	baseURL = fmt.Sprintf("http://%s:%d", host, port)
	inited = true
	return nil
}

func Register(inst Instance) error {
	mu.Lock()
	defer mu.Unlock()

	if !inited {
		return fmt.Errorf("discovery: call Init first")
	}

	g := inst.Group
	if g == "" {
		g = group
	}

	form := url.Values{}
	form.Set("ip", inst.IP)
	form.Set("port", strconv.FormatUint(inst.Port, 10))
	form.Set("serviceName", joinGroupService(g, inst.ServiceName))
	form.Set("groupName", g)
	form.Set("weight", "10")
	form.Set("enable", "true")
	form.Set("healthy", "true")
	form.Set("ephemeral", "true")
	form.Set("clusterName", "DEFAULT")
	if namespace != "" {
		form.Set("namespaceId", namespace)
	}

	if err := doForm(http.MethodPost, "/nacos/v1/ns/instance", form); err != nil {
		return fmt.Errorf("discovery: register: %w", err)
	}

	current = inst
	current.Group = g
	registered = true

	if stopBeat != nil {
		close(stopBeat)
	}
	stopBeat = make(chan struct{})
	go heartbeatLoop(stopBeat, current)

	return nil
}

func Deregister() error {
	mu.Lock()
	defer mu.Unlock()

	if !inited || !registered {
		return nil
	}

	if stopBeat != nil {
		close(stopBeat)
		stopBeat = nil
	}

	form := url.Values{}
	form.Set("ip", current.IP)
	form.Set("port", strconv.FormatUint(current.Port, 10))
	form.Set("serviceName", joinGroupService(current.Group, current.ServiceName))
	form.Set("groupName", current.Group)
	form.Set("clusterName", "DEFAULT")
	form.Set("ephemeral", "true")
	if namespace != "" {
		form.Set("namespaceId", namespace)
	}

	if err := doForm(http.MethodDelete, "/nacos/v1/ns/instance", form); err != nil {
		return fmt.Errorf("discovery: deregister: %w", err)
	}

	registered = false
	return nil
}

func GetInstances(serviceName string, groupName ...string) ([]ServiceInstance, error) {
	mu.Lock()
	defer mu.Unlock()

	if !inited {
		return nil, fmt.Errorf("discovery: call Init first")
	}

	g := group
	if len(groupName) > 0 && groupName[0] != "" {
		g = groupName[0]
	}

	q := url.Values{}
	q.Set("serviceName", joinGroupService(g, serviceName))
	q.Set("groupName", g)
	q.Set("healthyOnly", "true")
	if namespace != "" {
		q.Set("namespaceId", namespace)
	}

	body, err := doGet("/nacos/v1/ns/instance/list", q)
	if err != nil {
		return nil, fmt.Errorf("discovery: select instances: %w", err)
	}

	var resp listResponse
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, fmt.Errorf("discovery: decode instances: %w", err)
	}

	out := make([]ServiceInstance, 0, len(resp.Hosts))
	for _, h := range resp.Hosts {
		out = append(out, ServiceInstance{
			IP:          h.IP,
			Port:        uint64(h.Port),
			ServiceName: h.ServiceName,
			Healthy:     h.Healthy,
			Weight:      h.Weight,
			Enabled:     h.Enabled,
			Ephemeral:   h.Ephemeral,
			ClusterName: h.ClusterName,
			Metadata:    h.Metadata,
		})
	}
	return out, nil
}

// Pick returns one healthy instance for serviceName using round-robin.
func Pick(serviceName string, groupName ...string) (*ServiceInstance, error) {
	list, err := GetInstances(serviceName, groupName...)
	if err != nil {
		return nil, err
	}
	if len(list) == 0 {
		return nil, fmt.Errorf("discovery: no healthy instance for %q", serviceName)
	}

	pickMu.Lock()
	idx := pickRR[serviceName] % uint64(len(list))
	pickRR[serviceName]++
	pickMu.Unlock()

	inst := list[idx]
	return &inst, nil
}

type listResponse struct {
	Hosts []hostInfo `json:"hosts"`
}

type hostInfo struct {
	IP          string            `json:"ip"`
	Port        float64           `json:"port"`
	ServiceName string            `json:"serviceName"`
	Healthy     bool              `json:"healthy"`
	Weight      float64           `json:"weight"`
	Enabled     bool              `json:"enabled"`
	Ephemeral   bool              `json:"ephemeral"`
	ClusterName string            `json:"clusterName"`
	Metadata    map[string]string `json:"metadata"`
}

type beatResponse struct {
	ClientBeatInterval int64 `json:"clientBeatInterval"`
}

func heartbeatLoop(stop <-chan struct{}, inst Instance) {
	interval := 5 * time.Second
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-stop:
			return
		case <-ticker.C:
			if d := sendBeat(inst); d > 0 {
				ticker.Reset(d)
			}
		}
	}
}

func sendBeat(inst Instance) time.Duration {
	groupedName := joinGroupService(inst.Group, inst.ServiceName)
	beat, _ := json.Marshal(map[string]any{
		"ip":          inst.IP,
		"port":        inst.Port,
		"serviceName": groupedName,
		"weight":      10,
		"ephemeral":   true,
		"cluster":     "DEFAULT",
	})

	q := url.Values{}
	q.Set("serviceName", groupedName)
	q.Set("groupName", inst.Group)
	q.Set("ip", inst.IP)
	q.Set("port", strconv.FormatUint(inst.Port, 10))
	q.Set("beat", string(beat))
	if namespace != "" {
		q.Set("namespaceId", namespace)
	}
	req, err := http.NewRequest(http.MethodPut, baseURL+"/nacos/v1/ns/instance/beat?"+q.Encode(), nil)
	if err != nil {
		return 0
	}
	resp, err := httpClient.Do(req)
	if err != nil {
		return 0
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode >= 300 {
		return 0
	}

	var br beatResponse
	if json.Unmarshal(body, &br) == nil && br.ClientBeatInterval > 0 {
		return time.Duration(br.ClientBeatInterval) * time.Millisecond
	}
	return 0
}

func doForm(method, path string, form url.Values) error {
	req, err := http.NewRequest(method, baseURL+path, strings.NewReader(form.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	resp, err := httpClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode >= 300 {
		return fmt.Errorf("http %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}
	if s := strings.TrimSpace(string(body)); s != "" && s != "ok" && s != "true" {
		// rnacos / nacos 成功时通常返回 "ok"
		var ok bool
		if err := json.Unmarshal(body, &ok); err == nil && ok {
			return nil
		}
		if s != "ok" {
			return fmt.Errorf("unexpected response: %s", s)
		}
	}
	return nil
}

func doGet(path string, q url.Values) ([]byte, error) {
	resp, err := httpClient.Get(baseURL + path + "?" + q.Encode())
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= 300 {
		return nil, fmt.Errorf("http %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
	}
	return body, nil
}

func parseHostPort(addr string) (string, uint64, error) {
	host, portStr, ok := strings.Cut(addr, ":")
	if !ok || host == "" || portStr == "" {
		return "", 0, fmt.Errorf("discovery: invalid host %q, want host:port", addr)
	}
	port, err := strconv.ParseUint(portStr, 10, 64)
	if err != nil {
		return "", 0, fmt.Errorf("discovery: invalid port in %q: %w", addr, err)
	}
	return host, port, nil
}

func joinGroupService(groupName, serviceName string) string {
	if groupName == "" || strings.Contains(serviceName, "@@") {
		return serviceName
	}
	return groupName + "@@" + serviceName
}