iotlcm/server/utility/queue/queue.go

package queue

import (
	"context"
	"encoding/json"
	"errors"
	"sync"

	"github.com/gogf/gf/v2/frame/g"
)

var ErrNotFoundTopic = errors.New("not found topic")

type Payload struct {
	Ts       int64           `json:"ts"`        // 毫秒时间戳
	TraceId  string          `json:"trace_id"`  // 链路ID
	Topic    string          `json:"topic"`     // 消费主题 队列名
	Data     json.RawMessage `json:"data"`      // 数据
	Expire   int64           `json:"expire"`    // 过期时长毫秒 0:不过期 业务内自行处理
	ClientId string          `json:"client_id"` // 客户端ID
}

type Config struct {
	RedisName string `json:"redis_name"` // redis名称
	Workers   int    `json:"workers"`    // 工作线程数
	ClientId  string `json:"client_id"`  // 客户端ID
}

type Process interface {
	GetTopic() string                                  // 获取消费主题
	GetType() QueueType                                // 获取类型
	Handle(ctx context.Context, p Payload) (err error) // 处理过程的方法
}

type QueueType string

const (
	QueueTypeMemory         QueueType = "memory_queue"
	QueueTypeRedis          QueueType = "redis_queue"
	QueueTypeRedisBroadcast QueueType = "redis_broadcast"
)

type sProducer interface {
	GetType() QueueType                                                                 // 获取类型
	GetTopic() string                                                                   // 获取主题
	Push(ctx context.Context, topic string, data interface{}, expire int64) (err error) // 推送消息
}

type sConsumer interface {
	GetType() QueueType
	GetTopic() string                       // 获取消费主题
	Handle(ctx context.Context) (err error) // 处理过程的方法
}

var (
	rwLock   sync.RWMutex
	producer = make(map[string]sProducer)
	consumer = make(map[string]sConsumer)
)

func NewProducer(ctx context.Context, t QueueType, topic string, config *Config) error {
	if config == nil {
		config = &Config{}
	}
	if _, ok := producer[topic]; ok {
		return errors.New("producer already exists")
	}
	var p sProducer
	switch t {
	case QueueTypeMemory:
		p = RegisterMemoryQueue(ctx, topic, nil, config)
	case QueueTypeRedis:
		p = RegisterRedisQueue(ctx, topic, nil, config)
	case QueueTypeRedisBroadcast:
		p = RegisterRedisBroadcast(ctx, topic, nil, config)
	}
	rwLock.Lock()
	producer[topic] = p
	rwLock.Unlock()
	return nil
}

func NewConsumer(ctx context.Context, p Process, config *Config) error {
	if config == nil {
		config = &Config{}
	}
	if _, ok := consumer[p.GetTopic()+"__"+config.ClientId]; ok {
		return errors.New("consumer already exists")
	}
	var c sConsumer
	switch p.GetType() {
	case QueueTypeMemory:
		c = RegisterMemoryQueue(ctx, p.GetTopic(), p, config)
	case QueueTypeRedis:
		c = RegisterRedisQueue(ctx, p.GetTopic(), p, config)
	case QueueTypeRedisBroadcast:
		c = RegisterRedisBroadcast(ctx, p.GetTopic(), p, config)
	}
	rwLock.Lock()
	consumer[p.GetTopic()+"__"+config.ClientId] = c
	rwLock.Unlock()
	return nil
}

func UnregisterProducer(ctx context.Context, topic string) error {
	rwLock.Lock()
	defer rwLock.Unlock()
	p, ok := producer[topic]
	if !ok {
		return errors.New("producer not found")
	}
	switch p.GetType() {
	case QueueTypeMemory:
		UnregisterMemoryQueue(ctx, topic)
	case QueueTypeRedis:
		UnregisterRedisQueue(ctx, topic)
	case QueueTypeRedisBroadcast:
		UnregisterRedisBroadcast(ctx, topic, "")
	}
	delete(producer, topic)
	return nil
}

func UnregisterConsumer(ctx context.Context, topic string, groupId string) error {
	rwLock.Lock()
	defer rwLock.Unlock()
	c, ok := consumer[topic+"__"+groupId]
	if !ok {
		return ErrNotFoundTopic
	}
	switch c.GetType() {
	case QueueTypeMemory:
		UnregisterMemoryQueue(ctx, topic)
	case QueueTypeRedis:
		UnregisterRedisQueue(ctx, topic)
	case QueueTypeRedisBroadcast:
		UnregisterRedisBroadcast(ctx, topic, groupId)
	}
	delete(consumer, topic+"__"+groupId)
	return nil
}

func ProducerTopicExists(topic string) bool {
	rwLock.RLock()
	_, ok := producer[topic]
	rwLock.RUnlock()
	return ok
}

func ConsumerTopicExists(topic string) bool {
	rwLock.RLock()
	_, ok := consumer[topic]
	rwLock.RUnlock()
	return ok
}

func Push(ctx context.Context, topic string, data interface{}) (err error) {
	rwLock.RLock()
	p, ok := producer[topic]
	rwLock.RUnlock()
	if !ok {
		g.Log().Errorf(ctx, "topic not found. topic: %s", topic)
		return ErrNotFoundTopic
	}
	return p.Push(ctx, topic, data, 0)
}

func PushWithExpire(ctx context.Context, topic string, data interface{}, expire int64) (err error) {
	rwLock.RLock()
	p, ok := producer[topic]
	rwLock.RUnlock()
	if !ok {
		g.Log().Errorf(ctx, "topic not found. topic: %s", topic)
		return ErrNotFoundTopic
	}
	return p.Push(ctx, topic, data, expire)
}