From 2f3bf6ee63bdd296adff167169299fca86eb90a3 Mon Sep 17 00:00:00 2001
From: Alok Menghrajani <alok@squareup.com>
Date: Thu, 18 Oct 2018 12:51:57 -0700
Subject: [PATCH] Fix bug in transactions cache handling

---
 backend/electrum_backend.go      | 28 +++++++------------
 backend/electrum_backend_test.go | 46 ++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 19 deletions(-)
 create mode 100644 backend/electrum_backend_test.go

diff --git a/backend/electrum_backend.go b/backend/electrum_backend.go
index 31373a0..3ba2cd6 100644
--- a/backend/electrum_backend.go
+++ b/backend/electrum_backend.go
@@ -327,14 +327,7 @@ func (eb *ElectrumBackend) processTxRequest(node *electrum.Node, txHash string)
 		eb.txRequests <- txHash
 		return err
 	}
-	height, err := eb.getTxHeight(txHash)
-	if err != nil {
-		// a transaction thas is being asked for is not in the cache, yet (most likely)
-		// TODO: we should have a retry counter and fail gracefully if a transaction fails
-		//       too many times.
-		eb.txRequests <- txHash
-		return nil
-	}
+	height := eb.getTxHeight(txHash)
 
 	eb.txResponses <- &TxResponse{
 		Hash:   txHash,
@@ -345,15 +338,15 @@ func (eb *ElectrumBackend) processTxRequest(node *electrum.Node, txHash string)
 	return nil
 }
 
-func (eb *ElectrumBackend) getTxHeight(txHash string) (int64, error) {
+func (eb *ElectrumBackend) getTxHeight(txHash string) int64 {
 	eb.transactionsMu.Lock()
 	defer eb.transactionsMu.Unlock()
 
 	height, exists := eb.transactions[txHash]
 	if !exists {
-		return -1, fmt.Errorf("no block height information for transaction %s yet", txHash)
+		log.Panicf("transactions cache miss for %s", txHash)
 	}
-	return height, nil
+	return height
 }
 
 func (eb *ElectrumBackend) processAddrRequest(node *electrum.Node, addr *deriver.Address) error {
@@ -386,18 +379,15 @@ func (eb *ElectrumBackend) processAddrRequest(node *electrum.Node, addr *deriver
 }
 
 func (eb *ElectrumBackend) cacheTxs(txs []*electrum.Transaction) {
+	eb.transactionsMu.Lock()
+	defer eb.transactionsMu.Unlock()
+
 	for _, tx := range txs {
-		eb.transactionsMu.Lock()
 		height, exists := eb.transactions[tx.Hash]
-		if exists {
-			if height != int64(tx.Height) {
-				panic(fmt.Sprintf("inconsistent cache: %s %d != %d", tx.Hash, height, tx.Height))
-			}
-			eb.transactionsMu.Unlock()
-			return
+		if exists && (height != int64(tx.Height)) {
+			panic(fmt.Sprintf("inconsistent cache: %s %d != %d", tx.Hash, height, tx.Height))
 		}
 		eb.transactions[tx.Hash] = int64(tx.Height)
-		eb.transactionsMu.Unlock()
 	}
 }
 
diff --git a/backend/electrum_backend_test.go b/backend/electrum_backend_test.go
new file mode 100644
index 0000000..96bb83f
--- /dev/null
+++ b/backend/electrum_backend_test.go
@@ -0,0 +1,46 @@
+package backend
+
+import (
+	"github.com/square/beancounter/backend/electrum"
+	"github.com/square/beancounter/deriver"
+	"github.com/square/beancounter/utils"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestTransactionCache(t *testing.T) {
+	// TODO: refactor ElectrumBackend to make it easier to test
+
+	eb := &ElectrumBackend{
+		nodes:            make(map[string]*electrum.Node),
+		blacklistedNodes: make(map[string]struct{}),
+		network:          utils.Testnet,
+		addrRequests:     make(chan *deriver.Address, 2*maxPeers),
+		addrResponses:    make(chan *AddrResponse, 2*maxPeers),
+		txRequests:       make(chan string, 2*maxPeers),
+		txResponses:      make(chan *TxResponse, 2*maxPeers),
+
+		peersRequests: make(chan struct{}),
+		transactions:  make(map[string]int64),
+		doneCh:        make(chan bool),
+	}
+
+	tx1 := electrum.Transaction{Hash: "aaaaaa", Height: 100}
+	tx2 := electrum.Transaction{Hash: "bbbbbb", Height: 100}
+	tx3 := electrum.Transaction{Hash: "cccccc", Height: 101}
+	badTx := electrum.Transaction{Hash: "aaaaaa", Height: 102}
+
+	eb.cacheTxs([]*electrum.Transaction{&tx1, &tx2})
+
+	assert.Equal(t, int64(tx1.Height), eb.getTxHeight(tx1.Hash))
+	assert.Equal(t, int64(tx2.Height), eb.getTxHeight(tx2.Hash))
+	assert.Panics(t, func() { eb.getTxHeight(tx3.Hash) })
+
+	eb.cacheTxs([]*electrum.Transaction{&tx2, &tx3})
+
+	assert.Equal(t, int64(tx1.Height), eb.getTxHeight(tx1.Hash))
+	assert.Equal(t, int64(tx2.Height), eb.getTxHeight(tx2.Hash))
+	assert.Equal(t, int64(tx3.Height), eb.getTxHeight(tx3.Hash))
+
+	assert.Panics(t, func() { eb.cacheTxs([]*electrum.Transaction{&badTx}) })
+}