package discover import ( "fmt" "net" "net/http" "os" "os/signal" "regexp" "strings" "sync" "sync/atomic" "syscall" "time" "apigo.cc/go/cast" gohttp "apigo.cc/go/http" "apigo.cc/go/id" "apigo.cc/go/log" "apigo.cc/go/redis" "apigo.cc/go/safe" ) // Discoverer 发现服务实例 type Discoverer struct { config Config registry string app string serverRedisPool *redis.Redis clientRedisPool *redis.Redis pubsubRedisPool *redis.Redis isServer bool isClient bool daemonRunning atomic.Bool myAddr string logger *log.Logger daemonStopSignal chan struct{} daemonDoneSignal chan struct{} appLock sync.RWMutex calls map[string]*callInfoType appNodes map[string]map[string]*NodeInfo appSubscribed map[string]bool appClientPools map[string]*gohttp.Client appClientPoolsLock sync.RWMutex settedRoute func(*AppClient, *http.Request) settedLoadBalancer LoadBalancer } type callInfoType struct { Timeout time.Duration HttpVersion int Token *safe.SafeBuf SSL bool } // IsServer 返回当前节点是否作为服务端运行 func (d *Discoverer) IsServer() bool { return d.isServer } // IsClient 返回当前节点是否作为客户端运行 func (d *Discoverer) IsClient() bool { return d.isClient } func (d *Discoverer) logError(msg string, extra ...any) { d.logger.Error("Discover: "+msg, append(extra, "app", d.app, "addr", d.myAddr)...) } func (d *Discoverer) logInfo(msg string, extra ...any) { d.logger.Info("Discover: "+msg, append(extra, "app", d.app, "addr", d.myAddr)...) } // SetLogger 设置 Discover 使用的全局 Logger func (d *Discoverer) SetLogger(logger *log.Logger) { d.logger = logger } // New 创建一个新的发现服务实例 func New(confs ...Config) *Discoverer { var conf Config if len(confs) > 0 { conf = confs[0] } if conf.CallRetryTimes <= 0 { conf.CallRetryTimes = 10 } if conf.Weight <= 0 { conf.Weight = 100 } d := &Discoverer{ config: conf, calls: make(map[string]*callInfoType), appNodes: make(map[string]map[string]*NodeInfo), appSubscribed: make(map[string]bool), appClientPools: make(map[string]*gohttp.Client), settedLoadBalancer: &DefaultLoadBalancer{}, daemonStopSignal: make(chan struct{}), daemonDoneSignal: make(chan struct{}), } if conf.Logger != nil { d.logger = conf.Logger } else { d.logger = log.New(id.MakeID(12)) } return d } // Start 启动服务发现,指定当前节点的外部访问地址 func Start(registry, app, addr string, confs ...Config) *Discoverer { d := New(confs...) d.registry = registry d.app = app d.myAddr = addr d.isServer = d.app != "" && d.config.Weight > 0 if d.isServer && d.registry != "" { d.serverRedisPool = redis.GetRedis(d.registry, d.logger) if d.serverRedisPool.Error != nil { d.logError(d.serverRedisPool.Error.Error()) } // 注册节点 if d.serverRedisPool.Do("HSET", d.app, addr, d.config.Weight).Error == nil { d.serverRedisPool.Do("SETEX", d.app+"_"+addr, 10, "1") d.logInfo("registered") d.serverRedisPool.PUBLISH("CH_"+d.app, fmt.Sprintf("%s %d", addr, d.config.Weight)) d.daemonRunning.Store(true) d.daemonStopSignal = make(chan struct{}) d.daemonDoneSignal = make(chan struct{}) go d.daemon() } else { d.logError("register failed") } } calls := d.config.Calls if len(calls) > 0 { for callApp, c := range calls { d.addApp(callApp, c, false) } if !d.startSub() { return d } } return d } // Open 启动服务发现仅作为客户端 func Open(registry string, confs ...Config) *Discoverer { d := New(confs...) d.registry = registry calls := d.config.Calls if len(calls) > 0 { for callApp, c := range calls { d.addApp(callApp, c, false) } } d.startSub() return d } func (d *Discoverer) daemon() { d.logInfo("daemon thread started") ticker := time.NewTicker(5 * time.Second) defer ticker.Stop() for d.daemonRunning.Load() { select { case <-ticker.C: if !d.daemonRunning.Load() { break } if d.isServer && d.serverRedisPool != nil { if !d.serverRedisPool.Do("HEXISTS", d.app, d.myAddr).Bool() { d.logInfo("lost app registered info, re-registering") if d.serverRedisPool.Do("HSET", d.app, d.myAddr, d.config.Weight).Error == nil { d.serverRedisPool.Do("SETEX", d.app+"_"+d.myAddr, 10, "1") d.serverRedisPool.PUBLISH("CH_"+d.app, fmt.Sprintf("%s %d", d.myAddr, d.config.Weight)) } } else { d.serverRedisPool.Do("SETEX", d.app+"_"+d.myAddr, 10, "1") } } case <-d.daemonStopSignal: goto done } } done: d.logInfo("daemon thread stopped") close(d.daemonDoneSignal) } func (d *Discoverer) startSub() bool { if d.registry == "" { return true } d.appLock.Lock() if d.clientRedisPool == nil { d.clientRedisPool = redis.GetRedis(d.registry, d.logger) } if d.pubsubRedisPool == nil { d.pubsubRedisPool = redis.GetRedis(d.registry, d.logger.New(id.MakeID(12))) // 订阅所有已注册的应用 for app := range d.appSubscribed { d.subscribeApp(app) } // 必须在释放锁之前完成配置,但在释放锁之后启动,避免死锁 d.appLock.Unlock() d.pubsubRedisPool.Start() d.appLock.Lock() } d.isClient = true d.appLock.Unlock() return true } func (d *Discoverer) subscribeApp(app string) { if d.pubsubRedisPool == nil { d.appSubscribed[app] = true return } d.pubsubRedisPool.Subscribe("CH_"+app, func() { d.fetchApp(app) }, func(data []byte) { a := strings.Split(string(data), " ") addr := a[0] weight := 0 if len(a) == 2 { weight = cast.Int(a[1]) } d.logInfo("received node update", "app", app, "addr", addr, "weight", weight) d.pushNode(app, addr, weight) }) } func (d *Discoverer) subscribeAppWithLock(app string) { d.appLock.Lock() defer d.appLock.Unlock() d.subscribeApp(app) } // Stop 停止 Discover 并从注册中心注销当前节点 func (d *Discoverer) Stop() { d.appLock.Lock() // 1. 提取需要的状态,提前修改标志位 isClient := d.isClient pubsub := d.pubsubRedisPool d.isClient = false isServer := d.isServer serverPool := d.serverRedisPool myAddr := d.myAddr if isServer { d.daemonRunning.Store(false) if d.daemonStopSignal != nil { close(d.daemonStopSignal) } d.isServer = false } // 核心修复:在这里尽早释放锁!避免与 pushNode 回调发生死锁 d.appLock.Unlock() // 2. 在无锁状态下进行耗时的网络和停止操作 if isClient && pubsub != nil { pubsub.Stop() } if isServer && serverPool != nil { serverPool.Do("HDEL", d.app, myAddr) serverPool.Do("DEL", d.app+"_"+myAddr) serverPool.PUBLISH("CH_"+d.app, fmt.Sprintf("%s %d", myAddr, 0)) } // 3. 释放 HTTP 连接池 d.appClientPoolsLock.Lock() for _, client := range d.appClientPools { client.Destroy() } d.appClientPools = make(map[string]*gohttp.Client) d.appClientPoolsLock.Unlock() // 4. 重置状态以支持重新启动 d.appLock.Lock() d.serverRedisPool = nil d.clientRedisPool = nil d.pubsubRedisPool = nil d.appNodes = make(map[string]map[string]*NodeInfo) d.appSubscribed = make(map[string]bool) d.appLock.Unlock() } // Wait 等待守护进程退出 func (d *Discoverer) Wait() { if d.daemonDoneSignal != nil { <-d.daemonDoneSignal } } // EasyStart 自动根据本地网卡信息启动 Discover func EasyStart(registry, app string, confs ...Config) (*Discoverer, string, int) { port := 0 if listen := os.Getenv("DISCOVER_LISTEN"); listen != "" { if _, p, err := net.SplitHostPort(listen); err == nil { port = cast.Int(p) } else { port = cast.Int(listen) } } ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port)) if err != nil { return nil, "", 0 } addrInfo := ln.Addr().(*net.TCPAddr) _ = ln.Close() port = addrInfo.Port var conf Config if len(confs) > 0 { conf = confs[0] } ip := addrInfo.IP if !ip.IsGlobalUnicast() { addrs, _ := net.InterfaceAddrs() for _, a := range addrs { if an, ok := a.(*net.IPNet); ok { ip4 := an.IP.To4() if ip4 == nil || !ip4.IsGlobalUnicast() { continue } if conf.IpPrefix != "" && strings.HasPrefix(ip4.String(), conf.IpPrefix) { ip = ip4 break } if !strings.HasPrefix(ip4.String(), "172.17.") { ip = ip4 } } } } addr := fmt.Sprintf("%s:%d", ip.String(), port) d := Start(registry, app, addr, conf) sigChan := make(chan os.Signal, 1) signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) go func() { <-sigChan d.Stop() }() return d, ip.String(), port } // AddExternalApp 动态添加需要发现的外部应用 func (d *Discoverer) AddExternalApp(app, callConf string) bool { if d.addApp(app, callConf, true) { if !d.isClient { d.startSub() } else { d.subscribeAppWithLock(app) } d.fetchApp(app) // 同步拉取一次 return true } return false } // SetNode 手动设置某个服务的节点信息 func (d *Discoverer) SetNode(app, addr string, weight int) { d.pushNode(app, addr, weight) } func (d *Discoverer) getCallInfo(app string) *callInfoType { d.appLock.RLock() defer d.appLock.RUnlock() return d.calls[app] } var numberMatcher = regexp.MustCompile(`^\d+(s|ms|us|µs|ns?)?$`) func (d *Discoverer) addApp(app, callConf string, fetch bool) bool { d.appLock.Lock() // 1. 写时复制(Copy-on-Write):创建一个全新的 Map 避免影响读操作 newCalls := make(map[string]string) for k, v := range d.config.Calls { newCalls[k] = v } if newCalls[app] == callConf && d.appNodes[app] != nil { d.appLock.Unlock() return false } newCalls[app] = callConf d.config.Calls = newCalls // 将新的 Map 赋值给 Config callInfo := &callInfoType{ Timeout: 10 * time.Second, HttpVersion: 2, SSL: false, } for _, v := range cast.Split(callConf, ":") { switch v { case "1": callInfo.HttpVersion = 1 case "2": callInfo.HttpVersion = 2 case "s", "https": callInfo.SSL = true callInfo.HttpVersion = 2 case "http": callInfo.SSL = false callInfo.HttpVersion = 1 case "h2c": callInfo.SSL = false callInfo.HttpVersion = 2 default: if numberMatcher.MatchString(v) { callInfo.Timeout = cast.Duration(v) } else { callInfo.Token = safe.NewSafeBuf([]byte(v)) } } } d.calls[app] = callInfo if d.appNodes[app] == nil { d.appNodes[app] = make(map[string]*NodeInfo) } d.appSubscribed[app] = true d.appLock.Unlock() if fetch && d.isClient { d.fetchApp(app) } return true } func (d *Discoverer) fetchApp(app string) { d.appLock.RLock() pool := d.clientRedisPool d.appLock.RUnlock() if pool == nil { return } results := pool.Do("HGETALL", app).ResultMap() // 检查存活 for addr := range results { if !pool.Do("EXISTS", app+"_"+addr).Bool() { pool.Do("HDEL", app, addr) delete(results, addr) } } currentNodes := d.getAppNodes(app) if currentNodes != nil { for addr := range currentNodes { if _, ok := results[addr]; !ok { d.pushNode(app, addr, 0) } } } for addr, res := range results { d.pushNode(app, addr, res.Int()) } } func (d *Discoverer) getAppNodes(app string) map[string]*NodeInfo { d.appLock.RLock() defer d.appLock.RUnlock() if d.appNodes[app] == nil { return nil } nodes := make(map[string]*NodeInfo) for k, v := range d.appNodes[app] { nodes[k] = v } return nodes } func (d *Discoverer) getCalls() map[string]string { calls := make(map[string]string) for k, v := range d.config.Calls { calls[k] = v } return calls } // GetAppNodes 获取某个应用的所有节点列表 func (d *Discoverer) GetAppNodes(app string) map[string]*NodeInfo { return d.getAppNodes(app) } func (d *Discoverer) pushNode(app, addr string, weight int) { d.appLock.Lock() defer d.appLock.Unlock() if weight <= 0 { if d.appNodes[app] != nil { delete(d.appNodes[app], addr) } return } if d.appNodes[app] == nil { d.appNodes[app] = make(map[string]*NodeInfo) } if node, ok := d.appNodes[app][addr]; ok { if node.Weight != weight { used := node.UsedTimes.Load() node.UsedTimes.Store(uint64(float64(used) / float64(node.Weight) * float64(weight))) node.Weight = weight } } else { var avgUsed uint64 = 0 if len(d.appNodes[app]) > 0 { var totalScore float64 for _, n := range d.appNodes[app] { totalScore += float64(n.UsedTimes.Load()) / float64(n.Weight) } avgUsed = uint64(totalScore / float64(len(d.appNodes[app])) * float64(weight)) } node := &NodeInfo{ Addr: addr, Weight: weight, } node.UsedTimes.Store(avgUsed) d.appNodes[app][addr] = node } }