From 4f7e871b42b5f359de2bc70163dcd53e708a39f2 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Thu, 30 Mar 2023 12:00:52 +0200
Subject: [PATCH] wtclient: add clarifying comments to the AddressIterator

---
 watchtower/wtclient/addr_iterator.go | 39 +++++++++++++++++++++++++++-
 1 file changed, 38 insertions(+), 1 deletion(-)

diff --git a/watchtower/wtclient/addr_iterator.go b/watchtower/wtclient/addr_iterator.go
index 8abcb3891..5ada7d5de 100644
--- a/watchtower/wtclient/addr_iterator.go
+++ b/watchtower/wtclient/addr_iterator.go
@@ -162,11 +162,15 @@ func (a *addressIterator) next(lock bool) (net.Addr, error) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
 
+	// In-case currentTopAddr is nil (meaning that Reset has not yet been
+	// called), return an error indicating this.
 	if a.currentTopAddr == nil {
 		return nil, ErrAddressesExhausted
 	}
 
-	// Set the next candidate to the subsequent element.
+	// Set the next candidate to the subsequent element. If we are at the
+	// end of the address list, this could mean setting currentTopAddr to
+	// nil.
 	a.currentTopAddr = a.currentTopAddr.Next()
 
 	for a.currentTopAddr != nil {
@@ -176,9 +180,26 @@ func (a *addressIterator) next(lock bool) (net.Addr, error) {
 		// Check whether this address is still considered a candidate.
 		// If it's not, we'll proceed to the next.
 		candidate, ok := a.candidates[addrID]
+
+		// If the address cannot be found in the candidate set, then
+		// this must mean that the Remove method was called for the
+		// address. The Remove method would have checked that the
+		// address is not the last one in the iterator and that it has
+		// no locks on it. It is therefor safe to remove.
 		if !ok {
+			// Grab the next address candidate. This might be nil
+			// if the iterator is on the last item in the list.
 			nextCandidate := a.currentTopAddr.Next()
+
+			// Remove the address from the list that is no longer
+			// in the candidate set.
 			a.addrList.Remove(a.currentTopAddr)
+
+			// Set the current top to the next candidate. This might
+			// mean setting it to nil if the iterator is on its last
+			// item in which case the loop will be exited and an
+			// ErrAddressesExhausted exhausted error will be
+			// returned.
 			a.currentTopAddr = nextCandidate
 			continue
 		}
@@ -234,9 +255,25 @@ func (a *addressIterator) peek(lock bool) net.Addr {
 
 		addrID := a.currentTopAddr.Value.(string)
 		candidate, ok := a.candidates[addrID]
+
+		// If the address cannot be found in the candidate set, then
+		// this must mean that the Remove method was called for the
+		// address. The Remove method would have checked that the
+		// address is not the last one in the iterator and that it has
+		// no locks on it. It is therefor safe to remove.
 		if !ok {
+			// Grab the next address candidate. This might be nil
+			// if the iterator is on the last item in the list.
 			nextCandidate := a.currentTopAddr.Next()
+
+			// Remove the address from the list that is no longer
+			// in the candidate set.
 			a.addrList.Remove(a.currentTopAddr)
+
+			// Set the current top to the next candidate. This might
+			// mean setting it to nil if the iterator is on its last
+			// item but this will be reset at the top of the for
+			// loop.
 			a.currentTopAddr = nextCandidate
 			continue
 		}