chainntnfs: Fix btcdnotify race condition causing double heap add

chainntnfs/btcdnotify/btcd.go

@@ -295,7 +295,7 @@ out:
 				// If the notification can be partially or
 				// fully dispatched, then we can skip the first
 				// phase for ntfns.
-				if b.attemptHistoricalDispatch(msg, currentHeight) {
+				if b.attemptHistoricalDispatch(msg) {
 					continue
 				}
 
@@ -418,8 +418,10 @@ out:
 // attemptHistoricalDispatch tries to use historical information to decide if a
 // notification ca be dispatched immediately, or is partially confirmed so it
 // can skip straight to the confirmations heap.
-func (b *BtcdNotifier) attemptHistoricalDispatch(msg *confirmationsNotification,
-	currentHeight int32) bool {
+//
+// Returns true if the transaction was either partially or completely confirmed
+func (b *BtcdNotifier) attemptHistoricalDispatch(
+	msg *confirmationsNotification) bool {
 
 	chainntnfs.Log.Infof("Attempting to trigger dispatch for %v from "+
 		"historical chain", msg.txid)
@@ -480,16 +482,16 @@ func (b *BtcdNotifier) attemptHistoricalDispatch(msg *confirmationsNotification,
 
 	// Otherwise, the transaction has only been *partially* confirmed, so
 	// we need to insert it into the confirmation heap.
-	confsLeft := msg.numConfirmations - uint32(tx.Confirmations)
-	confHeight := uint32(currentHeight) + confsLeft
+	// Find the block height at which this transaction will be confirmed
+	confHeight := uint32(block.Height) + msg.numConfirmations - 1
 	heapEntry := &confEntry{
 		msg,
 		confDetails,
 		confHeight,
 	}
 	heap.Push(b.confHeap, heapEntry)
 
-	return false
+	return true
 }
 
 // notifyBlockEpochs notifies all registered block epoch clients of the newly

chainntnfs/interface_test.go

@@ -604,6 +604,86 @@ func testTxConfirmedBeforeNtfnRegistration(miner *rpctest.Harness,
 	}
 }
 
+// Test the case of a notification consumer having forget or being delayed in
+// checking for a confirmation. This should not cause the notifier to stop
+// working
+func testLazyNtfnConsumer(miner *rpctest.Harness,
+	notifier chainntnfs.ChainNotifier, t *testing.T) {
+
+	// Create a transaction to be notified about. We'll register for
+	// notifications on this transaction but won't be prompt in checking them
+	txid, err := getTestTxId(miner)
+	if err != nil {
+		t.Fatalf("unable to create test tx: %v", err)
+	}
+
+	_, currentHeight, err := miner.Node.GetBestBlock()
+	if err != nil {
+		t.Fatalf("unable to get current height: %v", err)
+	}
+
+	numConfs := uint32(3)
+
+	// Add a block right before registering, this makes race conditions
+	// between the historical dispatcher and the normal dispatcher more obvious
+	if _, err := miner.Node.Generate(1); err != nil {
+		t.Fatalf("unable to generate blocks: %v", err)
+	}
+
+	firstConfIntent, err := notifier.RegisterConfirmationsNtfn(txid, numConfs,
+		uint32(currentHeight))
+	if err != nil {
+		t.Fatalf("unable to register ntfn: %v", err)
+	}
+
+	// Generate another 2 blocks, this should dispatch the confirm notification
+	if _, err := miner.Node.Generate(2); err != nil {
+		t.Fatalf("unable to generate blocks: %v", err)
+	}
+
+	// Now make another transaction, just because we haven't checked to see
+	// if the first transaction has confirmed doesn't mean that we shouldn't
+	// be able to see if this transaction confirms first
+	txid, err = getTestTxId(miner)
+	if err != nil {
+		t.Fatalf("unable to create test tx: %v", err)
+	}
+
+	_, currentHeight, err = miner.Node.GetBestBlock()
+	if err != nil {
+		t.Fatalf("unable to get current height: %v", err)
+	}
+
+	numConfs = 1
+
+	secondConfIntent, err := notifier.RegisterConfirmationsNtfn(txid, numConfs,
+		uint32(currentHeight))
+
+	if err != nil {
+		t.Fatalf("unable to register ntfn: %v", err)
+	}
+
+	if _, err := miner.Node.Generate(1); err != nil {
+		t.Fatalf("unable to generate blocks: %v", err)
+	}
+
+	select {
+	case <-secondConfIntent.Confirmed:
+		// Successfully receive the second notification
+		break
+	case <-time.After(30 * time.Second):
+		t.Fatalf("Second confirmation notification never received")
+	}
+
+	// Make sure the first tx confirmed successfully
+	select {
+	case <-firstConfIntent.Confirmed:
+		break
+	case <-time.After(30 * time.Second):
+		t.Fatalf("First confirmation notification never received")
+	}
+}
+
 // Tests the case in which a spend notification is requested for a spend that
 // has already been included in a block. In this case, the spend notification
 // should be dispatched immediately.
@@ -900,6 +980,10 @@ var ntfnTests = []testCase{
 		name: "cancel epoch ntfn",
 		test: testCancelEpochNtfn,
 	},
+	{
+		name: "lazy ntfn consumer",
+		test: testLazyNtfnConsumer,
+	},
 }
 
 // TestInterfaces tests all registered interfaces with a unified set of tests