staticaddr: quit deposit fsm handler when finialized

Author: hieblmi

staticaddr/deposit/fsm.go

@@ -141,6 +141,15 @@ type FSM struct {
 
 	blockNtfnChan chan uint32
 
+	// quitChan stops after the FSM stops consuming blockNtfnChan.
+	quitChan chan struct{}
+
+	// spentChan is used to signal that the deposit has been spent and that
+	// the FSM should stop processing block notifications and exit.
+	spentChan chan struct{}
+
+	// finalizedDepositChan is used to signal that the deposit has been
+	// finalized and the FSM can be removed from the manager's memory.
 	finalizedDepositChan chan wire.OutPoint
 }
 
@@ -167,7 +176,9 @@ func NewFSM(ctx context.Context, deposit *Deposit, cfg *ManagerConfig,
 		params:               params,
 		address:              address,
 		blockNtfnChan:        make(chan uint32),
+		quitChan:             make(chan struct{}),
 		finalizedDepositChan: finalizedDepositChan,
+		spentChan:            make(chan struct{}),
 	}
 
 	depositStates := depoFsm.DepositStatesV0()
@@ -191,19 +202,24 @@ func NewFSM(ctx context.Context, deposit *Deposit, cfg *ManagerConfig,
 
 	depoFsm.ActionEntryFunc = depoFsm.updateDeposit
 
-	go func() {
+	go func(fsm *FSM) {
+		defer close(fsm.quitChan)
+
 		for {
 			select {
-			case currentHeight := <-depoFsm.blockNtfnChan:
+			case currentHeight := <-fsm.blockNtfnChan:
 				depoFsm.handleBlockNotification(
 					ctx, currentHeight,
 				)
 
+			case <-fsm.spentChan:
+				return
+
 			case <-ctx.Done():
 				return
 			}
 		}
-	}()
+	}(depoFsm)
 
 	return depoFsm, nil
 }

staticaddr/deposit/manager.go

@@ -136,19 +136,44 @@ func (m *Manager) Run(ctx context.Context, currentHeight uint32) error {
 		select {
 		case height := <-newBlockChan:
 			// Inform all active deposits about a new block arrival.
+			m.mu.Lock()
+			activeDeposits := make([]*FSM, 0, len(m.activeDeposits))
 			for _, fsm := range m.activeDeposits {
+				activeDeposits = append(activeDeposits, fsm)
+			}
+			m.mu.Unlock()
+
+			for _, fsm := range activeDeposits {
 				select {
 				case fsm.blockNtfnChan <- uint32(height):
 
+				case <-fsm.quitChan:
+					continue
+
 				case <-ctx.Done():
 					return ctx.Err()
 				}
 			}
+
 		case outpoint := <-m.finalizedDepositChan:
-			// If deposits notify us about their finalization, we
-			// update the manager's internal state and flush the
-			// finalized deposit from memory.
+			// If deposits notify us about their finalization, flush
+			// the finalized deposit from memory and end the fsm's
+			// block notif goroutine.
+			m.mu.Lock()
+			fsm := m.activeDeposits[outpoint]
+			m.mu.Unlock()
+
+			// End the fsm's block notif go routine.
+			select {
+			case fsm.spentChan <- struct{}{}:
+
+			case <-ctx.Done():
+				return ctx.Err()
+			}
+
+			m.mu.Lock()
 			delete(m.activeDeposits, outpoint)
+			m.mu.Unlock()
 
 		case err = <-newBlockErrChan:
 			return err
@@ -197,7 +222,9 @@ func (m *Manager) recoverDeposits(ctx context.Context) error {
 			}
 		}()
 
+		m.mu.Lock()
 		m.activeDeposits[d.OutPoint] = fsm
+		m.mu.Unlock()
 	}
 
 	return nil