From a4a3c4192483a7e3b79bc676cb712b2a3ce55ee7 Mon Sep 17 00:00:00 2001
From: Joost Jager <joost.jager@gmail.com>
Date: Wed, 27 Nov 2019 14:19:15 +0100
Subject: [PATCH] channeldb: split cancel and add htlc updates

Previously the cancel and add actions were combined in a single map.
Nil values implictly signaled cancel actions. This wasn't very obvious.
Furthermore this split prepares for processing the adds and cancels
separately, which is more efficient if there are already two maps.
---
 channeldb/invoice_test.go   |  2 +-
 channeldb/invoices.go       | 45 +++++++++++++++++++------------------
 invoices/invoiceregistry.go |  8 +++----
 invoices/update.go          |  4 +++-
 4 files changed, 31 insertions(+), 28 deletions(-)

diff --git a/channeldb/invoice_test.go b/channeldb/invoice_test.go
index 82d960826..ac19025c7 100644
--- a/channeldb/invoice_test.go
+++ b/channeldb/invoice_test.go
@@ -686,7 +686,7 @@ func getUpdateInvoice(amt lnwire.MilliSatoshi) InvoiceUpdateCallback {
 		update := &InvoiceUpdateDesc{
 			Preimage: invoice.Terms.PaymentPreimage,
 			State:    ContractSettled,
-			Htlcs: map[CircuitKey]*HtlcAcceptDesc{
+			AddHtlcs: map[CircuitKey]*HtlcAcceptDesc{
 				{}: {
 					Amt: amt,
 				},
diff --git a/channeldb/invoices.go b/channeldb/invoices.go
index 81e273941..e6472fae9 100644
--- a/channeldb/invoices.go
+++ b/channeldb/invoices.go
@@ -316,10 +316,12 @@ type InvoiceUpdateDesc struct {
 	// State is the new state that this invoice should progress to.
 	State ContractState
 
-	// Htlcs describes the changes that need to be made to the invoice htlcs
-	// in the database. Htlc map entries with their value set should be
-	// added. If the map value is nil, the htlc should be canceled.
-	Htlcs map[CircuitKey]*HtlcAcceptDesc
+	// CancelHtlcs describes the htlcs that need to be canceled.
+	CancelHtlcs map[CircuitKey]struct{}
+
+	// AddHtlcs describes the newly accepted htlcs that need to be added to
+	// the invoice.
+	AddHtlcs map[CircuitKey]*HtlcAcceptDesc
 
 	// Preimage must be set to the preimage when state is settled.
 	Preimage lntypes.Preimage
@@ -1251,26 +1253,25 @@ func (d *DB) updateInvoice(hash lntypes.Hash, invoices, settleIndex *bbolt.Bucke
 
 	now := d.Now()
 
-	// Update htlc set.
-	for key, htlcUpdate := range update.Htlcs {
+	// Process cancel actions from update descriptor.
+	for key := range update.CancelHtlcs {
 		htlc, ok := invoice.Htlcs[key]
-
-		// No update means the htlc needs to be canceled.
-		if htlcUpdate == nil {
-			if !ok {
-				return nil, fmt.Errorf("unknown htlc %v", key)
-			}
-			if htlc.State != HtlcStateAccepted {
-				return nil, fmt.Errorf("can only cancel " +
-					"accepted htlcs")
-			}
-
-			htlc.State = HtlcStateCanceled
-			htlc.ResolveTime = now
-			invoice.AmtPaid -= htlc.Amt
-
-			continue
+		if !ok {
+			return nil, fmt.Errorf("unknown htlc %v", key)
 		}
+		if htlc.State != HtlcStateAccepted {
+			return nil, fmt.Errorf("can only cancel " +
+				"accepted htlcs")
+		}
+
+		htlc.State = HtlcStateCanceled
+		htlc.ResolveTime = now
+		invoice.AmtPaid -= htlc.Amt
+	}
+
+	// Process add actions from update descriptor.
+	for key, htlcUpdate := range update.AddHtlcs {
+		htlc, ok := invoice.Htlcs[key]
 
 		// Add new htlc paying to the invoice.
 		if ok {
diff --git a/invoices/invoiceregistry.go b/invoices/invoiceregistry.go
index 62998b741..e36a961b1 100644
--- a/invoices/invoiceregistry.go
+++ b/invoices/invoiceregistry.go
@@ -598,7 +598,7 @@ func (i *InvoiceRegistry) CancelInvoice(payHash lntypes.Hash) error {
 
 		// Mark individual held htlcs as canceled.
 		canceledHtlcs := make(
-			map[channeldb.CircuitKey]*channeldb.HtlcAcceptDesc,
+			map[channeldb.CircuitKey]struct{},
 		)
 		for key, htlc := range invoice.Htlcs {
 			switch htlc.State {
@@ -615,13 +615,13 @@ func (i *InvoiceRegistry) CancelInvoice(payHash lntypes.Hash) error {
 				continue
 			}
 
-			canceledHtlcs[key] = nil
+			canceledHtlcs[key] = struct{}{}
 		}
 
 		// Move invoice to the canceled state.
 		return &channeldb.InvoiceUpdateDesc{
-			Htlcs: canceledHtlcs,
-			State: channeldb.ContractCanceled,
+			CancelHtlcs: canceledHtlcs,
+			State:       channeldb.ContractCanceled,
 		}, nil
 	}
 
diff --git a/invoices/update.go b/invoices/update.go
index 1b015ecf0..dd438bd8a 100644
--- a/invoices/update.go
+++ b/invoices/update.go
@@ -130,7 +130,9 @@ func updateInvoice(ctx *invoiceUpdateCtx, inv *channeldb.Invoice) (
 		},
 	}
 
-	update := channeldb.InvoiceUpdateDesc{Htlcs: newHtlcs}
+	update := channeldb.InvoiceUpdateDesc{
+		AddHtlcs: newHtlcs,
+	}
 
 	// Don't update invoice state if we are accepting a duplicate payment.
 	// We do accept or settle the HTLC.