Files
common/nats/nats.go

955 lines
26 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package nats
import (
"context"
"encoding/json"
"fmt"
"reflect"
"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
}
// ============ RPC 服务封装 ============
// 以下方法提供了完全抽象的 RPC 调用接口
// 调用方和响应方完全不需要知道底层使用的是 NATS 的发布订阅模式
// RPC 服务注册表
var (
rpcServices map[string]RPCHandler
rpcSubs map[string]*nats.Subscription // 服务名 -> 订阅
rpcServicesMu sync.RWMutex
queueRPCServices map[string]map[string]RPCHandler // queueName -> subject -> handler
queueRPCSubs map[string]map[string]*nats.Subscription // queueName -> serviceName -> 订阅
queueRPCMu sync.RWMutex
)
// RPCHandler RPC 处理函数类型
// 实现方只需要关注请求参数和返回值,无需了解底层 NATS 实现
type RPCHandler func(ctx context.Context, req []byte) ([]byte, error)
// RegisterRPCService 注册 RPC 服务(单实例)
// serviceName: 服务名称,调用方通过此名称调用服务
// handler: 服务处理函数,接收请求并返回响应
func RegisterRPCService(serviceName string, handler RPCHandler) error {
if !checkConnected() {
return fmt.Errorf("NATS 未连接")
}
rpcServicesMu.Lock()
if rpcServices == nil {
rpcServices = make(map[string]RPCHandler)
}
if rpcSubs == nil {
rpcSubs = make(map[string]*nats.Subscription)
}
// 如果已存在该服务,先取消之前的订阅
if oldSub, exists := rpcSubs[serviceName]; exists {
oldSub.Unsubscribe()
}
rpcServices[serviceName] = handler
rpcServicesMu.Unlock()
// 订阅服务主题
subject := fmt.Sprintf("rpc.%s", serviceName)
sub, err := nc.Subscribe(subject, func(msg *nats.Msg) {
ctx := context.Background()
response, err := handler(ctx, msg.Data)
if err != nil {
errMsg := fmt.Sprintf("处理失败: %v", err)
if err = msg.Respond([]byte(errMsg)); err != nil {
g.Log().Errorf(ctx, "RPC 错误响应失败: %v", err)
}
return
}
if err = msg.Respond(response); err != nil {
g.Log().Errorf(ctx, "RPC 响应失败: %v", err)
}
})
if err != nil {
return fmt.Errorf("注册 RPC 服务失败: %w", err)
}
rpcSubs[serviceName] = sub
metrics.SubscribeCount.Add(1)
g.Log().Infof(context.Background(), "✅ RPC 服务已注册: %s", serviceName)
return nil
}
// RegisterQueueRPCService 注册 RPC 服务(集群模式)
// 多个服务实例注册同一服务时,请求会自动负载均衡
// serviceName: 服务名称
// queueName: 队列组名,同一队列组的实例共享请求
// handler: 服务处理函数
func RegisterQueueRPCService(serviceName, queueName string, handler RPCHandler) error {
if !checkConnected() {
return fmt.Errorf("NATS 未连接")
}
queueRPCMu.Lock()
if queueRPCServices == nil {
queueRPCServices = make(map[string]map[string]RPCHandler)
}
if queueRPCSubs == nil {
queueRPCSubs = make(map[string]map[string]*nats.Subscription)
}
if queueRPCServices[queueName] == nil {
queueRPCServices[queueName] = make(map[string]RPCHandler)
}
if queueRPCSubs[queueName] == nil {
queueRPCSubs[queueName] = make(map[string]*nats.Subscription)
}
// 如果已存在该服务,先取消之前的订阅
if oldSub, exists := queueRPCSubs[queueName][serviceName]; exists {
oldSub.Unsubscribe()
}
queueRPCServices[queueName][serviceName] = handler
queueRPCMu.Unlock()
// 订阅服务主题(队列模式)
subject := fmt.Sprintf("rpc.%s", serviceName)
sub, err := nc.QueueSubscribe(subject, queueName, func(msg *nats.Msg) {
ctx := context.Background()
response, err := handler(ctx, msg.Data)
if err != nil {
errMsg := fmt.Sprintf("处理失败: %v", err)
if err = msg.Respond([]byte(errMsg)); err != nil {
g.Log().Errorf(ctx, "RPC 错误响应失败: %v", err)
}
return
}
if err = msg.Respond(response); err != nil {
g.Log().Errorf(ctx, "RPC 响应失败: %v", err)
}
})
if err != nil {
return fmt.Errorf("注册队列 RPC 服务失败: %w", err)
}
queueRPCMu.Lock()
queueRPCSubs[queueName][serviceName] = sub
queueRPCMu.Unlock()
metrics.SubscribeCount.Add(1)
g.Log().Infof(context.Background(), "✅ 队列 RPC 服务已注册: %s (队列组: %s)", serviceName, queueName)
return nil
}
// CallRPC 调用 RPC 服务
// serviceName: 服务名称
// req: 请求数据
// timeout: 超时时间
// 返回: 响应数据和错误
func CallRPC(ctx context.Context, serviceName string, req []byte, timeout time.Duration) ([]byte, error) {
if !checkConnected() {
return nil, fmt.Errorf("NATS 未连接")
}
metrics.RequestCount.Add(1)
// 检查本地是否有注册的单实例服务,如果有则直接调用(优化性能)
rpcServicesMu.RLock()
if localHandler, exists := rpcServices[serviceName]; exists {
rpcServicesMu.RUnlock()
// 本地直接调用,避免网络开销
response, err := localHandler(ctx, req)
if err != nil {
metrics.RequestError.Add(1)
return nil, fmt.Errorf("本地调用 RPC 服务失败 [%s]: %w", serviceName, err)
}
return response, nil
}
rpcServicesMu.RUnlock()
// 通过 NATS 网络调用远程服务
timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
subject := fmt.Sprintf("rpc.%s", serviceName)
msg, err := nc.RequestWithContext(timeoutCtx, subject, req)
if err != nil {
metrics.RequestError.Add(1)
return nil, fmt.Errorf("调用 RPC 服务失败 [%s]: %w", serviceName, err)
}
if msg == nil {
metrics.RequestError.Add(1)
return nil, fmt.Errorf("RPC 响应为空 [%s]", serviceName)
}
return msg.Data, nil
}
// RegisterServiceOption 注册选项类型
type RegisterServiceOption func(*registerServiceConfig)
type registerServiceConfig struct {
queueName string // 队列组名(用于集群模式)
excludeMethods []string
}
// WithQueueGroup 设置队列组名(集群模式)
func WithQueueGroup(queueName string) RegisterServiceOption {
return func(cfg *registerServiceConfig) {
cfg.queueName = queueName
}
}
// WithExcludeMethods 排除不需要注册的方法
func WithExcludeMethods(methods ...string) RegisterServiceOption {
return func(cfg *registerServiceConfig) {
cfg.excludeMethods = append(cfg.excludeMethods, methods...)
}
}
// registerService 注册单个服务的所有公开方法(内部函数)
func registerService(service interface{}, serviceNamePrefix string, options ...RegisterServiceOption) error {
if !checkConnected() {
return fmt.Errorf("NATS 未连接")
}
// 应用选项
cfg := &registerServiceConfig{}
for _, opt := range options {
opt(cfg)
}
// 创建排除方法集合
excludeSet := make(map[string]struct{})
for _, method := range cfg.excludeMethods {
excludeSet[method] = struct{}{}
}
// 获取 service 的类型
serviceType := reflect.TypeOf(service)
// 遍历所有方法
registeredCount := 0
for i := 0; i < serviceType.NumMethod(); i++ {
method := serviceType.Method(i)
// 只注册导出方法(首字母大写)
if !method.IsExported() {
continue
}
// 排除指定的方法
if _, exists := excludeSet[method.Name]; exists {
continue
}
// 检查方法签名:必须是 func(ctx context.Context, request) (response, error)
if method.Type.NumIn() < 2 {
g.Log().Warningf(context.Background(), "方法 %s 的参数数量不足,跳过注册", method.Name)
continue
}
// 第一个参数必须是 context.Context
if !method.Type.In(0).Implements(reflect.TypeOf((*context.Context)(nil)).Elem()) {
g.Log().Warningf(context.Background(), "方法 %s 的第一个参数必须是 context.Context跳过注册", method.Name)
continue
}
// 返回值必须是 (result, error) 或 error
if method.Type.NumOut() < 1 || method.Type.NumOut() > 2 {
g.Log().Warningf(context.Background(), "方法 %s 的返回值数量不正确,跳过注册", method.Name)
continue
}
if !method.Type.Out(method.Type.NumOut() - 1).Implements(reflect.TypeOf((*error)(nil)).Elem()) {
g.Log().Warningf(context.Background(), "方法 %s 的最后一个返回值必须是 error跳过注册", method.Name)
continue
}
// 生成服务名称:前缀.方法名(保持原始方法名)
serviceName := fmt.Sprintf("%s.%s", serviceNamePrefix, method.Name)
// 创建 RPC handler
handler := func(ctx context.Context, req []byte) ([]byte, error) {
// 准备方法调用参数
args := make([]reflect.Value, 2)
args[0] = reflect.ValueOf(ctx)
// 解析请求参数
if len(req) > 0 {
// 如果方法有第二个参数,尝试解析 JSON
if method.Type.NumIn() > 1 {
reqValuePtr := reflect.New(method.Type.In(1))
if err := json.Unmarshal(req, reqValuePtr.Interface()); err != nil {
return nil, fmt.Errorf("解析请求参数失败: %w", err)
}
args[1] = reqValuePtr.Elem()
}
} else if method.Type.NumIn() > 1 {
// 如果方法需要参数但请求为空,创建零值
args[1] = reflect.Zero(method.Type.In(1))
}
// 调用方法
results := method.Func.Call(args)
// 处理返回值
var err error
var result interface{}
if len(results) == 1 {
// 只有 error
if !results[0].IsNil() {
err = results[0].Interface().(error)
}
} else if len(results) == 2 {
// (result, error)
result = results[0].Interface()
if !results[1].IsNil() {
err = results[1].Interface().(error)
}
}
if err != nil {
return nil, err
}
// 序列化返回值
if result == nil || (reflect.ValueOf(result).Kind() == reflect.Ptr && reflect.ValueOf(result).IsNil()) {
return []byte("{}"), nil
}
return json.Marshal(result)
}
// 注册 RPC 服务
var err error
if cfg.queueName != "" {
err = RegisterQueueRPCService(serviceName, cfg.queueName, handler)
} else {
err = RegisterRPCService(serviceName, handler)
}
if err != nil {
g.Log().Errorf(context.Background(), "注册服务 %s 失败: %v", serviceName, err)
continue
}
registeredCount++
g.Log().Infof(context.Background(), "✅ 已自动注册 RPC 服务: %s -> %s", serviceName, method.Name)
}
if registeredCount == 0 {
g.Log().Warningf(context.Background(), "未注册任何方法,请检查 %v 的方法签名", serviceNamePrefix)
return fmt.Errorf("未找到可注册的方法")
}
g.Log().Infof(context.Background(), "✅ Service %v 共注册了 %d 个 RPC 方法", serviceNamePrefix, registeredCount)
return nil
}
// AutoRegisterServices 自动注册多个服务的所有公开方法
// serviceInstances: map[包名]service实例如 map[string]interface{}{"user": userService, "order": orderService}
// options: 注册选项(可选)
// 示例:
//
// AutoRegisterServices(map[string]interface{}{
// "user": userService,
// "order": orderService,
// })
// 或
// AutoRegisterServices(map[string]interface{}{
// "order": orderService,
// }, WithQueueGroup("order-group"))
func AutoRegisterServices(serviceInstances map[string]interface{}, options ...RegisterServiceOption) error {
if len(serviceInstances) == 0 {
return fmt.Errorf("service 实例列表不能为空")
}
totalRegistered := 0
// 遍历每个 service 实例
for pkgName, serviceInstance := range serviceInstances {
// 注册服务
err := registerService(serviceInstance, pkgName, options...)
if err != nil {
g.Log().Errorf(context.Background(), "注册 %s 服务失败: %v", pkgName, err)
continue
}
totalRegistered++
g.Log().Infof(context.Background(), "✅ %s 服务已自动注册", pkgName)
}
if totalRegistered == 0 {
return fmt.Errorf("未能注册任何服务")
}
g.Log().Infof(context.Background(), "✅ 共自动注册了 %d 个服务", totalRegistered)
return nil
}
// Shutdown 优雅关闭:自动注销所有已注册的服务并关闭 NATS 连接
func Shutdown() error {
ctx := context.Background()
g.Log().Info(ctx, "开始优雅关闭 NATS RPC 服务...")
// 注销所有单实例服务
rpcServicesMu.Lock()
singleServiceCount := len(rpcServices)
for serviceName := range rpcServices {
if sub, exists := rpcSubs[serviceName]; exists {
if err := sub.Unsubscribe(); err != nil {
g.Log().Errorf(ctx, "注销服务 %s 失败: %v", serviceName, err)
}
}
delete(rpcSubs, serviceName)
delete(rpcServices, serviceName)
}
rpcServicesMu.Unlock()
// 注销所有队列服务
queueRPCMu.Lock()
queueServiceCount := 0
for queueName, servicesMap := range queueRPCServices {
queueServiceCount += len(servicesMap)
for serviceName, sub := range queueRPCSubs[queueName] {
if err := sub.Unsubscribe(); err != nil {
g.Log().Errorf(ctx, "注销队列服务 %s (队列: %s) 失败: %v", serviceName, queueName, err)
}
}
delete(queueRPCSubs, queueName)
delete(queueRPCServices, queueName)
}
queueRPCMu.Unlock()
g.Log().Infof(ctx, "已注销 %d 个单实例服务和 %d 个队列服务", singleServiceCount, queueServiceCount)
mu.Lock()
defer mu.Unlock()
// 停止健康检查协程
if healthCancel != nil {
healthCancel()
}
// 关闭连接
if nc != nil && !nc.IsClosed() {
nc.Close()
connected = false
inited = false
}
g.Log().Info(ctx, "NATS RPC 服务已优雅关闭")
return nil
}