iotlcm/server/utility/queue/redis_broadcast.go

package queue

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"sync"
	"time"

	"git.5iot.com/iotlcm-public/iotlcm-app/utility/gfredis"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/gtrace"
	"github.com/gogf/gf/v2/os/gctx"
	"github.com/gogf/gf/v2/os/grpool"
	"github.com/redis/go-redis/v9"
)

type tRedisBroadcast struct {
	p         Process
	RedisName string
	isRunning bool
	grpool    *grpool.Pool
	clientId  string
	pubSub    *redis.PubSub
}

var (
	redisBroadcastRwLock sync.RWMutex
	redisBroadcastQueues = make(map[string]*tRedisBroadcast)
)

func RegisterRedisBroadcast(ctx context.Context, topic string, p Process, config *Config) *tRedisBroadcast {
	redisBroadcastRwLock.Lock()
	defer redisBroadcastRwLock.Unlock()
	q, ok := redisBroadcastQueues[topic+"__"+config.ClientId]
	if ok {
		q.p = p
		q.RedisName = config.RedisName
		q.isRunning = true
		q.clientId = config.ClientId
	} else {
		q = &tRedisBroadcast{
			p:         p,
			RedisName: config.RedisName,
			isRunning: true,
			clientId:  config.ClientId,
		}
		redisBroadcastQueues[topic+"__"+config.ClientId] = q
	}
	workers := 1
	if config != nil && config.Workers > 0 {
		workers = config.Workers
	}
	if p != nil {
		q.p = p
		q.grpool = grpool.New(workers)

		q.pubSub, _ = q.getConn(ctx)
		go q.Handle(ctx)
	}
	return q
}

func UnregisterRedisBroadcast(ctx context.Context, topic string, clientId string) {
	redisBroadcastRwLock.Lock()
	defer redisBroadcastRwLock.Unlock()
	q, ok := redisBroadcastQueues[topic+"__"+clientId]
	if !ok {
		return
	}
	q.isRunning = false
	delete(redisBroadcastQueues, topic+"__"+clientId)
}

func (q *tRedisBroadcast) GetType() QueueType {
	return QueueTypeRedis
}

func (q *tRedisBroadcast) GetTopic() string {
	return q.p.GetTopic()
}

func (q *tRedisBroadcast) Push(ctx context.Context, topic string, data interface{}, expire int64) (err error) {
	jsonData, err := json.Marshal(data)
	if err != nil {
		return err
	}
	payload := Payload{
		Ts:      time.Now().UnixMilli(),
		TraceId: gctx.CtxId(ctx),
		Topic:   topic,
		Data:    json.RawMessage(jsonData),
		Expire:  expire,
	}
	payloadJson, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	num, err := gfredis.Instance(q.RedisName).Publish(ctx, topic, payloadJson)
	if err != nil {
		return err
	}
	g.Log().Debugf(ctx, "redisqueue Push %s %d", topic, num)
	return
}

func (q *tRedisBroadcast) getConn(ctx context.Context) (conn *redis.PubSub, err error) {

	conn, err = gfredis.Instance(q.RedisName).Subscribe(ctx, q.p.GetTopic())
	if err != nil || conn == nil {
		g.Log().Errorf(ctx, "tRedisBroadcast.Handle Conn %s error: %v", q.p.GetTopic(), err)
		return
	}
	g.Log().Debugf(ctx, "tRedisBroadcast.Handle Subscribe %s 成功, sub:%+v", q.p.GetTopic(), conn)
	return
}

func (q *tRedisBroadcast) Handle(ctx context.Context) (err error) {
	g.Log().Debugf(ctx, "tRedisBroadcast.Handle 开始处理 %s, group_id: %s", q.p.GetTopic(), q.clientId)
	topic := q.p.GetTopic()

	for {
		if !q.isRunning {
			break
		}
		if q.pubSub == nil {
			q.pubSub, err = q.getConn(ctx)
			if err != nil {
				g.Log().Errorf(ctx, "tRedisBroadcast.Handle getConn %s error: %v", topic, err)
				time.Sleep(2 * time.Second)
				continue
			}
		}
		data, err := q.pubSub.ReceiveMessage(ctx)
		if err != nil {
			g.Log().Errorf(ctx, "tRedisBroadcast.Handle ReceiveMessage %s error: %v", topic, err)
			q.pubSub.Close()
			q.pubSub = nil
			continue
		}
		var payload Payload
		err = json.Unmarshal([]byte(data.Payload), &payload)
		if err != nil {
			g.Log().Errorf(ctx, "tRedisBroadcast.Handle unmarshal error: %v", data)
			continue
		}
		payload.ClientId = q.clientId
		subCtx := gctx.New()
		gtrace.WithTraceID(subCtx, payload.TraceId)
		err = q.grpool.Add(subCtx, func(ctx context.Context) {
			defer func() {
				if r := recover(); r != nil {
					g.Log().Errorf(ctx, "tRedisBroadcast.Handle process panic: %v, topic: %s, trace_id: %s, len(data): %d, hex(data): %v",
						r, q.p.GetTopic(), payload.TraceId, len(payload.Data), hex.EncodeToString(payload.Data))
				}
			}()
			err = q.p.Handle(subCtx, payload)
			if err != nil {
				g.Log().Errorf(ctx, "tRedisBroadcast.Handle process error: %v, topic: %s, trace_id: %s, len(data): %d, hex(data): %v",
					err, q.p.GetTopic(), payload.TraceId, len(payload.Data), hex.EncodeToString(payload.Data))
				return
			}
		})
	}
	redisBroadcastRwLock.Lock()
	delete(redisBroadcastQueues, topic)
	redisBroadcastRwLock.Unlock()
	if q.pubSub != nil {
		q.pubSub.Close()
	}
	return
}