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

package consul

import (
	"math"
	"sig-pub/pkg/utils/collect"
	"sig-pub/pkg/utils/retry"
	"sig-pub/pkg/zlog"
	"time"

	"github.com/hashicorp/consul/api"
	"github.com/hashicorp/consul/api/watch"
)

// 定义watcher
type Watcher struct {
	client   *api.Client
	wp       *watch.Plan                           // 总的Services变化对应的Plan
	watchers *collect.SyncMap[string, *watch.Plan] // 对已经进行监控的service作个记录
	handler  func(serviceName string, passing bool)
}

func NewWatcher(client *api.Client, handler func(serviceName string, passing bool)) *Watcher {
	return &Watcher{
		client:   client,
		watchers: collect.NewSyncMap[string, *watch.Plan](),
		handler:  handler,
	}
}

// WatchServices
// watch services doc: https://www.consul.io/docs/dynamic-app-config/watches#services
func (w *Watcher) WatchServices() (err error) {
	w.wp, err = watch.Parse(map[string]any{
		"type": "services",
	})
	if err != nil {
		return
	}
	w.wp.Handler = func(i uint64, data any) {
		switch d := data.(type) {
		// services
		case map[string][]string:
			services := make([]string, 0, len(d))
			for i := range d {
				if i == "consul" {
					continue
				}
				services = append(services, i)
				if _, loaded := w.watchers.LoadOrStore(i, nil); !loaded {
					w.registerServiceWatcher(i)
				}
			}
			zlog.Infof("consul services update: %v", services)

			// remove unknown services from watchers
			// var dels []string
			// w.watchers.Range(func(k string, wp *watch.Plan) bool {
			// 	wp.Stop()
			// 	dels = append(dels, k)
			// 	return true
			// })
			// for _, svc := range dels {
			// 	w.watchers.Delete(svc)
			// }
		default:
			zlog.Debugf("unknown consul watch type: %#v", &d)
		}
	}

	go func() {
		if err := w.wp.RunWithClientAndHclog(w.client, nil); err != nil {
			zlog.Error("watch services error: ", err)
			retry.DoWithFixDelay(math.MaxUint32, 3*time.Second, func(retryTimes uint32) (_ struct{}, err error) {
				err = w.WatchServices()
				return
			})
		}
	}()
	return
}

// 将consul新增的service加入，并监控
// doc: https://www.consul.io/docs/dynamic-app-config/watches#service
func (w *Watcher) registerServiceWatcher(serviceName string) error {
	wp, err := watch.Parse(map[string]any{
		"type":    "service",
		"service": serviceName,
	})
	if err != nil {
		return err
	}

	// 定义service变化后所执行的程序(函数)handler
	wp.Handler = func(idx uint64, data any) {
		switch d := data.(type) {
		case []*api.ServiceEntry:
			for _, i := range d {
				if w.handler != nil {
					w.handler(i.Service.Service, i.Checks.AggregatedStatus() == "passing") // status: passing / critical
				}
				// fmt.Printf("service %s 已变化", i.Service.Service)
				// // 打印service的状态
				// fmt.Println("service status: ", i.Checks.AggregatedStatus())
			}
		}
	}
	// 启动监控
	go wp.RunWithClientAndHclog(w.client, nil)

	w.watchers.Store(serviceName, wp)
	return nil
}