JCS-pub/datamap/internal/mq/mq.go

package mq

import (
	"fmt"
	jsoniter "github.com/json-iterator/go"
	"github.com/streadway/amqp"
	stgmod "gitlink.org.cn/cloudream/storage/common/models"
	"gitlink.org.cn/cloudream/storage/datamap/internal/config"
	"gitlink.org.cn/cloudream/storage/datamap/internal/models"
	"log"
)

func InitMQ(cfg config.RabbitMQConfig) (*amqp.Connection, error) {
	conn, err := amqp.Dial(fmt.Sprintf("amqp://%s:%s@%s:%s/",
		cfg.User, cfg.Password, cfg.Host, cfg.Port))
	if err != nil {
		return nil, err
	}

	// 启动队列监听
	go listenQueues(conn)

	return conn, nil
}

func listenQueues(conn *amqp.Connection) {
	queues := []string{
		"datamap_hubinfo",
		"datamap_storageinfo",
		"datamap_storagestats",
		"datamap_hubtransferstats",
		"datamap_hubstoragetransferstats",
		"datamap_blocktransfer",
		"datamap_blockdistribution",
	}

	for _, queue := range queues {
		go func(q string) {
			ch, err := conn.Channel()
			if err != nil {
				log.Printf("Failed to open channel for queue %s: %v", q, err)
				return
			}
			defer ch.Close()

			msgs, err := ch.Consume(q, "", true, false, false, false, nil)
			if err != nil {
				log.Printf("Failed to register consumer for queue %s: %v", q, err)
				return
			}

			for msg := range msgs {
				processMessage(q, msg.Body)
			}
		}(queue)
	}
}

func processMessage(queue string, body []byte) {
	switch queue {
	case "datamap_hubinfo":
		var data stgmod.HubInfo
		if err := jsoniter.Unmarshal(body, &data); err != nil {
			log.Printf("Failed to unmarshal HubInfo: %v, body: %s", err, body)
			return
		}
		models.ProcessHubInfo(data)
	case "datamap_storageinfo":
		var data stgmod.StorageInfo
		if err := jsoniter.Unmarshal(body, &data); err != nil {
			log.Printf("Failed to unmarshal StorageInfo: %v, body: %s", err, body)
			return
		}
		models.ProcessStorageInfo(data)
	case "datamap_storagestats":
		var data stgmod.StorageStats
		if err := jsoniter.Unmarshal(body, &data); err != nil {
			log.Printf("Failed to unmarshal StorageStats: %v, body: %s", err, body)
			return
		}
		models.ProcessStorageStats(data)
	case "datamap_hubtransferstats":
		var data stgmod.HubTransferStats
		err := jsoniter.Unmarshal(body, &data)
		if err != nil {
			log.Printf("Failed to unmarshal HubTransferStats: %v, body: %s", err, body)
			return
		}
		models.ProcessHubTransfer(data)
	case "datamap_hubstoragetransferstats":
		var data stgmod.HubStorageTransferStats
		err := jsoniter.Unmarshal(body, &data)
		if err != nil {
			log.Printf("Failed to unmarshal HubStorageTransferStats: %v, body: %s", err, body)
			return
		}
		models.ProcessHubStorageTransfer(data)
	case "datamap_blocktransfer":
		var data stgmod.BlockTransfer
		err := jsoniter.Unmarshal(body, &data)
		if err != nil {
			log.Printf("Failed to unmarshal BlockTransfer: %v, body: %s", err, body)
			return
		}
		models.ProcessBlockTransfer(data)
	case "datamap_blockdistribution":
		var data stgmod.BlockDistribution
		err := jsoniter.Unmarshal(body, &data)
		if err != nil {
			log.Printf("Failed to unmarshal BlockDistribution: %v, body: %s", err, body)
			return
		}
		models.ProcessBlockDistribution(data)
	default:
		log.Printf("Unknown queue: %s", queue)
	}
}