From 90a59fe70f00a50033d859c1f5edbe9e46a87fa8 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Thu, 24 Sep 2020 09:53:07 +0200
Subject: [PATCH] router: call CleanStore on startup

---
 htlcswitch/switch.go |  8 ++++++++
 routing/mock_test.go |  7 +++++++
 routing/router.go    | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/htlcswitch/switch.go b/htlcswitch/switch.go
index 1e7db9fee5..72eebf59dc 100644
--- a/htlcswitch/switch.go
+++ b/htlcswitch/switch.go
@@ -438,6 +438,14 @@ func (s *Switch) GetPaymentResult(paymentID uint64, paymentHash lntypes.Hash,
 	return resultChan, nil
 }
 
+// CleanStore calls the underlying result store, telling it is safe to delete
+// all entries except the ones in the keepPids map. This should be called
+// preiodically to let the switch clean up payment results that we have
+// handled.
+func (s *Switch) CleanStore(keepPids map[uint64]struct{}) error {
+	return s.networkResults.cleanStore(keepPids)
+}
+
 // SendHTLC is used by other subsystems which aren't belong to htlc switch
 // package in order to send the htlc update. The paymentID used MUST be unique
 // for this HTLC, and MUST be used only once, otherwise the switch might reject
diff --git a/routing/mock_test.go b/routing/mock_test.go
index f5a6d3823c..a284cf57d7 100644
--- a/routing/mock_test.go
+++ b/routing/mock_test.go
@@ -72,6 +72,9 @@ func (m *mockPaymentAttemptDispatcher) GetPaymentResult(paymentID uint64,
 	return c, nil
 
 }
+func (m *mockPaymentAttemptDispatcher) CleanStore(map[uint64]struct{}) error {
+	return nil
+}
 
 func (m *mockPaymentAttemptDispatcher) setPaymentResult(
 	f func(firstHop lnwire.ShortChannelID) ([32]byte, error)) {
@@ -187,6 +190,10 @@ func (m *mockPayer) GetPaymentResult(paymentID uint64, _ lntypes.Hash,
 	}
 }
 
+func (m *mockPayer) CleanStore(pids map[uint64]struct{}) error {
+	return nil
+}
+
 type initArgs struct {
 	c *channeldb.PaymentCreationInfo
 }
diff --git a/routing/router.go b/routing/router.go
index 84fb9d0b18..af1e8d75f7 100644
--- a/routing/router.go
+++ b/routing/router.go
@@ -171,6 +171,14 @@ type PaymentAttemptDispatcher interface {
 	GetPaymentResult(paymentID uint64, paymentHash lntypes.Hash,
 		deobfuscator htlcswitch.ErrorDecrypter) (
 		<-chan *htlcswitch.PaymentResult, error)
+
+	// CleanStore calls the underlying result store, telling it is safe to
+	// delete all entries except the ones in the keepPids map. This should
+	// be called preiodically to let the switch clean up payment results
+	// that we have handled.
+	// NOTE: New payment attempts MUST NOT be made after the keepPids map
+	// has been created and this method has returned.
+	CleanStore(keepPids map[uint64]struct{}) error
 }
 
 // PaymentSessionSource is an interface that defines a source for the router to
@@ -538,6 +546,30 @@ func (r *ChannelRouter) Start() error {
 		return err
 	}
 
+	// Before we restart existing payments and start accepting more
+	// payments to be made, we clean the network result store of the
+	// Switch. We do this here at startup to ensure no more payments can be
+	// made concurrently, so we know the toKeep map will be up-to-date
+	// until the cleaning has finished.
+	toKeep := make(map[uint64]struct{})
+	for _, p := range payments {
+		payment, err := r.cfg.Control.FetchPayment(
+			p.Info.PaymentHash,
+		)
+		if err != nil {
+			return err
+		}
+
+		for _, a := range payment.HTLCs {
+			toKeep[a.AttemptID] = struct{}{}
+		}
+	}
+
+	log.Debugf("Cleaning network result store.")
+	if err := r.cfg.Payer.CleanStore(toKeep); err != nil {
+		return err
+	}
+
 	for _, payment := range payments {
 		log.Infof("Resuming payment with hash %v", payment.Info.PaymentHash)
 		r.wg.Add(1)