github.com/gitbundle/modules@v0.0.0-20231025071548-85b91c5c3b01/queue/unique_queue_disk.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/gitbundle/modules/nosql"
    12  
    13  	"gitea.com/lunny/levelqueue"
    14  )
    15  
    16  // LevelUniqueQueueType is the type for level queue
    17  const LevelUniqueQueueType Type = "unique-level"
    18  
    19  // LevelUniqueQueueConfiguration is the configuration for a LevelUniqueQueue
    20  type LevelUniqueQueueConfiguration struct {
    21  	ByteFIFOQueueConfiguration
    22  	DataDir          string
    23  	ConnectionString string
    24  	QueueName        string
    25  }
    26  
    27  // LevelUniqueQueue implements a disk library queue
    28  type LevelUniqueQueue struct {
    29  	*ByteFIFOUniqueQueue
    30  }
    31  
    32  // NewLevelUniqueQueue creates a ledis local queue
    33  //
    34  // Please note that this Queue does not guarantee that a particular
    35  // task cannot be processed twice or more at the same time. Uniqueness is
    36  // only guaranteed whilst the task is waiting in the queue.
    37  func NewLevelUniqueQueue(handle HandlerFunc, cfg, exemplar interface{}) (Queue, error) {
    38  	configInterface, err := toConfig(LevelUniqueQueueConfiguration{}, cfg)
    39  	if err != nil {
    40  		return nil, err
    41  	}
    42  	config := configInterface.(LevelUniqueQueueConfiguration)
    43  
    44  	if len(config.ConnectionString) == 0 {
    45  		config.ConnectionString = config.DataDir
    46  	}
    47  	config.WaitOnEmpty = true
    48  
    49  	byteFIFO, err := NewLevelUniqueQueueByteFIFO(config.ConnectionString, config.QueueName)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  
    54  	byteFIFOQueue, err := NewByteFIFOUniqueQueue(LevelUniqueQueueType, byteFIFO, handle, config.ByteFIFOQueueConfiguration, exemplar)
    55  	if err != nil {
    56  		return nil, err
    57  	}
    58  
    59  	queue := &LevelUniqueQueue{
    60  		ByteFIFOUniqueQueue: byteFIFOQueue,
    61  	}
    62  	queue.qid = GetManager().Add(queue, LevelUniqueQueueType, config, exemplar)
    63  	return queue, nil
    64  }
    65  
    66  var _ UniqueByteFIFO = &LevelUniqueQueueByteFIFO{}
    67  
    68  // LevelUniqueQueueByteFIFO represents a ByteFIFO formed from a LevelUniqueQueue
    69  type LevelUniqueQueueByteFIFO struct {
    70  	internal   *levelqueue.UniqueQueue
    71  	connection string
    72  }
    73  
    74  // NewLevelUniqueQueueByteFIFO creates a new ByteFIFO formed from a LevelUniqueQueue
    75  func NewLevelUniqueQueueByteFIFO(connection, prefix string) (*LevelUniqueQueueByteFIFO, error) {
    76  	db, err := nosql.GetManager().GetLevelDB(connection)
    77  	if err != nil {
    78  		return nil, err
    79  	}
    80  
    81  	internal, err := levelqueue.NewUniqueQueue(db, []byte(prefix), []byte(prefix+"-unique"), false)
    82  	if err != nil {
    83  		return nil, err
    84  	}
    85  
    86  	return &LevelUniqueQueueByteFIFO{
    87  		connection: connection,
    88  		internal:   internal,
    89  	}, nil
    90  }
    91  
    92  // PushFunc pushes data to the end of the fifo and calls the callback if it is added
    93  func (fifo *LevelUniqueQueueByteFIFO) PushFunc(ctx context.Context, data []byte, fn func() error) error {
    94  	return fifo.internal.LPushFunc(data, fn)
    95  }
    96  
    97  // PushBack pushes data to the top of the fifo
    98  func (fifo *LevelUniqueQueueByteFIFO) PushBack(ctx context.Context, data []byte) error {
    99  	return fifo.internal.RPush(data)
   100  }
   101  
   102  // Pop pops data from the start of the fifo
   103  func (fifo *LevelUniqueQueueByteFIFO) Pop(ctx context.Context) ([]byte, error) {
   104  	data, err := fifo.internal.RPop()
   105  	if err != nil && err != levelqueue.ErrNotFound {
   106  		return nil, err
   107  	}
   108  	return data, nil
   109  }
   110  
   111  // Len returns the length of the fifo
   112  func (fifo *LevelUniqueQueueByteFIFO) Len(ctx context.Context) int64 {
   113  	return fifo.internal.Len()
   114  }
   115  
   116  // Has returns whether the fifo contains this data
   117  func (fifo *LevelUniqueQueueByteFIFO) Has(ctx context.Context, data []byte) (bool, error) {
   118  	return fifo.internal.Has(data)
   119  }
   120  
   121  // Close this fifo
   122  func (fifo *LevelUniqueQueueByteFIFO) Close() error {
   123  	err := fifo.internal.Close()
   124  	_ = nosql.GetManager().CloseLevelDB(fifo.connection)
   125  	return err
   126  }
   127  
   128  func init() {
   129  	queuesMap[LevelUniqueQueueType] = NewLevelUniqueQueue
   130  }