wtclient: replay pending tasks on sessionQueue stop

This commit does a few things:
- First, it gives the sessionQueue access to the TowerClient task
  pipeline so that it can replay backup tasks onto the pipeline on Stop.
- Given that the above is done, the ForceQuit functionality of the
  sessionQueue and TowerClient can be removed.
- The bug demonstrated in a prior commit is now fixed due to the above
  changes.

watchtower/wtclient/session_queue.go

@@ -16,17 +16,21 @@ import (
 	"github.com/lightningnetwork/lnd/watchtower/wtwire"
 )
 
-// reserveStatus is an enum that signals how full a particular session is.
-type reserveStatus uint8
+// sessionQueueStatus is an enum that signals how full a particular session is.
+type sessionQueueStatus uint8
 
 const (
-	// reserveAvailable indicates that the session has space for at least
-	// one more backup.
-	reserveAvailable reserveStatus = iota
+	// sessionQueueAvailable indicates that the session has space for at
+	// least one more backup.
+	sessionQueueAvailable sessionQueueStatus = iota
 
-	// reserveExhausted indicates that all slots in the session have been
-	// allocated.
-	reserveExhausted
+	// sessionQueueExhausted indicates that all slots in the session have
+	// been allocated.
+	sessionQueueExhausted
+
+	// sessionQueueShuttingDown indicates that the session queue is
+	// shutting down and so is no longer accepting any more backups.
+	sessionQueueShuttingDown
 )
 
 // sessionQueueConfig bundles the resources required by the sessionQueue to
