github.com/gitbundle/modules@v0.0.0-20231025071548-85b91c5c3b01/queue/unique_queue_redis.go (about)

     1  // Copyright 2023 The GitBundle Inc. All rights reserved.
     2  // Copyright 2017 The Gitea Authors. All rights reserved.
     3  // Use of this source code is governed by a MIT-style
     4  // license that can be found in the LICENSE file.
     5  
     6  package queue
     7  
     8  import (
     9  	"context"
    10  
    11  	"github.com/redis/go-redis/v9"
    12  )
    13  
    14  // RedisUniqueQueueType is the type for redis queue
    15  const RedisUniqueQueueType Type = "unique-redis"
    16  
    17  // RedisUniqueQueue redis queue
    18  type RedisUniqueQueue struct {
    19  	*ByteFIFOUniqueQueue
    20  }
    21  
    22  // RedisUniqueQueueConfiguration is the configuration for the redis queue
    23  type RedisUniqueQueueConfiguration struct {
    24  	ByteFIFOQueueConfiguration
    25  	RedisUniqueByteFIFOConfiguration
    26  }
    27  
    28  // NewRedisUniqueQueue creates single redis or cluster redis queue.
    29  //
    30  // Please note that this Queue does not guarantee that a particular
    31  // task cannot be processed twice or more at the same time. Uniqueness is
    32  // only guaranteed whilst the task is waiting in the queue.
    33  func NewRedisUniqueQueue(handle HandlerFunc, cfg, exemplar interface{}) (Queue, error) {
    34  	configInterface, err := toConfig(RedisUniqueQueueConfiguration{}, cfg)
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  	config := configInterface.(RedisUniqueQueueConfiguration)
    39  
    40  	byteFIFO, err := NewRedisUniqueByteFIFO(config.RedisUniqueByteFIFOConfiguration)
    41  	if err != nil {
    42  		return nil, err
    43  	}
    44  
    45  	if len(byteFIFO.setName) == 0 {
    46  		byteFIFO.setName = byteFIFO.queueName + "_unique"
    47  	}
    48  
    49  	byteFIFOQueue, err := NewByteFIFOUniqueQueue(RedisUniqueQueueType, byteFIFO, handle, config.ByteFIFOQueueConfiguration, exemplar)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  
    54  	queue := &RedisUniqueQueue{
    55  		ByteFIFOUniqueQueue: byteFIFOQueue,
    56  	}
    57  
    58  	queue.qid = GetManager().Add(queue, RedisUniqueQueueType, config, exemplar)
    59  
    60  	return queue, nil
    61  }
    62  
    63  var _ UniqueByteFIFO = &RedisUniqueByteFIFO{}
    64  
    65  // RedisUniqueByteFIFO represents a UniqueByteFIFO formed from a redisClient
    66  type RedisUniqueByteFIFO struct {
    67  	RedisByteFIFO
    68  	setName string
    69  }
    70  
    71  // RedisUniqueByteFIFOConfiguration is the configuration for the RedisUniqueByteFIFO
    72  type RedisUniqueByteFIFOConfiguration struct {
    73  	RedisByteFIFOConfiguration
    74  	SetName string
    75  }
    76  
    77  // NewRedisUniqueByteFIFO creates a UniqueByteFIFO formed from a redisClient
    78  func NewRedisUniqueByteFIFO(config RedisUniqueByteFIFOConfiguration) (*RedisUniqueByteFIFO, error) {
    79  	internal, err := NewRedisByteFIFO(config.RedisByteFIFOConfiguration)
    80  	if err != nil {
    81  		return nil, err
    82  	}
    83  
    84  	fifo := &RedisUniqueByteFIFO{
    85  		RedisByteFIFO: *internal,
    86  		setName:       config.SetName,
    87  	}
    88  
    89  	return fifo, nil
    90  }
    91  
    92  // PushFunc pushes data to the end of the fifo and calls the callback if it is added
    93  func (fifo *RedisUniqueByteFIFO) PushFunc(ctx context.Context, data []byte, fn func() error) error {
    94  	added, err := fifo.client.SAdd(ctx, fifo.setName, data).Result()
    95  	if err != nil {
    96  		return err
    97  	}
    98  	if added == 0 {
    99  		return ErrAlreadyInQueue
   100  	}
   101  	if fn != nil {
   102  		if err := fn(); err != nil {
   103  			return err
   104  		}
   105  	}
   106  	return fifo.client.RPush(ctx, fifo.queueName, data).Err()
   107  }
   108  
   109  // PushBack pushes data to the top of the fifo
   110  func (fifo *RedisUniqueByteFIFO) PushBack(ctx context.Context, data []byte) error {
   111  	added, err := fifo.client.SAdd(ctx, fifo.setName, data).Result()
   112  	if err != nil {
   113  		return err
   114  	}
   115  	if added == 0 {
   116  		return ErrAlreadyInQueue
   117  	}
   118  	return fifo.client.LPush(ctx, fifo.queueName, data).Err()
   119  }
   120  
   121  // Pop pops data from the start of the fifo
   122  func (fifo *RedisUniqueByteFIFO) Pop(ctx context.Context) ([]byte, error) {
   123  	data, err := fifo.client.LPop(ctx, fifo.queueName).Bytes()
   124  	if err != nil && err != redis.Nil {
   125  		return data, err
   126  	}
   127  
   128  	if len(data) == 0 {
   129  		return data, nil
   130  	}
   131  
   132  	err = fifo.client.SRem(ctx, fifo.setName, data).Err()
   133  	return data, err
   134  }
   135  
   136  // Has returns whether the fifo contains this data
   137  func (fifo *RedisUniqueByteFIFO) Has(ctx context.Context, data []byte) (bool, error) {
   138  	return fifo.client.SIsMember(ctx, fifo.setName, data).Result()
   139  }
   140  
   141  func init() {
   142  	queuesMap[RedisUniqueQueueType] = NewRedisUniqueQueue
   143  }