github.com/gitbundle/modules@v0.0.0-20231025071548-85b91c5c3b01/queue/unique_queue_disk_channel.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  	"runtime/pprof"
    11  	"sync"
    12  	"time"
    13  
    14  	"github.com/gitbundle/modules/log"
    15  )
    16  
    17  // PersistableChannelUniqueQueueType is the type for persistable queue
    18  const PersistableChannelUniqueQueueType Type = "unique-persistable-channel"
    19  
    20  // PersistableChannelUniqueQueueConfiguration is the configuration for a PersistableChannelUniqueQueue
    21  type PersistableChannelUniqueQueueConfiguration struct {
    22  	Name         string
    23  	DataDir      string
    24  	BatchLength  int
    25  	QueueLength  int
    26  	Timeout      time.Duration
    27  	MaxAttempts  int
    28  	Workers      int
    29  	MaxWorkers   int
    30  	BlockTimeout time.Duration
    31  	BoostTimeout time.Duration
    32  	BoostWorkers int
    33  }
    34  
    35  // PersistableChannelUniqueQueue wraps a channel queue and level queue together
    36  //
    37  // Please note that this Queue does not guarantee that a particular
    38  // task cannot be processed twice or more at the same time. Uniqueness is
    39  // only guaranteed whilst the task is waiting in the queue.
    40  type PersistableChannelUniqueQueue struct {
    41  	channelQueue *ChannelUniqueQueue
    42  	delayedStarter
    43  	lock   sync.Mutex
    44  	closed chan struct{}
    45  }
    46  
    47  // NewPersistableChannelUniqueQueue creates a wrapped batched channel queue with persistable level queue backend when shutting down
    48  // This differs from a wrapped queue in that the persistent queue is only used to persist at shutdown/terminate
    49  func NewPersistableChannelUniqueQueue(handle HandlerFunc, cfg, exemplar interface{}) (Queue, error) {
    50  	configInterface, err := toConfig(PersistableChannelUniqueQueueConfiguration{}, cfg)
    51  	if err != nil {
    52  		return nil, err
    53  	}
    54  	config := configInterface.(PersistableChannelUniqueQueueConfiguration)
    55  
    56  	queue := &PersistableChannelUniqueQueue{
    57  		closed: make(chan struct{}),
    58  	}
    59  
    60  	wrappedHandle := func(data ...Data) (failed []Data) {
    61  		for _, unhandled := range handle(data...) {
    62  			if fail := queue.PushBack(unhandled); fail != nil {
    63  				failed = append(failed, fail)
    64  			}
    65  		}
    66  		return
    67  	}
    68  
    69  	channelUniqueQueue, err := NewChannelUniqueQueue(wrappedHandle, ChannelUniqueQueueConfiguration{
    70  		WorkerPoolConfiguration: WorkerPoolConfiguration{
    71  			QueueLength:  config.QueueLength,
    72  			BatchLength:  config.BatchLength,
    73  			BlockTimeout: config.BlockTimeout,
    74  			BoostTimeout: config.BoostTimeout,
    75  			BoostWorkers: config.BoostWorkers,
    76  			MaxWorkers:   config.MaxWorkers,
    77  			Name:         config.Name + "-channel",
    78  		},
    79  		Workers: config.Workers,
    80  	}, exemplar)
    81  	if err != nil {
    82  		return nil, err
    83  	}
    84  
    85  	// the level backend only needs temporary workers to catch up with the previously dropped work
    86  	levelCfg := LevelUniqueQueueConfiguration{
    87  		ByteFIFOQueueConfiguration: ByteFIFOQueueConfiguration{
    88  			WorkerPoolConfiguration: WorkerPoolConfiguration{
    89  				QueueLength:  config.QueueLength,
    90  				BatchLength:  config.BatchLength,
    91  				BlockTimeout: 1 * time.Second,
    92  				BoostTimeout: 5 * time.Minute,
    93  				BoostWorkers: 1,
    94  				MaxWorkers:   5,
    95  				Name:         config.Name + "-level",
    96  			},
    97  			Workers: 0,
    98  		},
    99  		DataDir: config.DataDir,
   100  	}
   101  
   102  	queue.channelQueue = channelUniqueQueue.(*ChannelUniqueQueue)
   103  
   104  	levelQueue, err := NewLevelUniqueQueue(func(data ...Data) []Data {
   105  		for _, datum := range data {
   106  			err := queue.Push(datum)
   107  			if err != nil && err != ErrAlreadyInQueue {
   108  				log.Error("Unable push to channelled queue: %v", err)
   109  			}
   110  		}
   111  		return nil
   112  	}, levelCfg, exemplar)
   113  	if err == nil {
   114  		queue.delayedStarter = delayedStarter{
   115  			internal: levelQueue.(*LevelUniqueQueue),
   116  			name:     config.Name,
   117  		}
   118  
   119  		_ = GetManager().Add(queue, PersistableChannelUniqueQueueType, config, exemplar)
   120  		return queue, nil
   121  	}
   122  	if IsErrInvalidConfiguration(err) {
   123  		// Retrying ain't gonna make this any better...
   124  		return nil, ErrInvalidConfiguration{cfg: cfg}
   125  	}
   126  
   127  	queue.delayedStarter = delayedStarter{
   128  		cfg:         levelCfg,
   129  		underlying:  LevelUniqueQueueType,
   130  		timeout:     config.Timeout,
   131  		maxAttempts: config.MaxAttempts,
   132  		name:        config.Name,
   133  	}
   134  	_ = GetManager().Add(queue, PersistableChannelUniqueQueueType, config, exemplar)
   135  	return queue, nil
   136  }
   137  
   138  // Name returns the name of this queue
   139  func (q *PersistableChannelUniqueQueue) Name() string {
   140  	return q.delayedStarter.name
   141  }
   142  
   143  // Push will push the indexer data to queue
   144  func (q *PersistableChannelUniqueQueue) Push(data Data) error {
   145  	return q.PushFunc(data, nil)
   146  }
   147  
   148  // PushFunc will push the indexer data to queue
   149  func (q *PersistableChannelUniqueQueue) PushFunc(data Data, fn func() error) error {
   150  	select {
   151  	case <-q.closed:
   152  		return q.internal.(UniqueQueue).PushFunc(data, fn)
   153  	default:
   154  		return q.channelQueue.PushFunc(data, fn)
   155  	}
   156  }
   157  
   158  // PushBack will push the indexer data to queue
   159  func (q *PersistableChannelUniqueQueue) PushBack(data Data) error {
   160  	select {
   161  	case <-q.closed:
   162  		if pbr, ok := q.internal.(PushBackable); ok {
   163  			return pbr.PushBack(data)
   164  		}
   165  		return q.internal.Push(data)
   166  	default:
   167  		return q.channelQueue.Push(data)
   168  	}
   169  }
   170  
   171  // Has will test if the queue has the data
   172  func (q *PersistableChannelUniqueQueue) Has(data Data) (bool, error) {
   173  	// This is more difficult...
   174  	has, err := q.channelQueue.Has(data)
   175  	if err != nil || has {
   176  		return has, err
   177  	}
   178  	q.lock.Lock()
   179  	defer q.lock.Unlock()
   180  	if q.internal == nil {
   181  		return false, nil
   182  	}
   183  	return q.internal.(UniqueQueue).Has(data)
   184  }
   185  
   186  // Run starts to run the queue
   187  func (q *PersistableChannelUniqueQueue) Run(atShutdown, atTerminate func(func())) {
   188  	pprof.SetGoroutineLabels(q.channelQueue.baseCtx)
   189  	log.Debug("PersistableChannelUniqueQueue: %s Starting", q.delayedStarter.name)
   190  
   191  	q.lock.Lock()
   192  	if q.internal == nil {
   193  		err := q.setInternal(atShutdown, func(data ...Data) []Data {
   194  			for _, datum := range data {
   195  				err := q.Push(datum)
   196  				if err != nil && err != ErrAlreadyInQueue {
   197  					log.Error("Unable push to channelled queue: %v", err)
   198  				}
   199  			}
   200  			return nil
   201  		}, q.channelQueue.exemplar)
   202  		q.lock.Unlock()
   203  		if err != nil {
   204  			log.Fatal("Unable to create internal queue for %s Error: %v", q.Name(), err)
   205  			return
   206  		}
   207  	} else {
   208  		q.lock.Unlock()
   209  	}
   210  	atShutdown(q.Shutdown)
   211  	atTerminate(q.Terminate)
   212  	_ = q.channelQueue.AddWorkers(q.channelQueue.workers, 0)
   213  
   214  	if luq, ok := q.internal.(*LevelUniqueQueue); ok && luq.ByteFIFOUniqueQueue.byteFIFO.Len(luq.shutdownCtx) != 0 {
   215  		// Just run the level queue - we shut it down once it's flushed
   216  		go q.internal.Run(func(_ func()) {}, func(_ func()) {})
   217  		go func() {
   218  			_ = q.internal.Flush(0)
   219  			log.Debug("LevelUniqueQueue: %s flushed so shutting down", q.internal.(*LevelUniqueQueue).Name())
   220  			q.internal.(*LevelUniqueQueue).Shutdown()
   221  			GetManager().Remove(q.internal.(*LevelUniqueQueue).qid)
   222  		}()
   223  	} else {
   224  		log.Debug("PersistableChannelUniqueQueue: %s Skipping running the empty level queue", q.delayedStarter.name)
   225  		q.internal.(*LevelUniqueQueue).Shutdown()
   226  		GetManager().Remove(q.internal.(*LevelUniqueQueue).qid)
   227  	}
   228  }
   229  
   230  // Flush flushes the queue
   231  func (q *PersistableChannelUniqueQueue) Flush(timeout time.Duration) error {
   232  	return q.channelQueue.Flush(timeout)
   233  }
   234  
   235  // FlushWithContext flushes the queue
   236  func (q *PersistableChannelUniqueQueue) FlushWithContext(ctx context.Context) error {
   237  	return q.channelQueue.FlushWithContext(ctx)
   238  }
   239  
   240  // IsEmpty checks if a queue is empty
   241  func (q *PersistableChannelUniqueQueue) IsEmpty() bool {
   242  	return q.channelQueue.IsEmpty()
   243  }
   244  
   245  // IsPaused will return if the pool or queue is paused
   246  func (q *PersistableChannelUniqueQueue) IsPaused() bool {
   247  	return q.channelQueue.IsPaused()
   248  }
   249  
   250  // Pause will pause the pool or queue
   251  func (q *PersistableChannelUniqueQueue) Pause() {
   252  	q.channelQueue.Pause()
   253  }
   254  
   255  // Resume will resume the pool or queue
   256  func (q *PersistableChannelUniqueQueue) Resume() {
   257  	q.channelQueue.Resume()
   258  }
   259  
   260  // IsPausedIsResumed will return a bool indicating if the pool or queue is paused and a channel that will be closed when it is resumed
   261  func (q *PersistableChannelUniqueQueue) IsPausedIsResumed() (paused, resumed <-chan struct{}) {
   262  	return q.channelQueue.IsPausedIsResumed()
   263  }
   264  
   265  // Shutdown processing this queue
   266  func (q *PersistableChannelUniqueQueue) Shutdown() {
   267  	log.Trace("PersistableChannelUniqueQueue: %s Shutting down", q.delayedStarter.name)
   268  	q.lock.Lock()
   269  	select {
   270  	case <-q.closed:
   271  		q.lock.Unlock()
   272  		return
   273  	default:
   274  		if q.internal != nil {
   275  			q.internal.(*LevelUniqueQueue).Shutdown()
   276  		}
   277  		close(q.closed)
   278  		q.lock.Unlock()
   279  	}
   280  
   281  	log.Trace("PersistableChannelUniqueQueue: %s Cancelling pools", q.delayedStarter.name)
   282  	q.internal.(*LevelUniqueQueue).baseCtxCancel()
   283  	q.channelQueue.baseCtxCancel()
   284  	log.Trace("PersistableChannelUniqueQueue: %s Waiting til done", q.delayedStarter.name)
   285  	q.channelQueue.Wait()
   286  	q.internal.(*LevelUniqueQueue).Wait()
   287  	// Redirect all remaining data in the chan to the internal channel
   288  	close(q.channelQueue.dataChan)
   289  	log.Trace("PersistableChannelUniqueQueue: %s Redirecting remaining data", q.delayedStarter.name)
   290  	for data := range q.channelQueue.dataChan {
   291  		_ = q.internal.Push(data)
   292  	}
   293  	log.Trace("PersistableChannelUniqueQueue: %s Done Redirecting remaining data", q.delayedStarter.name)
   294  
   295  	log.Debug("PersistableChannelUniqueQueue: %s Shutdown", q.delayedStarter.name)
   296  }
   297  
   298  // Terminate this queue and close the queue
   299  func (q *PersistableChannelUniqueQueue) Terminate() {
   300  	log.Trace("PersistableChannelUniqueQueue: %s Terminating", q.delayedStarter.name)
   301  	q.Shutdown()
   302  	q.lock.Lock()
   303  	defer q.lock.Unlock()
   304  	if q.internal != nil {
   305  		q.internal.(*LevelUniqueQueue).Terminate()
   306  	}
   307  	q.channelQueue.baseCtxFinished()
   308  	log.Debug("PersistableChannelUniqueQueue: %s Terminated", q.delayedStarter.name)
   309  }
   310  
   311  func init() {
   312  	queuesMap[PersistableChannelUniqueQueueType] = NewPersistableChannelUniqueQueue
   313  }