FranklinWH-Toolbox/service/cache/redisx/streamx.go

package redisx

import (
	"context"
	"time"

	"github.com/redis/go-redis/v9"
)

// Stream（流）封装

// 添加消息
func XAdd(ctx context.Context, stream string, values map[string]interface{}) (string, error) {
	return getClient().XAdd(ctx, &redis.XAddArgs{
		Stream: stream,
		Values: values,
	}).Result()
}

// 读取消息
func XRead(ctx context.Context, streams map[string]string, count int64) ([]redis.XStream, error) {
	args := &redis.XReadArgs{
		Streams: flattenStreams(streams),
		Count:   count,
		Block:   0,
	}
	return getClient().XRead(ctx, args).Result()
}

// 创建消费者组
func XGroupCreate(ctx context.Context, stream, group, start string) error {
	return getClient().XGroupCreateMkStream(ctx, stream, group, start).Err()
}

// 从消费者组读取
func XReadGroup(ctx context.Context, group, consumer string, streams map[string]string, count int64) ([]redis.XStream, error) {
	args := &redis.XReadGroupArgs{
		Group:    group,
		Consumer: consumer,
		Streams:  flattenStreams(streams),
		Count:    count,
		Block:    5 * time.Second,
	}
	return getClient().XReadGroup(ctx, args).Result()
}

// 从消费者组读取滞留消息
func XAutoClaim(ctx context.Context, stream string, group string, consumer string, idle time.Duration, count int64) (messages []redis.XMessage, start string, err error) {
	args := &redis.XAutoClaimArgs{
		Stream:   stream,
		Group:    group,
		Consumer: consumer,
		Start:    "-", // 从最小ID开始
		Count:    count,
		MinIdle:  idle,
	}
	return getClient().XAutoClaim(ctx, args).Result()
}

// 从消费者组读取指定ID的Pending状态信息
func XPendingExtOne(ctx context.Context, stream string, group string, consumer string, id string) (*redis.XPendingExt, error) {
	args := &redis.XPendingExtArgs{
		Stream:   stream,
		Group:    group,
		Consumer: consumer,
		Count:    1,
		Idle:     0,
		Start:    id,
		End:      id,
	}
	res, err := getClient().XPendingExt(ctx, args).Result()
	if err != nil {
		return nil, err
	}
	if len(res) == 0 {
		return nil, nil
	}
	return &res[0], nil
}

// 确认消息
func XAck(ctx context.Context, stream, group string, ids ...string) error {
	return getClient().XAck(ctx, stream, group, ids...).Err()
}

// 确认同时删除消息
func XAckDel(ctx context.Context, stream, group string, ids ...string) error {
	// 版本不支持合并写法
	// return getClient().XAckDel(ctx, stream, group, "DEL", ids...).Err()
	err := getClient().XAck(ctx, stream, group, ids...).Err()
	if err != nil {
		return err
	}
	return getClient().XDel(ctx, stream, ids...).Err()
}

// 查询范围内指定的消息信息
func XRangeN(ctx context.Context, stream, start, stop string, count int64) ([]redis.XMessage, error) {
	return getClient().XRangeN(ctx, stream, start, stop, count).Result()
}

// 获取完整的Stream信息
func XInfoStreamFull(ctx context.Context, key string, count int) (*redis.XInfoStreamFull, error) {
	return getClient().XInfoStreamFull(ctx, key, count).Result()
}

func flattenStreams(m map[string]string) []string {
	arr := make([]string, 0, len(m)*2)
	for k, v := range m {
		arr = append(arr, k, v)
	}
	return arr
}