RAGflow的init函数修改和mq

rabbitmq/client.go

@@ -0,0 +1,179 @@
+package rabbitmq
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/gogf/gf/v2/frame/g"
+	amqp "github.com/rabbitmq/amqp091-go"
+)
+
+var (
+	conn    *amqp.Connection
+	channel *amqp.Channel
+	once    sync.Once
+	mu      sync.RWMutex
+)
+
+// Config RabbitMQ 配置
+type Config struct {
+	Host     string
+	Port     int
+	Username string
+	Password string
+	VHost    string
+}
+
+// Init 初始化 RabbitMQ 连接
+func Init(ctx context.Context, cfg *Config) error {
+	var err error
+	once.Do(func() {
+		// 构建连接字符串
+		url := fmt.Sprintf("amqp://%s:%s@%s:%d/%s",
+			cfg.Username,
+			cfg.Password,
+			cfg.Host,
+			cfg.Port,
+			cfg.VHost,
+		)
+
+		// 创建连接
+		conn, err = amqp.Dial(url)
+		if err != nil {
+			g.Log().Errorf(ctx, "RabbitMQ 连接失败: %v", err)
+			return
+		}
+
+		// 创建 Channel
+		channel, err = conn.Channel()
+		if err != nil {
+			g.Log().Errorf(ctx, "创建 RabbitMQ Channel 失败: %v", err)
+			return
+		}
+
+		// 监听连接关闭
+		go handleConnectionClose(ctx)
+
+		g.Log().Info(ctx, "RabbitMQ 连接成功")
+	})
+
+	return err
+}
+
+// InitFromConfig 从配置文件初始化
+func InitFromConfig(ctx context.Context) error {
+	cfg := &Config{
+		Host:     g.Cfg().MustGet(ctx, "rabbitmq.host").String(),
+		Port:     g.Cfg().MustGet(ctx, "rabbitmq.port").Int(),
+		Username: g.Cfg().MustGet(ctx, "rabbitmq.username").String(),
+		Password: g.Cfg().MustGet(ctx, "rabbitmq.password").String(),
+		VHost:    g.Cfg().MustGet(ctx, "rabbitmq.vhost", "/").String(),
+	}
+
+	return Init(ctx, cfg)
+}
+
+// GetChannel 获取 Channel
+func GetChannel() (*amqp.Channel, error) {
+	mu.RLock()
+	defer mu.RUnlock()
+
+	if channel == nil || channel.IsClosed() {
+		return nil, fmt.Errorf("RabbitMQ Channel 未初始化或已关闭")
+	}
+
+	return channel, nil
+}
+
+// GetConnection 获取连接
+func GetConnection() (*amqp.Connection, error) {
+	mu.RLock()
+	defer mu.RUnlock()
+
+	if conn == nil || conn.IsClosed() {
+		return nil, fmt.Errorf("RabbitMQ 连接未初始化或已关闭")
+	}
+
+	return conn, nil
+}
+
+// handleConnectionClose 监听连接关闭并重连
+func handleConnectionClose(ctx context.Context) {
+	closeErr := make(chan *amqp.Error)
+	conn.NotifyClose(closeErr)
+
+	err := <-closeErr
+	if err != nil {
+		g.Log().Errorf(ctx, "RabbitMQ 连接关闭: %v，尝试重连...", err)
+		reconnect(ctx)
+	}
+}
+
+// reconnect 重新连接
+func reconnect(ctx context.Context) {
+	mu.Lock()
+	defer mu.Unlock()
+
+	for i := 0; i < 10; i++ {
+		time.Sleep(time.Duration(i+1) * time.Second)
+
+		cfg := &Config{
+			Host:     g.Cfg().MustGet(ctx, "rabbitmq.host").String(),
+			Port:     g.Cfg().MustGet(ctx, "rabbitmq.port").Int(),
+			Username: g.Cfg().MustGet(ctx, "rabbitmq.username").String(),
+			Password: g.Cfg().MustGet(ctx, "rabbitmq.password").String(),
+			VHost:    g.Cfg().MustGet(ctx, "rabbitmq.vhost", "/").String(),
+		}
+
+		url := fmt.Sprintf("amqp://%s:%s@%s:%d/%s",
+			cfg.Username,
+			cfg.Password,
+			cfg.Host,
+			cfg.Port,
+			cfg.VHost,
+		)
+
+		var err error
+		conn, err = amqp.Dial(url)
+		if err != nil {
+			g.Log().Errorf(ctx, "重连失败 (尝试 %d/10): %v", i+1, err)
+			continue
+		}
+
+		channel, err = conn.Channel()
+		if err != nil {
+			g.Log().Errorf(ctx, "创建 Channel 失败 (尝试 %d/10): %v", i+1, err)
+			continue
+		}
+
+		g.Log().Info(ctx, "RabbitMQ 重连成功")
+		go handleConnectionClose(ctx)
+		return
+	}
+
+	g.Log().Fatal(ctx, "RabbitMQ 重连失败，已达到最大重试次数")
+}
+
+// Close 关闭连接
+func Close(ctx context.Context) error {
+	mu.Lock()
+	defer mu.Unlock()
+
+	if channel != nil {
+		if err := channel.Close(); err != nil {
+			g.Log().Errorf(ctx, "关闭 RabbitMQ Channel 失败: %v", err)
+		}
+	}
+
+	if conn != nil {
+		if err := conn.Close(); err != nil {
+			g.Log().Errorf(ctx, "关闭 RabbitMQ 连接失败: %v", err)
+			return err
+		}
+	}
+
+	g.Log().Info(ctx, "RabbitMQ 连接已关闭")
+	return nil
+}