package nats import ( "context" "fmt" "sync" "sync/atomic" "time" "github.com/gogf/gf/v2/frame/g" "github.com/nats-io/nats.go" "github.com/nats-io/nats.go/jetstream" ) var ( nc *nats.Conn js jetstream.JetStream inited bool mu sync.RWMutex natsURL string healthCtx context.Context healthCancel context.CancelFunc connected bool reconnectChan chan struct{} // 连接状态变化监听器 connStateListeners []ConnStateListener connListenersMu sync.RWMutex // 监控指标 metrics Metrics ) // Metrics 监控指标 type Metrics struct { PublishCount atomic.Int64 PublishError atomic.Int64 SubscribeCount atomic.Int64 RequestCount atomic.Int64 RequestError atomic.Int64 ConsumeCount atomic.Int64 ConsumeError atomic.Int64 } // ConnState 连接状态 type ConnState int const ( ConnStateDisconnected ConnState = iota ConnStateConnecting ConnStateConnected ConnStateReconnecting ConnStateClosed ) // ConnStateListener 连接状态监听器 type ConnStateListener func(state ConnState, err error) // GetMetrics 获取监控指标 func GetMetrics() Metrics { return metrics } // RegisterConnStateListener 注册连接状态监听器 func RegisterConnStateListener(listener ConnStateListener) { connListenersMu.Lock() defer connListenersMu.Unlock() connStateListeners = append(connStateListeners, listener) } // UnregisterConnStateListener 取消注册连接状态监听器 func UnregisterConnStateListener(listener ConnStateListener) { connListenersMu.Lock() defer connListenersMu.Unlock() for i, l := range connStateListeners { if l != nil && &l == &listener { connStateListeners = append(connStateListeners[:i], connStateListeners[i+1:]...) break } } } // notifyConnState 通知所有监听器连接状态变化 func notifyConnState(state ConnState, err error) { connListenersMu.RLock() listeners := make([]ConnStateListener, len(connStateListeners)) copy(listeners, connStateListeners) connListenersMu.RUnlock() for _, listener := range listeners { if listener != nil { listener(state, err) } } } // init 初始化 NATS 连接 func init() { // 从配置文件读取 NATS 地址 natsURL = g.Cfg().MustGet(context.Background(), "nats.url").String() if natsURL == "" { // 默认使用本地地址 natsURL = nats.DefaultURL } // 创建健康检查上下文 healthCtx, healthCancel = context.WithCancel(context.Background()) // 创建重连通知通道(增大缓冲区避免丢失通知) reconnectChan = make(chan struct{}, 10) // 启动连接 go initConnection() // 启动健康检查协程 go healthCheck() } // initConnection 初始化连接 func initConnection() { ctx := context.Background() notifyConnState(ConnStateConnecting, nil) if err := connect(ctx); err != nil { g.Log().Errorf(ctx, "NATS 初始连接失败: %v", err) notifyConnState(ConnStateDisconnected, err) } } // connect 建立 NATS 连接 func connect(ctx context.Context) error { mu.Lock() defer mu.Unlock() if nc != nil && !nc.IsClosed() { nc.Close() } // 连接选项配置 opts := []nats.Option{ nats.Name("goframe-nats-client"), nats.ReconnectWait(2 * time.Second), nats.MaxReconnects(-1), // 无限重连 nats.PingInterval(10 * time.Second), nats.MaxPingsOutstanding(5), nats.ReconnectHandler(func(nc *nats.Conn) { g.Log().Infof(ctx, "✅ NATS 重连成功: %s", nc.ConnectedUrl()) connected = true // 重新创建 JetStream 实例 if newJS, err := jetstream.New(nc); err == nil { js = newJS } // 通知重连成功 notifyConnState(ConnStateConnected, nil) // 使用非阻塞发送避免阻塞 select { case reconnectChan <- struct{}{}: default: // 通道已满,丢弃通知 } }), nats.DisconnectErrHandler(func(nc *nats.Conn, err error) { g.Log().Warningf(ctx, "⚠️ NATS 连接断开: %v, 准备重连...", err) connected = false notifyConnState(ConnStateReconnecting, err) }), nats.ClosedHandler(func(nc *nats.Conn) { g.Log().Infof(ctx, "NATS 连接已关闭: %s", nc.ConnectedUrl()) connected = false notifyConnState(ConnStateClosed, nil) }), nats.ErrorHandler(func(nc *nats.Conn, sub *nats.Subscription, err error) { g.Log().Errorf(ctx, "NATS 错误: %v", err) }), } var err error nc, err = nats.Connect(natsURL, opts...) if err != nil { return fmt.Errorf("NATS 连接失败: %w", err) } // 等待连接就绪 if nc.Status() != nats.CONNECTED { select { case <-time.After(5 * time.Second): notifyConnState(ConnStateDisconnected, fmt.Errorf("连接超时")) return fmt.Errorf("NATS 连接超时") case <-nc.StatusChanged(nats.CONNECTED): } } // 创建 JetStream 实例 js, err = jetstream.New(nc) if err != nil { return fmt.Errorf("创建 JetStream 失败: %w", err) } connected = true inited = true g.Log().Infof(ctx, "✅ NATS 连接成功: %s", nc.ConnectedUrl()) notifyConnState(ConnStateConnected, nil) return nil } // healthCheck 健康检查协程(仅作为备用检查) func healthCheck() { ticker := time.NewTicker(30 * time.Second) defer ticker.Stop() for { select { case <-healthCtx.Done(): return case <-ticker.C: mu.RLock() currentConnected := connected currentConn := nc mu.RUnlock() if !currentConnected || currentConn == nil || currentConn.IsClosed() { // 仅记录日志,不尝试重连(NATS 已有自动重连机制) g.Log().Warning(context.Background(), "NATS 连接断开,等待 NATS 自动重连...") } case <-reconnectChan: // 重连成功的通知(仅记录日志) g.Log().Info(context.Background(), "收到重连成功通知") } } } // checkConnected 检查连接状态 func checkConnected() bool { mu.RLock() defer mu.RUnlock() return connected && nc != nil && !nc.IsClosed() } // IsConnected 检查 NATS 是否已连接 func IsConnected() bool { return checkConnected() } // GetConnState 获取当前连接状态 func GetConnState() ConnState { mu.RLock() defer mu.RUnlock() if nc == nil { return ConnStateDisconnected } if nc.IsClosed() { return ConnStateClosed } if connected { return ConnStateConnected } return ConnStateDisconnected } // CreateTaskStream 创建任务消息队列流 // 存储策略: 文件存储 // 工作队列模式: 工作队列策略 func CreateTaskStream(ctx context.Context, streamName string, subjects []string) error { if !checkConnected() { return fmt.Errorf("NATS 未连接") } stream, err := js.Stream(ctx, streamName) if err == nil { // 流已存在,更新配置 _, err = js.UpdateStream(ctx, jetstream.StreamConfig{ Name: streamName, Subjects: subjects, Storage: jetstream.FileStorage, Retention: jetstream.WorkQueuePolicy, }) if err != nil { return fmt.Errorf("更新任务流失败: %w", err) } g.Log().Infof(ctx, "✅ 任务消息队列流已更新: %s", stream.CachedInfo().Config.Name) return nil } // 创建新流 stream, err = js.CreateStream(ctx, jetstream.StreamConfig{ Name: streamName, Subjects: subjects, Storage: jetstream.FileStorage, Retention: jetstream.WorkQueuePolicy, }) if err != nil { return fmt.Errorf("创建任务流失败: %w", err) } g.Log().Infof(ctx, "✅ 任务消息队列流创建成功: %s (文件存储+工作队列策略)", stream.CachedInfo().Config.Name) return nil } // CreateLogStream 创建日志流 // 存储策略: 内存存储 // 副本数: 单副本 (1) // 消息留存: 短时留存 (1小时) func CreateLogStream(ctx context.Context, streamName string, subjects []string) error { if !checkConnected() { return fmt.Errorf("NATS 未连接") } maxAge := 1 * time.Hour stream, err := js.Stream(ctx, streamName) if err == nil { // 流已存在,更新配置 _, err = js.UpdateStream(ctx, jetstream.StreamConfig{ Name: streamName, Subjects: subjects, Storage: jetstream.MemoryStorage, Replicas: 1, MaxAge: maxAge, }) if err != nil { return fmt.Errorf("更新日志流失败: %w", err) } g.Log().Infof(ctx, "✅ 日志流已更新: %s", stream.CachedInfo().Config.Name) return nil } // 创建新流 stream, err = js.CreateStream(ctx, jetstream.StreamConfig{ Name: streamName, Subjects: subjects, Storage: jetstream.MemoryStorage, Replicas: 1, MaxAge: maxAge, }) if err != nil { return fmt.Errorf("创建日志流失败: %w", err) } g.Log().Infof(ctx, "✅ 日志流创建成功: %s (内存存储+单副本+短时留存1小时)", stream.CachedInfo().Config.Name) return nil } // CreateTradeStream 创建交易业务流 // 存储策略: 文件存储 // 副本数: 3副本 // 同步刷盘: 启用 func CreateTradeStream(ctx context.Context, streamName string, subjects []string) error { if !checkConnected() { return fmt.Errorf("NATS 未连接") } stream, err := js.Stream(ctx, streamName) if err == nil { // 流已存在,更新配置 _, err = js.UpdateStream(ctx, jetstream.StreamConfig{ Name: streamName, Subjects: subjects, Storage: jetstream.FileStorage, Replicas: 3, RePublish: nil, Duplicates: 0, }) if err != nil { return fmt.Errorf("更新交易流失败: %w", err) } g.Log().Infof(ctx, "✅ 交易业务流已更新: %s", stream.CachedInfo().Config.Name) return nil } // 创建新流 stream, err = js.CreateStream(ctx, jetstream.StreamConfig{ Name: streamName, Subjects: subjects, Storage: jetstream.FileStorage, Replicas: 3, RePublish: nil, Duplicates: 0, }) if err != nil { return fmt.Errorf("创建交易流失败: %w", err) } g.Log().Infof(ctx, "✅ 交易业务流创建成功: %s (文件存储+3副本+同步刷盘)", stream.CachedInfo().Config.Name) return nil } // Publish 发布消息到指定主题 func Publish(ctx context.Context, subject string, data []byte) error { if !checkConnected() { return fmt.Errorf("NATS 未连接") } metrics.PublishCount.Add(1) _, err := js.Publish(ctx, subject, data) if err != nil { metrics.PublishError.Add(1) return fmt.Errorf("发布消息失败: %w", err) } return nil } // GetStream 获取流信息 func GetStream(ctx context.Context, streamName string) (*jetstream.StreamInfo, error) { if !checkConnected() { return nil, fmt.Errorf("NATS 未连接") } stream, err := js.Stream(ctx, streamName) if err != nil { return nil, fmt.Errorf("获取流失败: %w", err) } info, err := stream.Info(ctx) if err != nil { return nil, fmt.Errorf("获取流信息失败: %w", err) } return info, nil } // ListStreams 列出所有流(简化实现) // 注意:由于 API 限制,此方法可能需要根据实际需求进一步实现 func ListStreams(ctx context.Context) ([]string, error) { if !checkConnected() { return nil, fmt.Errorf("NATS 未连接") } // TODO: 根据实际 NATS 版本实现完整的流列表功能 return []string{}, nil } // DeleteStream 删除流 func DeleteStream(ctx context.Context, streamName string) error { if !checkConnected() { return fmt.Errorf("NATS 未连接") } if err := js.DeleteStream(ctx, streamName); err != nil { return fmt.Errorf("删除流失败: %w", err) } g.Log().Infof(ctx, "✅ 流已删除: %s", streamName) return nil } // GetConsumer 获取消费者信息 func GetConsumer(ctx context.Context, streamName, consumerName string) (*jetstream.ConsumerInfo, error) { if !checkConnected() { return nil, fmt.Errorf("NATS 未连接") } consumer, err := js.Consumer(ctx, streamName, consumerName) if err != nil { return nil, fmt.Errorf("获取消费者失败: %w", err) } info, err := consumer.Info(ctx) if err != nil { return nil, fmt.Errorf("获取消费者信息失败: %w", err) } return info, nil } // ListConsumers 列出指定流的所有消费者(简化实现) // 注意:由于 API 限制,此方法可能需要根据实际需求进一步实现 func ListConsumers(ctx context.Context, streamName string) ([]string, error) { if !checkConnected() { return nil, fmt.Errorf("NATS 未连接") } // TODO: 根据实际 NATS 版本实现完整的消费者列表功能 return []string{}, nil } // DeleteConsumer 删除消费者 func DeleteConsumer(ctx context.Context, streamName, consumerName string) error { if !checkConnected() { return fmt.Errorf("NATS 未连接") } if err := js.DeleteConsumer(ctx, streamName, consumerName); err != nil { return fmt.Errorf("删除消费者失败: %w", err) } g.Log().Infof(ctx, "✅ 消费者已删除: %s/%s", streamName, consumerName) return nil } // CreateConsumer 创建消费者 func CreateConsumer(ctx context.Context, streamName, consumerName string, config jetstream.ConsumerConfig) (jetstream.Consumer, error) { if !checkConnected() { return nil, fmt.Errorf("NATS 未连接") } // 尝试获取现有消费者 consumer, err := js.Consumer(ctx, streamName, consumerName) if err == nil { return consumer, nil } // 创建新消费者 consumer, err = js.CreateConsumer(ctx, streamName, config) if err != nil { return nil, fmt.Errorf("创建消费者失败: %w", err) } return consumer, nil } // SubscribeRequest 订阅 RPC 请求 // B服务作为服务提供者,订阅主题并响应请求时使用此方法 // subject: 订阅的主题名,与 Request 调用时使用相同的 subject func SubscribeRequest(subject string, handler func(subject string, data []byte) ([]byte, error)) (*nats.Subscription, error) { if !checkConnected() { return nil, fmt.Errorf("NATS 未连接") } sub, err := nc.Subscribe(subject, func(msg *nats.Msg) { // 处理请求 response, err := handler(msg.Subject, msg.Data) if err != nil { // 处理错误,发送错误响应 errMsg := fmt.Sprintf("处理失败: %v", err) if err = msg.Respond([]byte(errMsg)); err != nil { g.Log().Errorf(context.Background(), "RPC 错误响应失败: %v", err) } return } // 发送成功响应 if err = msg.Respond(response); err != nil { g.Log().Errorf(context.Background(), "RPC 响应失败: %v", err) } }) if err != nil { return nil, fmt.Errorf("订阅 RPC 请求失败: %w", err) } return sub, nil } // SubscribeQueueRequest 订阅队列模式的 RPC 请求(负载均衡) // 多个服务实例订阅同一主题,实现负载均衡 // subject: 订阅的主题名,与 Request 调用时使用相同的 subject // queueName: 队列组名,同一队列组的实例之间实现负载均衡 func SubscribeQueueRequest(subject, queueName string, handler func(subject string, data []byte) ([]byte, error)) (*nats.Subscription, error) { if !checkConnected() { return nil, fmt.Errorf("NATS 未连接") } sub, err := nc.QueueSubscribe(subject, queueName, func(msg *nats.Msg) { // 处理请求 response, err := handler(msg.Subject, msg.Data) if err != nil { // 处理错误,发送错误响应 errMsg := fmt.Sprintf("处理失败: %v", err) if err = msg.Respond([]byte(errMsg)); err != nil { g.Log().Errorf(context.Background(), "RPC 错误响应失败: %v", err) } return } // 发送成功响应 if err := msg.Respond(response); err != nil { g.Log().Errorf(context.Background(), "RPC 响应失败: %v", err) } }) if err != nil { return nil, fmt.Errorf("订阅队列 RPC 请求失败: %w", err) } return sub, nil } // Request RPC 请求-响应模式 // A服务调用B服务查询接口时使用此方法 func Request(ctx context.Context, subject string, data []byte, timeout time.Duration) ([]byte, error) { if !checkConnected() { return nil, fmt.Errorf("NATS 未连接") } metrics.RequestCount.Add(1) // 使用 timeout 参数创建超时上下文 timeoutCtx, cancel := context.WithTimeout(ctx, timeout) defer cancel() msg, err := nc.RequestWithContext(timeoutCtx, subject, data) if err != nil { metrics.RequestError.Add(1) return nil, fmt.Errorf("RPC 请求失败: %w", err) } if msg == nil { metrics.RequestError.Add(1) return nil, fmt.Errorf("RPC 响应为空") } return msg.Data, nil } // Close 关闭 NATS 连接 func Close() error { mu.Lock() defer mu.Unlock() // 停止健康检查协程 if healthCancel != nil { healthCancel() } // 关闭连接 if nc != nil && !nc.IsClosed() { nc.Close() connected = false inited = false g.Log().Info(context.Background(), "NATS 连接已关闭") } return nil }