staticaddr: quit deposit fsm handler when finialized

Author: hieblmi

staticaddr/deposit/fsm.go

@@ -141,6 +141,11 @@ type FSM struct {
 
 	blockNtfnChan chan uint32
 
+	// quitChan stops after the FSM stops consuming blockNtfnChan.
+	quitChan 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,6 +172,7 @@ func NewFSM(ctx context.Context, deposit *Deposit, cfg *ManagerConfig,
 		params:               params,
 		address:              address,
 		blockNtfnChan:        make(chan uint32),
+		quitChan:             make(chan struct{}),
 		finalizedDepositChan: finalizedDepositChan,
 	}
 
@@ -191,10 +197,12 @@ 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,
 				)
@@ -203,7 +211,7 @@ func NewFSM(ctx context.Context, deposit *Deposit, cfg *ManagerConfig,
 				return
 			}
 		}
-	}()
+	}(depoFsm)
 
 	return depoFsm, nil
 }

staticaddr/deposit/manager.go

@@ -136,19 +136,31 @@ 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.
+			m.mu.Lock()
 			delete(m.activeDeposits, outpoint)
+			m.mu.Unlock()
 
 		case err = <-newBlockErrChan:
 			return err
@@ -197,7 +209,9 @@ func (m *Manager) recoverDeposits(ctx context.Context) error {
 			}
 		}()
 
+		m.mu.Lock()
 		m.activeDeposits[d.OutPoint] = fsm
+		m.mu.Unlock()
 	}
 
 	return nil