@@ -62,6 +66,10 @@ type sessionQueueConfig struct {
 	// certain revoked commitment height.
 	BuildBreachRetribution BreachRetributionBuilder
 
+	// TaskPipeline is a pipeline which the sessionQueue should use to send
+	// any unhandled tasks on shutdown of the queue.
+	TaskPipeline *DiskOverflowQueue[*wtdb.BackupID]
+
 	// DB provides access to the client's stable storage.
 	DB DB
 
@@ -85,10 +93,8 @@ type sessionQueueConfig struct {
 
 // sessionQueue implements a reliable queue that will encrypt and send accepted
 // backups to the watchtower specified in the config's ClientSession. Calling
-// Quit will attempt to perform a clean shutdown by receiving an ACK from the
-// tower for all pending backups before exiting. The clean shutdown can be
-// aborted by using ForceQuit, which will attempt to shut down the queue
-// immediately.
+// Stop will attempt to perform a clean shutdown replaying any un-committed
+// pending updates to the TowerClient's main task pipeline.
 type sessionQueue struct {
 	started sync.Once
 	stopped sync.Once
@@ -109,9 +115,8 @@ type sessionQueue struct {
 
 	retryBackoff time.Duration
 
-	quit      chan struct{}
-	forceQuit chan struct{}
-	shutdown  chan struct{}
+	quit chan struct{}
+	wg   sync.WaitGroup
 }
 
 // newSessionQueue initializes a fresh sessionQueue.
@@ -133,8 +138,6 @@ func newSessionQueue(cfg *sessionQueueConfig,
 		seqNum:       cfg.ClientSession.SeqNum,
 		retryBackoff: cfg.MinBackoff,
 		quit:         make(chan struct{}),
-		forceQuit:    make(chan struct{}),
-		shutdown:     make(chan struct{}),
 	}
 	sq.queueCond = sync.NewCond(&sq.queueMtx)
 
@@ -151,41 +154,77 @@ func newSessionQueue(cfg *sessionQueueConfig,
 // backups.
 func (q *sessionQueue) Start() {
 	q.started.Do(func() {
+		q.wg.Add(1)
 		go q.sessionManager()
 	})
 }
 
 // Stop idempotently stops the sessionQueue by initiating a clean shutdown that
 // will clear all pending tasks in the queue before returning to the caller.
-func (q *sessionQueue) Stop() {
+func (q *sessionQueue) Stop() error {
+	var returnErr error
 	q.stopped.Do(func() {
 		q.log.Debugf("SessionQueue(%s) stopping ...", q.ID())
 
 		close(q.quit)
-		q.signalUntilShutdown()
 
-		// Skip log if we also force quit.
-		select {
-		case <-q.forceQuit:
+		shutdown := make(chan struct{})
+		go func() {
+			for {
+				select {
+				case <-time.After(time.Millisecond):
+					q.queueCond.Signal()
+				case <-shutdown:
+					return
+				}
+			}
+		}()
+
+		q.wg.Wait()
+		close(shutdown)
+
+		// Now, for any task in the pending queue that we have not yet
+		// created a CommittedUpdate for, re-add the task to the main
+		// task pipeline.
+		updates, err := q.cfg.DB.FetchSessionCommittedUpdates(q.ID())
+		if err != nil {
+			returnErr = err
 			return
-		default:
 		}
 
+		unAckedUpdates := make(map[wtdb.BackupID]bool)
+		for _, update := range updates {
+			unAckedUpdates[update.BackupID] = true
+		}
+
+		// Push any task that was on the pending queue that there is
+		// not yet a committed update for back to the main task
+		// pipeline.
+		q.queueCond.L.Lock()
+		for q.pendingQueue.Len() > 0 {
+			next := q.pendingQueue.Front()
+			q.pendingQueue.Remove(next)
+
+			//nolint:forcetypeassert
+			task := next.Value.(*backupTask)
+
+			if unAckedUpdates[task.id] {
+				continue
+			}
+
+			err := q.cfg.TaskPipeline.QueueBackupID(&task.id)
+			if err != nil {
+				log.Errorf("could not re-queue backup task: "+
+					"%v", err)
+				continue
+			}
+		}
+		q.queueCond.L.Unlock()
+
 		q.log.Debugf("SessionQueue(%s) stopped", q.ID())
 	})
-}
 
-// ForceQuit idempotently aborts any clean shutdown in progress and returns to
-// he caller after all lingering goroutines have spun down.
-func (q *sessionQueue) ForceQuit() {
-	q.forced.Do(func() {
-		q.log.Infof("SessionQueue(%s) force quitting...", q.ID())
-
-		close(q.forceQuit)
-		q.signalUntilShutdown()
-
-		q.log.Infof("SessionQueue(%s) force quit", q.ID())
-	})
+	return returnErr
 }
 
 // ID returns the wtdb.SessionID for the queue, which can be used to uniquely
@@ -196,10 +235,28 @@ func (q *sessionQueue) ID() *wtdb.SessionID {
 
 // AcceptTask attempts to queue a backupTask for delivery to the sessionQueue's
 // tower. The session will only be accepted if the queue is not already
-// exhausted and the task is successfully bound to the ClientSession.
-func (q *sessionQueue) AcceptTask(task *backupTask) (reserveStatus, bool) {
+// exhausted or shutting down and the task is successfully bound to the
+// ClientSession.
+func (q *sessionQueue) AcceptTask(task *backupTask) (sessionQueueStatus, bool) {
+	// Exit early if the queue has started shutting down.
+	select {
+	case <-q.quit:
+		return sessionQueueShuttingDown, false
+	default:
+	}
+
 	q.queueCond.L.Lock()
 
+	// There is a chance that sessionQueue started shutting down between
+	// the last quit channel check and waiting for the lock. So check one
+	// more time here.
+	select {
+	case <-q.quit:
+		q.queueCond.L.Unlock()
+		return sessionQueueShuttingDown, false
+	default:
+	}
+
 	numPending := uint32(q.pendingQueue.Len())
 	maxUpdates := q.cfg.ClientSession.Policy.MaxUpdates
 	q.log.Debugf("SessionQueue(%s) deciding to accept %v seqnum=%d "+
@@ -207,14 +264,14 @@ func (q *sessionQueue) AcceptTask(task *backupTask) (reserveStatus, bool) {
 		q.ID(), task.id, q.seqNum, numPending, maxUpdates)
 
 	// Examine the current reserve status of the session queue.
-	curStatus := q.reserveStatus()
+	curStatus := q.status()
 
 	switch curStatus {
 
 	// The session queue is exhausted, and cannot accept the task because it
 	// is full. Reject the task such that it can be tried against a
 	// different session.
-	case reserveExhausted:
+	case sessionQueueExhausted:
 		q.queueCond.L.Unlock()
 		return curStatus, false
 
@@ -224,7 +281,7 @@ func (q *sessionQueue) AcceptTask(task *backupTask) (reserveStatus, bool) {
 	// tried again.
 	//
 	// TODO(conner): queue backups and retry with different session params.
-	case reserveAvailable:
+	case sessionQueueAvailable:
 		err := task.bindSession(
 			&q.cfg.ClientSession.ClientSessionBody,
 			q.cfg.BuildBreachRetribution,
@@ -244,7 +301,7 @@ func (q *sessionQueue) AcceptTask(task *backupTask) (reserveStatus, bool) {
 	// Finally, compute the session's *new* reserve status. This will be
 	// used by the client to determine if it can continue using this session
 	// queue, or if it should negotiate a new one.
-	newStatus := q.reserveStatus()
+	newStatus := q.status()
 	q.queueCond.L.Unlock()
 
 	q.queueCond.Signal()
@@ -255,7 +312,7 @@ func (q *sessionQueue) AcceptTask(task *backupTask) (reserveStatus, bool) {
 // sessionManager is the primary event loop for the sessionQueue, and is
 // responsible for encrypting and sending accepted tasks to the tower.
 func (q *sessionQueue) sessionManager() {
-	defer close(q.shutdown)
+	defer q.wg.Done()
 
 	for {
 		q.queueCond.L.Lock()
@@ -266,12 +323,6 @@ func (q *sessionQueue) sessionManager() {
 
 			select {
 			case <-q.quit:
-				if q.commitQueue.Len() == 0 &&
-					q.pendingQueue.Len() == 0 {
-					q.queueCond.L.Unlock()
-					return
-				}
-			case <-q.forceQuit:
 				q.queueCond.L.Unlock()
 				return
 			default:
@@ -279,12 +330,9 @@ func (q *sessionQueue) sessionManager() {
 		}
 		q.queueCond.L.Unlock()
 
-		// Exit immediately if a force quit has been requested.  If
-		// either of the queues still has state updates to send to the
-		// tower, we may never exit in the above case if we are unable
-		// to reach the tower for some reason.
+		// Exit immediately if the sessionQueue has been stopped.
 		select {
-		case <-q.forceQuit:
+		case <-q.quit:
 			return
 		default:
 		}
@@ -333,7 +381,7 @@ func (q *sessionQueue) drainBackups() {
 			q.increaseBackoff()
 			select {
 			case <-time.After(q.retryBackoff):
-			case <-q.forceQuit:
+			case <-q.quit:
 			}
 			return
 		}
@@ -366,7 +414,7 @@ func (q *sessionQueue) drainBackups() {
 			q.increaseBackoff()
 			select {
 			case <-time.After(q.retryBackoff):
-			case <-q.forceQuit:
+			case <-q.quit:
 			}
 			return
 		}
@@ -388,7 +436,7 @@ func (q *sessionQueue) drainBackups() {
 		// when we will do so.
 		select {
 		case <-time.After(time.Millisecond):
-		case <-q.forceQuit:
+		case <-q.quit:
 			return
 		}
 	}
@@ -635,21 +683,21 @@ func (q *sessionQueue) sendStateUpdate(conn wtserver.Peer,
 	return nil
 }
 
-// reserveStatus returns a reserveStatus indicating whether the sessionQueue can
-// accept another task. reserveAvailable is returned when a task can be
-// accepted, and reserveExhausted is returned if the all slots in the session
-// have been allocated.
+// status returns a sessionQueueStatus indicating whether the sessionQueue can
+// accept another task. sessionQueueAvailable is returned when a task can be
+// accepted, and sessionQueueExhausted is returned if the all slots in the
+// session have been allocated.
 //
 // NOTE: This method MUST be called with queueCond's exclusive lock held.
-func (q *sessionQueue) reserveStatus() reserveStatus {
+func (q *sessionQueue) status() sessionQueueStatus {
 	numPending := uint32(q.pendingQueue.Len())
 	maxUpdates := uint32(q.cfg.ClientSession.Policy.MaxUpdates)
 
 	if uint32(q.seqNum)+numPending < maxUpdates {
-		return reserveAvailable
+		return sessionQueueAvailable
 	}
 
-	return reserveExhausted
+	return sessionQueueExhausted
 
 }
 
@@ -667,19 +715,6 @@ func (q *sessionQueue) increaseBackoff() {
 	}
 }
 
-// signalUntilShutdown strobes the sessionQueue's condition variable until the
-// main event loop exits.
-func (q *sessionQueue) signalUntilShutdown() {
-	for {
-		select {
-		case <-time.After(time.Millisecond):
-			q.queueCond.Signal()
-		case <-q.shutdown:
-			return
-		}
-	}
-}
-
 // sessionQueueSet maintains a mapping of SessionIDs to their corresponding
 // sessionQueue.
 type sessionQueueSet struct {
@@ -706,18 +741,18 @@ func (s *sessionQueueSet) AddAndStart(sessionQueue *sessionQueue) {
 
 // StopAndRemove stops the given session queue and removes it from the
 // sessionQueueSet.
-func (s *sessionQueueSet) StopAndRemove(id wtdb.SessionID) {
+func (s *sessionQueueSet) StopAndRemove(id wtdb.SessionID) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	queue, ok := s.queues[id]
 	if !ok {
-		return
+		return nil
 	}
 
-	queue.Stop()
-
 	delete(s.queues, id)
+
+	return queue.Stop()
 }
 
 // Get fetches and returns the sessionQueue with the given ID.