github.com/gitbundle/modules@v0.0.0-20231025071548-85b91c5c3b01/queue/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  // LevelQueueType is the type for level queue
    17  const LevelQueueType Type = "level"
    18  
    19  // LevelQueueConfiguration is the configuration for a LevelQueue
    20  type LevelQueueConfiguration struct {
    21  	ByteFIFOQueueConfiguration
    22  	DataDir          string
    23  	ConnectionString string
    24  	QueueName        string
    25  }
    26  
    27  // LevelQueue implements a disk library queue
    28  type LevelQueue struct {
    29  	*ByteFIFOQueue
    30  }
    31  
    32  // NewLevelQueue creates a ledis local queue
    33  func NewLevelQueue(handle HandlerFunc, cfg, exemplar interface{}) (Queue, error) {
    34  	configInterface, err := toConfig(LevelQueueConfiguration{}, cfg)
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  	config := configInterface.(LevelQueueConfiguration)
    39  
    40  	if len(config.ConnectionString) == 0 {
    41  		config.ConnectionString = config.DataDir
    42  	}
    43  	config.WaitOnEmpty = true
    44  
    45  	byteFIFO, err := NewLevelQueueByteFIFO(config.ConnectionString, config.QueueName)
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  
    50  	byteFIFOQueue, err := NewByteFIFOQueue(LevelQueueType, byteFIFO, handle, config.ByteFIFOQueueConfiguration, exemplar)
    51  	if err != nil {
    52  		return nil, err
    53  	}
    54  
    55  	queue := &LevelQueue{
    56  		ByteFIFOQueue: byteFIFOQueue,
    57  	}
    58  	queue.qid = GetManager().Add(queue, LevelQueueType, config, exemplar)
    59  	return queue, nil
    60  }
    61  
    62  var _ ByteFIFO = &LevelQueueByteFIFO{}
    63  
    64  // LevelQueueByteFIFO represents a ByteFIFO formed from a LevelQueue
    65  type LevelQueueByteFIFO struct {
    66  	internal   *levelqueue.Queue
    67  	connection string
    68  }
    69  
    70  // NewLevelQueueByteFIFO creates a ByteFIFO formed from a LevelQueue
    71  func NewLevelQueueByteFIFO(connection, prefix string) (*LevelQueueByteFIFO, error) {
    72  	db, err := nosql.GetManager().GetLevelDB(connection)
    73  	if err != nil {
    74  		return nil, err
    75  	}
    76  
    77  	internal, err := levelqueue.NewQueue(db, []byte(prefix), false)
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  
    82  	return &LevelQueueByteFIFO{
    83  		connection: connection,
    84  		internal:   internal,
    85  	}, nil
    86  }
    87  
    88  // PushFunc will push data into the fifo
    89  func (fifo *LevelQueueByteFIFO) PushFunc(ctx context.Context, data []byte, fn func() error) error {
    90  	if fn != nil {
    91  		if err := fn(); err != nil {
    92  			return err
    93  		}
    94  	}
    95  	return fifo.internal.LPush(data)
    96  }
    97  
    98  // PushBack pushes data to the top of the fifo
    99  func (fifo *LevelQueueByteFIFO) PushBack(ctx context.Context, data []byte) error {
   100  	return fifo.internal.RPush(data)
   101  }
   102  
   103  // Pop pops data from the start of the fifo
   104  func (fifo *LevelQueueByteFIFO) Pop(ctx context.Context) ([]byte, error) {
   105  	data, err := fifo.internal.RPop()
   106  	if err != nil && err != levelqueue.ErrNotFound {
   107  		return nil, err
   108  	}
   109  	return data, nil
   110  }
   111  
   112  // Close this fifo
   113  func (fifo *LevelQueueByteFIFO) Close() error {
   114  	err := fifo.internal.Close()
   115  	_ = nosql.GetManager().CloseLevelDB(fifo.connection)
   116  	return err
   117  }
   118  
   119  // Len returns the length of the fifo
   120  func (fifo *LevelQueueByteFIFO) Len(ctx context.Context) int64 {
   121  	return fifo.internal.Len()
   122  }
   123  
   124  func init() {
   125  	queuesMap[LevelQueueType] = NewLevelQueue
   126  }