iotlcm/server/utility/queue/memory_queue.go

package queue

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

	"github.com/gogf/gf/v2/container/gqueue"
	"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"
)

type tMemoryQueue struct {
	q      *gqueue.Queue
	p      Process
	grpool *grpool.Pool
}

var (
	memRwLock sync.RWMutex
	memQueues = make(map[string]*tMemoryQueue)
)

func RegisterMemoryQueue(ctx context.Context, topic string, p Process, config *Config) *tMemoryQueue {
	memRwLock.Lock()
	defer memRwLock.Unlock()
	q, ok := memQueues[topic]
	if !ok {
		q = &tMemoryQueue{
			q: gqueue.New(),
			p: p,
		}
		memQueues[topic] = q
	}
	workers := 1
	if config != nil && config.Workers > 0 {
		workers = config.Workers
	}
	if p != nil {
		q.p = p
		q.grpool = grpool.New(workers)
		go q.Handle(ctx)
	}
	g.Log().Infof(ctx, "RegisterMemoryQueue %s success, workers: %d, process: %v", topic, workers, p)
	return q
}

func UnregisterMemoryQueue(ctx context.Context, topic string) {
	memRwLock.Lock()
	defer memRwLock.Unlock()
	q, ok := memQueues[topic]
	if !ok {
		return
	}
	q.q.Close()
	delete(memQueues, topic)
}

func (q *tMemoryQueue) GetType() QueueType {
	return QueueTypeMemory
}

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

func (q *tMemoryQueue) 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,
	}
	q.q.Push(payload)
	return
}

func (q *tMemoryQueue) Handle(ctx context.Context) (err error) {
	if q.grpool == nil {
		g.Log().Errorf(ctx, "grpool is nil")
		return errors.New("grpool is nil")
	}
	p := q.p
	if p == nil {
		g.Log().Errorf(ctx, "process is nil")
		return errors.New("process is nil")
	}

	for {

		data := q.q.Pop()
		if data == nil {
			break
		}
		payload, ok := data.(Payload)
		if !ok {
			g.Log().Errorf(ctx, "memqueue process unmarshal error: %v", data)
			continue
		}
		subCtx := gctx.New()
		gtrace.WithTraceID(subCtx, payload.TraceId)
		q.grpool.Add(ctx, func(ctx context.Context) {
			defer func() {
				if r := recover(); r != nil {
					g.Log().Errorf(ctx, "memqueue process panic: %v, topic: %s, trace_id: %s, hex(data): %v",
						r, q.p.GetTopic(), payload.TraceId, hex.EncodeToString(payload.Data))
				}
			}()
			err := p.Handle(subCtx, payload)
			if err != nil {
				g.Log().Errorf(ctx, "memqueue process error: %v, topic: %s, trace_id: %s, hex(data): %v",
					err, q.p.GetTopic(), payload.TraceId, hex.EncodeToString(payload.Data))
			}
		})
	}
	return
}