From 3648f2fa4b6d0b689db188210d830974916e4a1a Mon Sep 17 00:00:00 2001
From: Andras Banki-Horvath <bhandras@gmail.com>
Date: Wed, 16 Jun 2021 22:20:13 +0200
Subject: [PATCH] payments: prefetch payment on hot paths

---
 channeldb/payment_control.go | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/channeldb/payment_control.go b/channeldb/payment_control.go
index d024bf6bf..ada24e0c2 100644
--- a/channeldb/payment_control.go
+++ b/channeldb/payment_control.go
@@ -131,6 +131,7 @@ func (p *PaymentControl) InitPayment(paymentHash lntypes.Hash,
 		// from a previous execution of the batched db transaction.
 		updateErr = nil
 
+		prefetchPayment(tx, paymentHash)
 		bucket, err := createPaymentBucket(tx, paymentHash)
 		if err != nil {
 			return err
@@ -292,6 +293,7 @@ func (p *PaymentControl) RegisterAttempt(paymentHash lntypes.Hash,
 
 	var payment *MPPayment
 	err = kvdb.Batch(p.db.Backend, func(tx kvdb.RwTx) error {
+		prefetchPayment(tx, paymentHash)
 		bucket, err := fetchPaymentBucketUpdate(tx, paymentHash)
 		if err != nil {
 			return err
@@ -430,6 +432,7 @@ func (p *PaymentControl) updateHtlcKey(paymentHash lntypes.Hash,
 	err := kvdb.Batch(p.db.Backend, func(tx kvdb.RwTx) error {
 		payment = nil
 
+		prefetchPayment(tx, paymentHash)
 		bucket, err := fetchPaymentBucketUpdate(tx, paymentHash)
 		if err != nil {
 			return err
@@ -500,6 +503,7 @@ func (p *PaymentControl) Fail(paymentHash lntypes.Hash,
 		updateErr = nil
 		payment = nil
 
+		prefetchPayment(tx, paymentHash)
 		bucket, err := fetchPaymentBucketUpdate(tx, paymentHash)
 		if err == ErrPaymentNotInitiated {
 			updateErr = ErrPaymentNotInitiated
@@ -550,6 +554,7 @@ func (p *PaymentControl) FetchPayment(paymentHash lntypes.Hash) (
 
 	var payment *MPPayment
 	err := kvdb.View(p.db, func(tx kvdb.RTx) error {
+		prefetchPayment(tx, paymentHash)
 		bucket, err := fetchPaymentBucket(tx, paymentHash)
 		if err != nil {
 			return err
@@ -568,6 +573,26 @@ func (p *PaymentControl) FetchPayment(paymentHash lntypes.Hash) (
 	return payment, nil
 }
 
+// prefetchPayment attempts to prefetch as much of the payment as possible to
+// reduce DB roundtrips.
+func prefetchPayment(tx kvdb.RTx, paymentHash lntypes.Hash) {
+	rb := kvdb.RootBucket(tx)
+	kvdb.Prefetch(
+		rb,
+		[]string{
+			// Prefetch all keys in the payment's bucket.
+			string(paymentsRootBucket),
+			string(paymentHash[:]),
+		},
+		[]string{
+			// Prefetch all keys in the payment's htlc bucket.
+			string(paymentsRootBucket),
+			string(paymentHash[:]),
+			string(paymentHtlcsBucket),
+		},
+	)
+}
+
 // createPaymentBucket creates or fetches the sub-bucket assigned to this
 // payment hash.
 func createPaymentBucket(tx kvdb.RwTx, paymentHash lntypes.Hash) (