diff --git a/mock.go b/mock.go index 1eccbea6..6c1f3fb6 100644 --- a/mock.go +++ b/mock.go @@ -123,6 +123,7 @@ func (m *mockNotfier) RegisterSpendNtfn(outpoint *wire.OutPoint, _ []byte, type mockSpendNotifier struct { *mockNotfier spendMap map[wire.OutPoint][]chan *chainntnfs.SpendDetail + spends map[wire.OutPoint]*chainntnfs.SpendDetail mtx sync.Mutex } @@ -132,6 +133,7 @@ func makeMockSpendNotifier() *mockSpendNotifier { confChannel: make(chan *chainntnfs.TxConfirmation), }, spendMap: make(map[wire.OutPoint][]chan *chainntnfs.SpendDetail), + spends: make(map[wire.OutPoint]*chainntnfs.SpendDetail), } } @@ -140,8 +142,22 @@ func (m *mockSpendNotifier) RegisterSpendNtfn(outpoint *wire.OutPoint, m.mtx.Lock() defer m.mtx.Unlock() - spendChan := make(chan *chainntnfs.SpendDetail) - m.spendMap[*outpoint] = append(m.spendMap[*outpoint], spendChan) + spendChan := make(chan *chainntnfs.SpendDetail, 1) + if detail, ok := m.spends[*outpoint]; ok { + // Deliver spend immediately if details are already known. + spendChan <- &chainntnfs.SpendDetail{ + SpentOutPoint: detail.SpentOutPoint, + SpendingHeight: detail.SpendingHeight, + SpendingTx: detail.SpendingTx, + SpenderTxHash: detail.SpenderTxHash, + SpenderInputIndex: detail.SpenderInputIndex, + } + } else { + // Otherwise, queue the notification for delivery if the spend + // is ever received. + m.spendMap[*outpoint] = append(m.spendMap[*outpoint], spendChan) + } + return &chainntnfs.SpendEvent{ Spend: spendChan, Cancel: func() { @@ -156,16 +172,30 @@ func (m *mockSpendNotifier) Spend(outpoint *wire.OutPoint, height int32, m.mtx.Lock() defer m.mtx.Unlock() + txnHash := txn.TxHash() + details := &chainntnfs.SpendDetail{ + SpentOutPoint: outpoint, + SpendingHeight: height, + SpendingTx: txn, + SpenderTxHash: &txnHash, + SpenderInputIndex: outpoint.Index, + } + + // Cache details in case of late registration. + if _, ok := m.spends[*outpoint]; !ok { + m.spends[*outpoint] = details + } + + // Deliver any backlogged spend notifications. if spendChans, ok := m.spendMap[*outpoint]; ok { delete(m.spendMap, *outpoint) for _, spendChan := range spendChans { - txnHash := txn.TxHash() spendChan <- &chainntnfs.SpendDetail{ - SpentOutPoint: outpoint, - SpendingHeight: height, - SpendingTx: txn, - SpenderTxHash: &txnHash, - SpenderInputIndex: outpoint.Index, + SpentOutPoint: details.SpentOutPoint, + SpendingHeight: details.SpendingHeight, + SpendingTx: details.SpendingTx, + SpenderTxHash: details.SpenderTxHash, + SpenderInputIndex: details.SpenderInputIndex, } } }