Merge pull request #2501 from cfromknecht/batch-preimage-writes

htlcswitch: batch preimage writes/consistency fix

channeldb/witness_cache.go

@@ -1,10 +1,10 @@
 package channeldb
 
 import (
-	"crypto/sha256"
 	"fmt"
 
 	"github.com/coreos/bbolt"
+	"github.com/lightningnetwork/lnd/lntypes"
 )
 
 var (
@@ -70,12 +70,42 @@ func (d *DB) NewWitnessCache() *WitnessCache {
 	}
 }
 
-// AddWitness adds a new witness of wType to the witness cache. The type of the
-// witness will be used to map the witness to the key that will be used to look
-// it up.
-//
-// TODO(roasbeef): fake closure to map instead a constructor?
-func (w *WitnessCache) AddWitness(wType WitnessType, witness []byte) error {
+// witnessEntry is a key-value struct that holds each key -> witness pair, used
+// when inserting records into the cache.
+type witnessEntry struct {
+	key     []byte
+	witness []byte
+}
+
+// AddSha256Witnesses adds a batch of new sha256 preimages into the witness
+// cache. This is an alias for AddWitnesses that uses Sha256HashWitness as the
+// preimages' witness type.
+func (w *WitnessCache) AddSha256Witnesses(preimages ...lntypes.Preimage) error {
+	// Optimistically compute the preimages' hashes before attempting to
+	// start the db transaction.
+	entries := make([]witnessEntry, 0, len(preimages))
+	for i := range preimages {
+		hash := preimages[i].Hash()
+		entries = append(entries, witnessEntry{
+			key:     hash[:],
+			witness: preimages[i][:],
+		})
+	}
+
+	return w.addWitnessEntries(Sha256HashWitness, entries)
+}
+
+// addWitnessEntries inserts the witnessEntry key-value pairs into the cache,
+// using the appropriate witness type to segment the namespace of possible
+// witness types.
+func (w *WitnessCache) addWitnessEntries(wType WitnessType,
+	entries []witnessEntry) error {
+
+	// Exit early if there are no witnesses to add.
+	if len(entries) == 0 {
+		return nil
+	}
+
 	return w.db.Batch(func(tx *bbolt.Tx) error {
 		witnessBucket, err := tx.CreateBucketIfNotExists(witnessBucketKey)
 		if err != nil {
@@ -93,23 +123,32 @@ func (w *WitnessCache) AddWitness(wType WitnessType, witness []byte) error {
 			return err
 		}
 
-		// Now that we have the proper bucket for this witness, we'll map the
-		// witness type to the proper key.
-		var witnessKey []byte
-		switch wType {
-		case Sha256HashWitness:
-			key := sha256.Sum256(witness)
-			witnessKey = key[:]
+		for _, entry := range entries {
+			err = witnessTypeBucket.Put(entry.key, entry.witness)
+			if err != nil {
+				return err
+			}
 		}
 
-		return witnessTypeBucket.Put(witnessKey, witness)
+		return nil
 	})
 }
 
-// LookupWitness attempts to lookup a witness according to its type and also
+// LookupSha256Witness attempts to lookup the preimage for a sha256 hash. If
+// the witness isn't found, ErrNoWitnesses will be returned.
+func (w *WitnessCache) LookupSha256Witness(hash lntypes.Hash) (lntypes.Preimage, error) {
+	witness, err := w.lookupWitness(Sha256HashWitness, hash[:])
+	if err != nil {
+		return lntypes.Preimage{}, err
+	}
+
+	return lntypes.MakePreimage(witness)
+}
+
+// lookupWitness attempts to lookup a witness according to its type and also
 // its witness key. In the case that the witness isn't found, ErrNoWitnesses
 // will be returned.
-func (w *WitnessCache) LookupWitness(wType WitnessType, witnessKey []byte) ([]byte, error) {
+func (w *WitnessCache) lookupWitness(wType WitnessType, witnessKey []byte) ([]byte, error) {
 	var witness []byte
 	err := w.db.View(func(tx *bbolt.Tx) error {
 		witnessBucket := tx.Bucket(witnessBucketKey)
@@ -143,8 +182,13 @@ func (w *WitnessCache) LookupWitness(wType WitnessType, witnessKey []byte) ([]by
 	return witness, nil
 }
 
-// DeleteWitness attempts to delete a particular witness from the database.
-func (w *WitnessCache) DeleteWitness(wType WitnessType, witnessKey []byte) error {
+// DeleteSha256Witness attempts to delete a sha256 preimage identified by hash.
+func (w *WitnessCache) DeleteSha256Witness(hash lntypes.Hash) error {
+	return w.deleteWitness(Sha256HashWitness, hash[:])
+}
+
+// deleteWitness attempts to delete a particular witness from the database.
+func (w *WitnessCache) deleteWitness(wType WitnessType, witnessKey []byte) error {
 	return w.db.Batch(func(tx *bbolt.Tx) error {
 		witnessBucket, err := tx.CreateBucketIfNotExists(witnessBucketKey)
 		if err != nil {

channeldb/witness_cache_test.go

@@ -2,13 +2,14 @@ package channeldb
 
 import (
 	"crypto/sha256"
-	"reflect"
 	"testing"
+
+	"github.com/lightningnetwork/lnd/lntypes"
 )
 
-// TestWitnessCacheRetrieval tests that we're able to add and lookup new
-// witnesses to the witness cache.
-func TestWitnessCacheRetrieval(t *testing.T) {
+// TestWitnessCacheSha256Retrieval tests that we're able to add and lookup new
+// sha256 preimages to the witness cache.
+func TestWitnessCacheSha256Retrieval(t *testing.T) {
 	t.Parallel()
 
 	cdb, cleanUp, err := makeTestDB()
@@ -19,33 +20,41 @@ func TestWitnessCacheRetrieval(t *testing.T) {
 
 	wCache := cdb.NewWitnessCache()
 
-	// We'll be attempting to add then lookup a d simple hash witness
+	// We'll be attempting to add then lookup two simple sha256 preimages
 	// within this test.
-	witness := rev[:]
-	witnessKey := sha256.Sum256(witness)
+	preimage1 := lntypes.Preimage(rev)
+	preimage2 := lntypes.Preimage(key)
 
-	// First, we'll attempt to add the witness to the database.
-	err = wCache.AddWitness(Sha256HashWitness, witness)
+	preimages := []lntypes.Preimage{preimage1, preimage2}
+	hashes := []lntypes.Hash{preimage1.Hash(), preimage2.Hash()}
+
+	// First, we'll attempt to add the preimages to the database.
+	err = wCache.AddSha256Witnesses(preimages...)
 	if err != nil {
 		t.Fatalf("unable to add witness: %v", err)
 	}
 
-	// With the witness stored, we'll now attempt to look it up. We should
-	// get back the *exact* same witness as we originally stored.
-	dbWitness, err := wCache.LookupWitness(Sha256HashWitness, witnessKey[:])
-	if err != nil {
-		t.Fatalf("unable to look up witness: %v", err)
-	}
+	// With the preimages stored, we'll now attempt to look them up.
+	for i, hash := range hashes {
+		preimage := preimages[i]
 
-	if !reflect.DeepEqual(witness, dbWitness[:]) {
-		t.Fatalf("witnesses don't match: expected %x, got %x",
-			witness[:], dbWitness[:])
+		// We should get back the *exact* same preimage as we originally
+		// stored.
+		dbPreimage, err := wCache.LookupSha256Witness(hash)
+		if err != nil {
+			t.Fatalf("unable to look up witness: %v", err)
+		}
+
+		if preimage != dbPreimage {
+			t.Fatalf("witnesses don't match: expected %x, got %x",
+				preimage[:], dbPreimage[:])
+		}
 	}
 }
 
-// TestWitnessCacheDeletion tests that we're able to delete a single witness,
-// and also a class of witnesses from the cache.
-func TestWitnessCacheDeletion(t *testing.T) {
+// TestWitnessCacheSha256Deletion tests that we're able to delete a single
+// sha256 preimage, and also a class of witnesses from the cache.
+func TestWitnessCacheSha256Deletion(t *testing.T) {
 	t.Parallel()
 
 	cdb, cleanUp, err := makeTestDB()
@@ -56,37 +65,39 @@ func TestWitnessCacheDeletion(t *testing.T) {
 
 	wCache := cdb.NewWitnessCache()
 
-	// We'll start by adding two witnesses to the cache.
-	witness1 := rev[:]
-	witness1Key := sha256.Sum256(witness1)
-	if err := wCache.AddWitness(Sha256HashWitness, witness1); err != nil {
+	// We'll start by adding two preimages to the cache.
+	preimage1 := lntypes.Preimage(key)
+	hash1 := preimage1.Hash()
+
+	preimage2 := lntypes.Preimage(rev)
+	hash2 := preimage2.Hash()
+
+	if err := wCache.AddSha256Witnesses(preimage1); err != nil {
 		t.Fatalf("unable to add witness: %v", err)
 	}
 
-	witness2 := key[:]
-	witness2Key := sha256.Sum256(witness2)
-	if err := wCache.AddWitness(Sha256HashWitness, witness2); err != nil {
+	if err := wCache.AddSha256Witnesses(preimage2); err != nil {
 		t.Fatalf("unable to add witness: %v", err)
 	}
 
-	// We'll now delete the first witness. If we attempt to look it up, we
+	// We'll now delete the first preimage. If we attempt to look it up, we
 	// should get ErrNoWitnesses.
-	err = wCache.DeleteWitness(Sha256HashWitness, witness1Key[:])
+	err = wCache.DeleteSha256Witness(hash1)
 	if err != nil {
 		t.Fatalf("unable to delete witness: %v", err)
 	}
-	_, err = wCache.LookupWitness(Sha256HashWitness, witness1Key[:])
+	_, err = wCache.LookupSha256Witness(hash1)
 	if err != ErrNoWitnesses {
 		t.Fatalf("expected ErrNoWitnesses instead got: %v", err)
 	}
 
 	// Next, we'll attempt to delete the entire witness class itself. When
-	// we try to lookup the second witness, we should again get
+	// we try to lookup the second preimage, we should again get
 	// ErrNoWitnesses.
 	if err := wCache.DeleteWitnessClass(Sha256HashWitness); err != nil {
 		t.Fatalf("unable to delete witness class: %v", err)
 	}
-	_, err = wCache.LookupWitness(Sha256HashWitness, witness2Key[:])
+	_, err = wCache.LookupSha256Witness(hash2)
 	if err != ErrNoWitnesses {
 		t.Fatalf("expected ErrNoWitnesses instead got: %v", err)
 	}
@@ -107,8 +118,121 @@ func TestWitnessCacheUnknownWitness(t *testing.T) {
 
 	// We'll attempt to add a new, undefined witness type to the database.
 	// We should get an error.
-	err = wCache.AddWitness(234, key[:])
+	err = wCache.legacyAddWitnesses(234, key[:])
 	if err != ErrUnknownWitnessType {
 		t.Fatalf("expected ErrUnknownWitnessType, got %v", err)
 	}
 }
+
+// TestAddSha256Witnesses tests that insertion using AddSha256Witnesses behaves
+// identically to the insertion via the generalized interface.
+func TestAddSha256Witnesses(t *testing.T) {
+	cdb, cleanUp, err := makeTestDB()
+	if err != nil {
+		t.Fatalf("unable to make test database: %v", err)
+	}
+	defer cleanUp()
+
+	wCache := cdb.NewWitnessCache()
+
+	// We'll start by adding a witnesses to the cache using the generic
+	// AddWitnesses method.
+	witness1 := rev[:]
+	preimage1 := lntypes.Preimage(rev)
+	hash1 := preimage1.Hash()
+
+	witness2 := key[:]
+	preimage2 := lntypes.Preimage(key)
+	hash2 := preimage2.Hash()
+
+	var (
+		witnesses = [][]byte{witness1, witness2}
+		preimages = []lntypes.Preimage{preimage1, preimage2}
+		hashes    = []lntypes.Hash{hash1, hash2}
+	)
+
+	err = wCache.legacyAddWitnesses(Sha256HashWitness, witnesses...)
+	if err != nil {
+		t.Fatalf("unable to add witness: %v", err)
+	}
+
+	for i, hash := range hashes {
+		preimage := preimages[i]
+
+		dbPreimage, err := wCache.LookupSha256Witness(hash)
+		if err != nil {
+			t.Fatalf("unable to lookup witness: %v", err)
+		}
+
+		// Assert that the retrieved witness matches the original.
+		if dbPreimage != preimage {
+			t.Fatalf("retrieved witness mismatch, want: %x, "+
+				"got: %x", preimage, dbPreimage)
+		}
+
+		// We'll now delete the witness, as we'll be reinserting it
+		// using the specialized AddSha256Witnesses method.
+		err = wCache.DeleteSha256Witness(hash)
+		if err != nil {
+			t.Fatalf("unable to delete witness: %v", err)
+		}
+	}
+
+	// Now, add the same witnesses using the type-safe interface for
+	// lntypes.Preimages..
+	err = wCache.AddSha256Witnesses(preimages...)
+	if err != nil {
+		t.Fatalf("unable to add sha256 preimage: %v", err)
+	}
+
+	// Finally, iterate over the keys and assert that the returned witnesses
+	// match the original witnesses. This asserts that the specialized
+	// insertion method behaves identically to the generalized interface.
+	for i, hash := range hashes {
+		preimage := preimages[i]
+
+		dbPreimage, err := wCache.LookupSha256Witness(hash)
+		if err != nil {
+			t.Fatalf("unable to lookup witness: %v", err)
+		}
+
+		// Assert that the retrieved witness matches the original.
+		if dbPreimage != preimage {
+			t.Fatalf("retrieved witness mismatch, want: %x, "+
+				"got: %x", preimage, dbPreimage)
+		}
+	}
+}
+
+// legacyAddWitnesses adds a batch of new witnesses of wType to the witness
+// cache. The type of the witness will be used to map each witness to the key
+// that will be used to look it up. All witnesses should be of the same
+// WitnessType.
+//
+// NOTE: Previously this method exposed a generic interface for adding
+// witnesses, which has since been deprecated in favor of a strongly typed
+// interface for each witness class. We keep this method around to assert the
+// correctness of specialized witness adding methods.
+func (w *WitnessCache) legacyAddWitnesses(wType WitnessType,
+	witnesses ...[]byte) error {
+
+	// Optimistically compute the witness keys before attempting to start
+	// the db transaction.
+	entries := make([]witnessEntry, 0, len(witnesses))
+	for _, witness := range witnesses {
+		// Map each witness to its key by applying the appropriate
+		// transformation for the given witness type.
+		switch wType {
+		case Sha256HashWitness:
+			key := sha256.Sum256(witness)
+			entries = append(entries, witnessEntry{
+				key:     key[:],
+				witness: witness,
+			})
+		default:
+			return ErrUnknownWitnessType
+		}
+	}
+
+	return w.addWitnessEntries(wType, entries)
+}