red-future/common
4
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 19 Wiki Activity

重构消息队列模块，新增NATS和RabbitMQ连接实现，移除旧版消息队列代码

Browse Source
This commit is contained in:
曲皓东
2026-01-31 05:17:14 +08:00
committed by 张斌
parent 6acdbb6e88
commit ee6d3c9033
17 changed files with 966 additions and 2391 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
message/nats_rpc.go
View File

@@ -38,11 +38,11 @@ var (
// 返回值可以是任意类型，会被自动序列化为 JSON
type rpcHandler func(ctx context.Context, req []byte) (any, error)
// RegisterRPCService 注册 RPC 服务（单实例）
// registerRPCService 注册 RPC 服务（单实例）
// serviceName: 服务名称，调用方通过此名称调用服务
// handler: 服务处理函数，接收请求并返回响应
func registerRPCService(serviceName string, handler rpcHandler) (err error) {
if !checkConnected() {
if !natsPing() {
return fmt.Errorf("NATS 未连接")
}
@@ -74,18 +74,17 @@ func registerRPCService(serviceName string, handler rpcHandler) (err error) {
}
rpcSubs[serviceName] = sub
metrics.SubscribeCount.Add(1)
g.Log().Infof(context.Background(), "✅ RPC 服务已注册: %s", serviceName)
return nil
}
// RegisterQueueRPCService 注册 RPC 服务（集群模式）
// registerQueueRPCService 注册 RPC 服务（集群模式）
// 多个服务实例注册同一服务时，请求会自动负载均衡
// serviceName: 服务名称
// queueName: 队列组名，同一队列组的实例共享请求
// handler: 服务处理函数
func registerQueueRPCService(serviceName, queueName string, handler rpcHandler) (err error) {
if !checkConnected() {
if !natsPing() {
return fmt.Errorf("NATS 未连接")
}
@@ -126,7 +125,6 @@ func registerQueueRPCService(serviceName, queueName string, handler rpcHandler)
queueRPCSubs[queueName][serviceName] = sub
queueRPCMu.Unlock()
metrics.SubscribeCount.Add(1)
g.Log().Infof(context.Background(), "✅ 队列 RPC 服务已注册: %s (队列组: %s)", serviceName, queueName)
return nil
}
@@ -138,16 +136,16 @@ func executeHandler(handler rpcHandler, msg *nats.Msg) {
// 从消息头重建上下文
ctx := headersToContext(context.Background(), msg.Header)
// 提取 TraceID，创建可取消的 context
ctx = createCancelContext(ctx, msg.Header.Get(TraceIDKey))
ctx = createCancelContext(ctx, msg.Header.Get(traceIDKey))
// 检查 context 是否已取消（在调用 handler 之前）
select {
case <-ctx.Done():
// context 已取消，返回取消错误
g.Log().Infof(ctx, "RPC 请求已取消，traceID: %s", msg.Header.Get(TraceIDKey))
g.Log().Infof(ctx, "RPC 请求已取消，traceID: %s", msg.Header.Get(traceIDKey))
// 仍然需要发送响应以避免客户端超时
respData = []byte(`{"_err":"请求已取消"}`)
// 清理取消映射表
cleanupTraceCancel(msg.Header.Get(TraceIDKey))
cleanupTraceCancel(msg.Header.Get(traceIDKey))
return
default:
}
@@ -176,7 +174,7 @@ func executeHandler(handler rpcHandler, msg *nats.Msg) {
g.Log().Errorf(ctx, "RPC 响应失败: %v", err)
}
// 请求结束，清理取消映射表
cleanupTraceCancel(msg.Header.Get(TraceIDKey))
cleanupTraceCancel(msg.Header.Get(traceIDKey))
}
// createCancelContext 创建可取消的 context 并注册到取消映射表
@@ -211,7 +209,7 @@ func createCancelContext(ctx context.Context, traceID string) context.Context {
//
// sub, err := nats.SetupCancelListener(ctx)
func setupCancelListener(ctx context.Context) (*nats.Subscription, error) {
if !checkConnected() {
if !natsPing() {
return nil, fmt.Errorf("NATS 未连接")
}
@@ -253,7 +251,6 @@ func setupCancelListener(ctx context.Context) (*nats.Subscription, error) {
return nil, fmt.Errorf("设置取消监听器失败: %w", err)
}
metrics.SubscribeCount.Add(1)
g.Log().Infof(ctx, "✅ 取消监听器已设置: %s", cancelSubject)
return sub, nil
}
@@ -264,7 +261,7 @@ func setupCancelListener(ctx context.Context) (*nats.Subscription, error) {
//
// err := nats.publishCancel(ctx, traceID)
func publishCancel(ctx context.Context, traceID string) error {
if !checkConnected() {
if !natsPing() {
return fmt.Errorf("NATS 未连接")
}
@@ -306,12 +303,10 @@ func cleanupTraceCancel(traceID string) {
// req: 请求数据
// 返回: 响应数据（任意类型）和错误
func CallRPC(ctx context.Context, serviceName string, req any, resp any) (err error) {
if !checkConnected() {
if !natsPing() {
return fmt.Errorf("NATS 未连接")
}
metrics.RequestCount.Add(1)
// 验证 resp 必须是指针类型
respValue := reflect.ValueOf(resp)
if respValue.Kind() != reflect.Ptr {
@@ -346,7 +341,6 @@ func CallRPC(ctx context.Context, serviceName string, req any, resp any) (err er
// 执行本地调用
var response interface{}
if response, err = localHandler(cancelCtx, reqBody); err != nil {
metrics.RequestError.Add(1)
return fmt.Errorf("本地调用 RPC 服务失败 [%s]: %w", serviceName, err)
}
@@ -357,7 +351,6 @@ func CallRPC(ctx context.Context, serviceName string, req any, resp any) (err er
var respMap map[string]any
if json.Unmarshal(response.([]byte), &respMap) == nil {
if errMsg, ok := respMap["_err"]; ok {
metrics.RequestError.Add(1)
return fmt.Errorf("%v", errMsg)
}
}
@@ -392,17 +385,17 @@ func CallRPC(ctx context.Context, serviceName string, req any, resp any) (err er
})
}
if msg.Header.Get(TraceIDKey) != "" {
if msg.Header.Get(traceIDKey) != "" {
go func() {
defer closeDone()
select {
case <-ctx.Done():
// context 被取消时，发送取消信号给服务端
if errors.Is(ctx.Err(), context.Canceled) {
if err := publishCancel(context.Background(), msg.Header.Get(TraceIDKey)); err != nil {
if err := publishCancel(context.Background(), msg.Header.Get(traceIDKey)); err != nil {
g.Log().Errorf(ctx, "发送 RPC 取消信号失败: %v", err)
} else {
g.Log().Infof(ctx, "RPC 调用已取消，traceID: %s", msg.Header.Get(TraceIDKey))
g.Log().Infof(ctx, "RPC 调用已取消，traceID: %s", msg.Header.Get(traceIDKey))
}
}
case <-done:
@@ -419,12 +412,10 @@ func CallRPC(ctx context.Context, serviceName string, req any, resp any) (err er
closeDone()
if err != nil {
metrics.RequestError.Add(1)
return fmt.Errorf("调用 RPC 服务失败 [%s]: %w", serviceName, err)
}
if responseMsg == nil {
metrics.RequestError.Add(1)
return fmt.Errorf("RPC 响应为空 [%s]", serviceName)
}
@@ -434,7 +425,6 @@ func CallRPC(ctx context.Context, serviceName string, req any, resp any) (err er
var respMap map[string]any
if json.Unmarshal(responseMsg.Data, &respMap) == nil {
if errMsg, ok := respMap["_err"]; ok {
metrics.RequestError.Add(1)
return fmt.Errorf("%v", errMsg)
}
}
@@ -449,7 +439,7 @@ func CallRPC(ctx context.Context, serviceName string, req any, resp any) (err er
}
// RegisterServiceOption 注册选项类型
type RegisterServiceOption func(*registerServiceConfig)
type registerServiceOption func(*registerServiceConfig)
type registerServiceConfig struct {
queueName string // 队列组名（用于集群模式）
@@ -457,14 +447,14 @@ type registerServiceConfig struct {
}
// WithQueueGroup 设置队列组名（集群模式）
func WithQueueGroup(queueName string) RegisterServiceOption {
func WithQueueGroup(queueName string) registerServiceOption {
return func(cfg *registerServiceConfig) {
cfg.queueName = queueName
}
}
// WithExcludeMethods 排除不需要注册的方法
func WithExcludeMethods(methods ...string) RegisterServiceOption {
func WithExcludeMethods(methods ...string) registerServiceOption {
return func(cfg *registerServiceConfig) {
cfg.excludeMethods = append(cfg.excludeMethods, methods...)
}
@@ -483,9 +473,9 @@ func WithExcludeMethods(methods ...string) RegisterServiceOption {
// AutoRegisterServices(map[string]interface{}{
// "order": orderService,
// }, WithQueueGroup("order-group"))
func AutoRegisterServices(ctx context.Context, serviceInstances map[string]interface{}, options ...RegisterServiceOption) error {
func AutoRegisterServices(ctx context.Context, serviceInstances map[string]interface{}, options ...registerServiceOption) error {
// 先注册 RPC 服务（如果 NATS 不可用则记录警告但不阻塞启动）
if !checkConnected() {
if !natsPing() {
return fmt.Errorf("NATS 未连接，RPC 服务未注册")
}
@@ -521,8 +511,8 @@ func AutoRegisterServices(ctx context.Context, serviceInstances map[string]inter
}
// registerService 注册单个服务的所有公开方法（内部函数）
func registerService(service interface{}, serviceNamePrefix string, options ...RegisterServiceOption) (err error) {
if !checkConnected() {
func registerService(service interface{}, serviceNamePrefix string, options ...registerServiceOption) (err error) {
if !natsPing() {
return fmt.Errorf("NATS 未连接")
}
@@ -676,10 +666,10 @@ func registerService(service interface{}, serviceNamePrefix string, options ...R
// ============ 上下文元数据工具函数 ============
// 以下函数用于在 context 和 NATS 消息头之间互转元数据
// 定义常见的上下文元数据 key
// 定义常见的上下文元数据 key（私有）
const (
TraceIDKey = "trace_id"
TokenKey = "token"
traceIDKey = "trace_id"
tokenKey = "token"
)
func getTraceID(ctx context.Context) (traceID string, err error) {
@@ -687,7 +677,7 @@ func getTraceID(ctx context.Context) (traceID string, err error) {
span := trace.SpanFromContext(ctx)
if span != nil && span.SpanContext().HasTraceID() {
traceID = span.SpanContext().TraceID().String()
} else if tid := ctx.Value(TraceIDKey); tid != nil {
} else if tid := ctx.Value(traceIDKey); tid != nil {
traceID = fmt.Sprintf("%v", tid)
}
if traceID == "" {
@@ -705,12 +695,12 @@ func contextToHeaders(ctx context.Context) (nats.Header, error) {
if traceID, err := getTraceID(ctx); err != nil {
return headers, err
} else {
headers.Set(TraceIDKey, traceID)
headers.Set(traceIDKey, traceID)
}
// 提取 token（优先级：context value > HTTP Authorization header）
token := ""
if t := ctx.Value(TokenKey); t != nil {
if t := ctx.Value(tokenKey); t != nil {
token = fmt.Sprintf("%v", t)
} else if r := g.RequestFromCtx(ctx); r != nil {
// 从 HTTP 请求的 Authorization header 中提取 token
@@ -725,7 +715,7 @@ func contextToHeaders(ctx context.Context) (nats.Header, error) {
}
}
if token != "" {
headers.Set(TokenKey, token)
headers.Set(tokenKey, token)
}
return headers, nil
@@ -739,13 +729,13 @@ func headersToContext(ctx context.Context, headers nats.Header) context.Context
}
// 恢复 trace_id
if traceID := headers.Get(TraceIDKey); traceID != "" {
ctx = context.WithValue(ctx, TraceIDKey, traceID)
if traceID := headers.Get(traceIDKey); traceID != "" {
ctx = context.WithValue(ctx, traceIDKey, traceID)
}
// 恢复 token
if token := headers.Get(TokenKey); token != "" {
ctx = context.WithValue(ctx, TokenKey, token)
if token := headers.Get(tokenKey); token != "" {
ctx = context.WithValue(ctx, tokenKey, token)
}
return ctx