watchtower/wtclient: generic disk overflow queue
In this commit, a new generic DiskOverflowQueue implementation is added. This allows a user to specify a maximum number of items that the queue can hold in-memory. Any new items will then overflow to disk. The producer and consumer of the queue items will interact with the queue just like a normal in-memory queue.
This commit is contained in:
3 files changed
+1005
No files matched your search
@@ -0,0 +1,566 @@
|
||||
package wtclient
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/btcsuite/btclog"
|
||||
"github.com/lightningnetwork/lnd/watchtower/wtdb"
|
||||
)
|
||||
|
||||
const (
|
||||
// dbErrorBackoff is the length of time we will back off before retrying
|
||||
// any DB action that failed.
|
||||
dbErrorBackoff = time.Second * 5
|
||||
)
|
||||
|
||||
// internalTask wraps a BackupID task with a success channel.
|
||||
type internalTask[T any] struct {
|
||||
task T
|
||||
success chan bool
|
||||
}
|
||||
|
||||
// newInternalTask creates a new internalTask with the given task.
|
||||
func newInternalTask[T any](task T) *internalTask[T] {
|
||||
return &internalTask[T]{
|
||||
task: task,
|
||||
success: make(chan bool),
|
||||
}
|
||||
}
|
||||
|
||||
// DiskOverflowQueue is a queue that must be initialised with a certain maximum
|
||||
// buffer size which represents the maximum number of elements that the queue
|
||||
// should hold in memory. If the queue is full, then any new elements added to
|
||||
// the queue will be persisted to disk instead. Once a consumer starts reading
|
||||
// from the front of the queue again then items on disk will be moved into the
|
||||
// queue again. The queue is also re-start safe. When it is stopped, any items
|
||||
// in the memory queue, will be persisted to disk. On start up, the queue will
|
||||
// be re-initialised with the items on disk.
|
||||
type DiskOverflowQueue[T any] struct {
|
||||
startOnce sync.Once
|
||||
stopOnce sync.Once
|
||||
|
||||
log btclog.Logger
|
||||
|
||||
// db is the database that will be used to persist queue items to disk.
|
||||
db wtdb.Queue[T]
|
||||
|
||||
// toDisk represents the current mode of operation of the queue.
|
||||
toDisk atomic.Bool
|
||||
|
||||
// We used an unbound list for the input of the queue so that producers
|
||||
// putting items into the queue are never blocked.
|
||||
inputListMu sync.Mutex
|
||||
inputListCond *sync.Cond
|
||||
inputList *list.List
|
||||
|
||||
// inputChan is an unbuffered channel used to pass items from
|
||||
// drainInputList to feedMemQueue.
|
||||
inputChan chan *internalTask[T]
|
||||
|
||||
// memQueue is a buffered channel used to pass items from
|
||||
// feedMemQueue to feedOutputChan.
|
||||
memQueue chan T
|
||||
|
||||
// outputChan is an unbuffered channel from which items at the head of
|
||||
// the queue can be read.
|
||||
outputChan chan T
|
||||
|
||||
// newDiskItemSignal is used to signal that there is a new item in the
|
||||
// main disk queue. There should only be one reader and one writer for
|
||||
// this channel.
|
||||
newDiskItemSignal chan struct{}
|
||||
|
||||
// leftOverItem1 will be a non-nil task on shutdown if the
|
||||
// feedOutputChan method was holding an unhandled tasks at shutdown
|
||||
// time. Since feedOutputChan handles the very head of the queue, this
|
||||
// item should be the first to be reloaded on restart.
|
||||
leftOverItem1 *T
|
||||
|
||||
// leftOverItems2 will be non-empty on shutdown if the feedMemQueue
|
||||
// method was holding any unhandled tasks at shutdown time. Since
|
||||
// feedMemQueue manages the input to the queue, the tasks should be
|
||||
// pushed to the head of the disk queue.
|
||||
leftOverItems2 []T
|
||||
|
||||
// leftOverItem3 will be non-nil on shutdown if drainInputList was
|
||||
// holding an unhandled task at shutdown time. This task should be put
|
||||
// at the tail of the disk queue but should come before any input list
|
||||
// task.
|
||||
leftOverItem3 *T
|
||||
|
||||
quit chan struct{}
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewDiskOverflowQueue constructs a new DiskOverflowQueue.
|
||||
func NewDiskOverflowQueue[T any](db wtdb.Queue[T], maxQueueSize uint64,
|
||||
logger btclog.Logger) (*DiskOverflowQueue[T], error) {
|
||||
|
||||
if maxQueueSize < 2 {
|
||||
return nil, errors.New("the in-memory queue buffer size " +
|
||||
"must be larger than 2")
|
||||
}
|
||||
|
||||
q := &DiskOverflowQueue[T]{
|
||||
log: logger,
|
||||
db: db,
|
||||
inputList: list.New(),
|
||||
newDiskItemSignal: make(chan struct{}, 1),
|
||||
inputChan: make(chan *internalTask[T]),
|
||||
memQueue: make(chan T, maxQueueSize-2),
|
||||
outputChan: make(chan T),
|
||||
quit: make(chan struct{}),
|
||||
}
|
||||
q.inputListCond = sync.NewCond(&q.inputListMu)
|
||||
|
||||
return q, nil
|
||||
}
|
||||
|
||||
// Start kicks off all the goroutines that are required to manage the queue.
|
||||
func (q *DiskOverflowQueue[T]) Start() error {
|
||||
var err error
|
||||
q.startOnce.Do(func() {
|
||||
err = q.start()
|
||||
})
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// start kicks off all the goroutines that are required to manage the queue.
|
||||
func (q *DiskOverflowQueue[T]) start() error {
|
||||
numDisk, err := q.db.Len()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if numDisk != 0 {
|
||||
q.toDisk.Store(true)
|
||||
}
|
||||
|
||||
// Kick off the three goroutines which will handle the input list, the
|
||||
// in-memory queue and the output channel.
|
||||
// The three goroutines are moving items according to the following
|
||||
// diagram:
|
||||
//
|
||||
// ┌─────────┐ drainInputList ┌──────────┐
|
||||
// │inputList├─────┬──────────►│disk/db │
|
||||
// └─────────┘ │ └──────────┘
|
||||
// │ (depending on mode)
|
||||
// │ ┌──────────┐
|
||||
// └──────────►│inputChan │
|
||||
// └──────────┘
|
||||
//
|
||||
// ┌─────────┐ feedMemQueue ┌──────────┐
|
||||
// │disk/db ├───────┬────────►│memQueue │
|
||||
// └─────────┘ │ └──────────┘
|
||||
// │ (depending on mode)
|
||||
// ┌─────────┐ │
|
||||
// │inputChan├───────┘
|
||||
// └─────────┘
|
||||
//
|
||||
// ┌─────────┐ feedOutputChan ┌──────────┐
|
||||
// │memQueue ├────────────────►│outputChan│
|
||||
// └─────────┘ └──────────┘
|
||||
//
|
||||
q.wg.Add(3)
|
||||
go q.drainInputList()
|
||||
go q.feedMemQueue()
|
||||
go q.feedOutputChan()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop stops the queue and persists any items in the memory queue to disk.
|
||||
func (q *DiskOverflowQueue[T]) Stop() error {
|
||||
var err error
|
||||
q.stopOnce.Do(func() {
|
||||
err = q.stop()
|
||||
})
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// stop the queue and persists any items in the memory queue to disk.
|
||||
func (q *DiskOverflowQueue[T]) stop() error {
|
||||
close(q.quit)
|
||||
|
||||
// Signal on the inputListCond until all the goroutines have returned.
|
||||
shutdown := make(chan struct{})
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case <-time.After(time.Millisecond):
|
||||
q.inputListCond.Signal()
|
||||
case <-shutdown:
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
q.wg.Wait()
|
||||
close(shutdown)
|
||||
|
||||
// queueHead will be the items that we will be pushed to the head of
|
||||
// the queue.
|
||||
var queueHead []T
|
||||
|
||||
// First, we append leftOverItem1 since this task is the current head
|
||||
// of the queue.
|
||||
if q.leftOverItem1 != nil {
|
||||
queueHead = append(queueHead, *q.leftOverItem1)
|
||||
}
|
||||
|
||||
// Next, drain the buffered queue.
|
||||
for {
|
||||
task, ok := <-q.memQueue
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
|
||||
queueHead = append(queueHead, task)
|
||||
}
|
||||
|
||||
// Then, any items held in leftOverItems2 would have been next to join
|
||||
// the memQueue. So those gets added next.
|
||||
if len(q.leftOverItems2) != 0 {
|
||||
queueHead = append(queueHead, q.leftOverItems2...)
|
||||
}
|
||||
|
||||
// Now, push these items to the head of the queue.
|
||||
err := q.db.PushHead(queueHead...)
|
||||
if err != nil {
|
||||
q.log.Errorf("Could not add tasks to queue head: %v", err)
|
||||
}
|
||||
|
||||
// Next we handle any items that need to be added to the main disk
|
||||
// queue.
|
||||
var diskQueue []T
|
||||
|
||||
// Any item in leftOverItem3 is the first item that should join the
|
||||
// disk queue.
|
||||
if q.leftOverItem3 != nil {
|
||||
diskQueue = append(diskQueue, *q.leftOverItem3)
|
||||
}
|
||||
|
||||
// Lastly, drain any items in the unbuffered input list.
|
||||
q.inputListCond.L.Lock()
|
||||
for q.inputList.Front() != nil {
|
||||
e := q.inputList.Front()
|
||||
|
||||
//nolint:forcetypeassert
|
||||
task := q.inputList.Remove(e).(T)
|
||||
|
||||
diskQueue = append(diskQueue, task)
|
||||
}
|
||||
q.inputListCond.L.Unlock()
|
||||
|
||||
// Now persist these items to the main disk queue.
|
||||
err = q.db.Push(diskQueue...)
|
||||
if err != nil {
|
||||
q.log.Errorf("Could not add tasks to queue tail: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// QueueBackupID adds a wtdb.BackupID to the queue. It will only return an error
|
||||
// if the queue has been stopped. It is non-blocking.
|
||||
func (q *DiskOverflowQueue[T]) QueueBackupID(item *wtdb.BackupID) error {
|
||||
// Return an error if the queue has been stopped
|
||||
select {
|
||||
case <-q.quit:
|
||||
return ErrClientExiting
|
||||
default:
|
||||
}
|
||||
|
||||
// Add the new item to the unbound input list.
|
||||
q.inputListCond.L.Lock()
|
||||
q.inputList.PushBack(item)
|
||||
q.inputListCond.L.Unlock()
|
||||
|
||||
// Signal that there is a new item in the input list.
|
||||
q.inputListCond.Signal()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// NextBackupID can be used to read from the head of the DiskOverflowQueue.
|
||||
func (q *DiskOverflowQueue[T]) NextBackupID() <-chan T {
|
||||
return q.outputChan
|
||||
}
|
||||
|
||||
// drainInputList handles the input to the DiskOverflowQueue. It takes from the
|
||||
// un-bounded input list and then, depending on what mode the queue is in,
|
||||
// either puts the new item straight onto the persisted disk queue or attempts
|
||||
// to feed it into the memQueue. On exit, any unhandled task will be assigned to
|
||||
// leftOverItem3.
|
||||
func (q *DiskOverflowQueue[T]) drainInputList() {
|
||||
defer q.wg.Done()
|
||||
|
||||
for {
|
||||
// Wait for the input list to not be empty.
|
||||
q.inputListCond.L.Lock()
|
||||
for q.inputList.Front() == nil {
|
||||
q.inputListCond.Wait()
|
||||
|
||||
select {
|
||||
case <-q.quit:
|
||||
q.inputListCond.L.Unlock()
|
||||
return
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// Pop the first element from the queue.
|
||||
e := q.inputList.Front()
|
||||
|
||||
//nolint:forcetypeassert
|
||||
task := q.inputList.Remove(e).(T)
|
||||
q.inputListCond.L.Unlock()
|
||||
|
||||
// What we do with this new item depends on what the mode of the
|
||||
// queue currently is.
|
||||
for q.pushToActiveQueue(task) {
|
||||
}
|
||||
|
||||
// If the above returned false because the quit channel was
|
||||
// closed, then we exit.
|
||||
select {
|
||||
case <-q.quit:
|
||||
return
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// pushToActiveQueue handles the input of a new task to the queue. It returns
|
||||
// true if the task should be retried and false if the task was handled or the
|
||||
// quit channel fired.
|
||||
func (q *DiskOverflowQueue[T]) pushToActiveQueue(task T) bool {
|
||||
// If the queue is in disk mode then any new items should be put
|
||||
// straight into the disk queue.
|
||||
if q.toDisk.Load() {
|
||||
err := q.db.Push(task)
|
||||
if err != nil {
|
||||
// Log and back off for a few seconds and then
|
||||
// try again with the same task.
|
||||
q.log.Errorf("could not persist %s to disk. "+
|
||||
"Retrying after backoff", task)
|
||||
|
||||
select {
|
||||
// Backoff for a bit and then re-check the mode
|
||||
// and try again to handle the task.
|
||||
case <-time.After(dbErrorBackoff):
|
||||
return true
|
||||
|
||||
// If the queue is quit at this moment, then the
|
||||
// unhandled task is assigned to leftOverItem3
|
||||
// so that it can be handled by the stop method.
|
||||
case <-q.quit:
|
||||
q.leftOverItem3 = &task
|
||||
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Send a signal that there is a new item in the main
|
||||
// disk queue.
|
||||
select {
|
||||
case q.newDiskItemSignal <- struct{}{}:
|
||||
case <-q.quit:
|
||||
|
||||
// Because there might already be a signal in the
|
||||
// newDiskItemSignal channel, we can skip sending another
|
||||
// signal. The channel only has a buffer of one, so we would
|
||||
// block here if we didn't have a default case.
|
||||
default:
|
||||
}
|
||||
|
||||
// If we got here, we were able to store the task in the disk
|
||||
// queue, so we can return false as no retry is necessary.
|
||||
return false
|
||||
}
|
||||
|
||||
// If the mode is memory mode, then try feed it to the feedMemQueue
|
||||
// handler via the un-buffered inputChan channel. We wrap it in an
|
||||
// internal task so that we can find out if feedMemQueue successfully
|
||||
// handled the item. If it did, we continue in memory mode and if not,
|
||||
// then we switch to disk mode so that we can persist the item to the
|
||||
// disk queue instead.
|
||||
it := newInternalTask(task)
|
||||
|
||||
select {
|
||||
// Try feed the task to the feedMemQueue handler. The handler, if it
|
||||
// does take the task, is guaranteed to respond via the success channel
|
||||
// of the task to indicate if the task was successfully added to the
|
||||
// in-mem queue. This is guaranteed even if the queue is being stopped.
|
||||
case q.inputChan <- it:
|
||||
|
||||
// If the queue is quit at this moment, then the unhandled task is
|
||||
// assigned to leftOverItem3 so that it can be handled by the stop
|
||||
// method.
|
||||
case <-q.quit:
|
||||
q.leftOverItem3 = &task
|
||||
|
||||
return false
|
||||
|
||||
default:
|
||||
// The task was not accepted. So maybe the mode changed.
|
||||
return true
|
||||
}
|
||||
|
||||
// If we get here, it means that the feedMemQueue handler took the task.
|
||||
// It is guaranteed to respond via the success channel, so we wait for
|
||||
// that response here.
|
||||
s := <-it.success
|
||||
if s {
|
||||
return false
|
||||
}
|
||||
|
||||
// If the task was not successfully handled by feedMemQueue, then we
|
||||
// switch to disk mode so that the task can be persisted in the disk
|
||||
// queue instead.
|
||||
q.toDisk.Store(true)
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// feedMemQueue manages which items should be fed onto the buffered
|
||||
// memQueue. If the queue is then in disk mode, then the handler will read new
|
||||
// tasks from the disk queue until it is empty. After that, it will switch
|
||||
// between reading from the input channel or the disk queue depending on the
|
||||
// queue mode.
|
||||
func (q *DiskOverflowQueue[T]) feedMemQueue() {
|
||||
defer func() {
|
||||
close(q.memQueue)
|
||||
q.wg.Done()
|
||||
}()
|
||||
|
||||
feedFromDisk := func() {
|
||||
select {
|
||||
case <-q.quit:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
for {
|
||||
// Ideally, we want to do batch reads from the DB. So
|
||||
// we check how much capacity there is in the memQueue
|
||||
// and fetch enough tasks to fill that capacity. If
|
||||
// there is no capacity, however, then we at least want
|
||||
// to fetch one task.
|
||||
numToPop := cap(q.memQueue) - len(q.memQueue)
|
||||
if numToPop == 0 {
|
||||
numToPop = 1
|
||||
}
|
||||
|
||||
tasks, err := q.db.PopUpTo(numToPop)
|
||||
if errors.Is(err, wtdb.ErrEmptyQueue) {
|
||||
q.toDisk.Store(false)
|
||||
|
||||
return
|
||||
} else if err != nil {
|
||||
q.log.Errorf("Could not load next task from " +
|
||||
"disk. Retrying.")
|
||||
|
||||
select {
|
||||
case <-time.After(dbErrorBackoff):
|
||||
continue
|
||||
case <-q.quit:
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
for i, task := range tasks {
|
||||
select {
|
||||
case q.memQueue <- task:
|
||||
|
||||
// If the queue is quit at this moment, then the
|
||||
// unhandled tasks are assigned to
|
||||
// leftOverItems2 so that they can be handled
|
||||
// by the stop method.
|
||||
case <-q.quit:
|
||||
q.leftOverItems2 = tasks[i:]
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If the queue is in disk mode, then the memQueue is fed with tasks
|
||||
// from the disk queue until it is empty.
|
||||
if q.toDisk.Load() {
|
||||
feedFromDisk()
|
||||
}
|
||||
|
||||
// Now the queue enters its normal operation.
|
||||
for {
|
||||
select {
|
||||
case <-q.quit:
|
||||
return
|
||||
|
||||
// If there is a signal that a new item has been added to disk
|
||||
// then we use the disk queue as the source of the next task
|
||||
// to feed into memQueue.
|
||||
case <-q.newDiskItemSignal:
|
||||
feedFromDisk()
|
||||
|
||||
// If any items come through on the inputChan, then we try feed
|
||||
// these directly into the memQueue. If there is space in the
|
||||
// memeQueue then we respond with success to the producer,
|
||||
// otherwise we respond with failure so that the producer can
|
||||
// instead persist the task to disk. After the producer,
|
||||
// drainInputList, has pushed an item to inputChan, it is
|
||||
// guaranteed to await a response on the task's success channel
|
||||
// before quiting. Therefore, it is not required to listen on
|
||||
// the quit channel here.
|
||||
case task := <-q.inputChan:
|
||||
select {
|
||||
case q.memQueue <- task.task:
|
||||
task.success <- true
|
||||
continue
|
||||
default:
|
||||
task.success <- false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// feedOutputChan will pop an item from the buffered memQueue and block until
|
||||
// the item is taken from the un-buffered outputChan. This is done repeatedly
|
||||
// for the lifetime of the DiskOverflowQueue. On shutdown of the queue, any
|
||||
// item not consumed by the outputChan but held by this method is assigned to
|
||||
// the leftOverItem1 member so that the Stop method can persist the item to
|
||||
// disk so that it is reloaded on restart.
|
||||
//
|
||||
// NOTE: This must be run as a goroutine.
|
||||
func (q *DiskOverflowQueue[T]) feedOutputChan() {
|
||||
defer func() {
|
||||
close(q.outputChan)
|
||||
q.wg.Done()
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case nextTask, ok := <-q.memQueue:
|
||||
// If the memQueue is closed, then the queue is
|
||||
// stopping.
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case q.outputChan <- nextTask:
|
||||
case <-q.quit:
|
||||
q.leftOverItem1 = &nextTask
|
||||
return
|
||||
}
|
||||
|
||||
case <-q.quit:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user