multi: thread context through to AddrsForNode

Author: ellemouton

chanbackup/backup.go

@@ -1,6 +1,7 @@
 package chanbackup
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/btcsuite/btcd/wire"
@@ -24,15 +25,17 @@ type LiveChannelSource interface {
 // passed open channel. The backup includes all information required to restore
 // the channel, as well as addressing information so we can find the peer and
 // reconnect to them to initiate the protocol.
-func assembleChanBackup(addrSource channeldb.AddrSource,
+func assembleChanBackup(ctx context.Context, addrSource channeldb.AddrSource,
 	openChan *channeldb.OpenChannel) (*Single, error) {
 
 	log.Debugf("Crafting backup for ChannelPoint(%v)",
 		openChan.FundingOutpoint)
 
 	// First, we'll query the channel source to obtain all the addresses
 	// that are associated with the peer for this channel.
-	known, nodeAddrs, err := addrSource.AddrsForNode(openChan.IdentityPub)
+	known, nodeAddrs, err := addrSource.AddrsForNode(
+		ctx, openChan.IdentityPub,
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -90,7 +93,8 @@ func buildCloseTxInputs(
 // FetchBackupForChan attempts to create a plaintext static channel backup for
 // the target channel identified by its channel point. If we're unable to find
 // the target channel, then an error will be returned.
-func FetchBackupForChan(chanPoint wire.OutPoint, chanSource LiveChannelSource,
+func FetchBackupForChan(ctx context.Context, chanPoint wire.OutPoint,
+	chanSource LiveChannelSource,
 	addrSource channeldb.AddrSource) (*Single, error) {
 
 	// First, we'll query the channel source to see if the channel is known
@@ -104,7 +108,7 @@ func FetchBackupForChan(chanPoint wire.OutPoint, chanSource LiveChannelSource,
 
 	// Once we have the target channel, we can assemble the backup using
 	// the source to obtain any extra information that we may need.
-	staticChanBackup, err := assembleChanBackup(addrSource, targetChan)
+	staticChanBackup, err := assembleChanBackup(ctx, addrSource, targetChan)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create chan backup: %w", err)
 	}
@@ -114,7 +118,7 @@ func FetchBackupForChan(chanPoint wire.OutPoint, chanSource LiveChannelSource,
 
 // FetchStaticChanBackups will return a plaintext static channel back up for
 // all known active/open channels within the passed channel source.
-func FetchStaticChanBackups(chanSource LiveChannelSource,
+func FetchStaticChanBackups(ctx context.Context, chanSource LiveChannelSource,
 	addrSource channeldb.AddrSource) ([]Single, error) {
 
 	// First, we'll query the backup source for information concerning all
@@ -129,7 +133,7 @@ func FetchStaticChanBackups(chanSource LiveChannelSource,
 	// channel.
 	staticChanBackups := make([]Single, 0, len(openChans))
 	for _, openChan := range openChans {
-		chanBackup, err := assembleChanBackup(addrSource, openChan)
+		chanBackup, err := assembleChanBackup(ctx, addrSource, openChan)
 		if err != nil {
 			return nil, err
 		}

chanbackup/backup_test.go

@@ -1,6 +1,7 @@
 package chanbackup
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"testing"
@@ -61,8 +62,8 @@ func (m *mockChannelSource) addAddrsForNode(nodePub *btcec.PublicKey, addrs []ne
 	m.addrs[nodeKey] = addrs
 }
 
-func (m *mockChannelSource) AddrsForNode(nodePub *btcec.PublicKey) (bool,
-	[]net.Addr, error) {
+func (m *mockChannelSource) AddrsForNode(_ context.Context,
+	nodePub *btcec.PublicKey) (bool, []net.Addr, error) {
 
 	if m.failQuery {
 		return false, nil, fmt.Errorf("fail")
@@ -120,7 +121,8 @@ func TestFetchBackupForChan(t *testing.T) {
 	}
 	for i, testCase := range testCases {
 		_, err := FetchBackupForChan(
-			testCase.chanPoint, chanSource, chanSource,
+			context.Background(), testCase.chanPoint, chanSource,
+			chanSource,
 		)
 		switch {
 		// If this is a valid test case, and we failed, then we'll
@@ -141,6 +143,7 @@ func TestFetchBackupForChan(t *testing.T) {
 // channel source for all channels and construct a Single for each channel.
 func TestFetchStaticChanBackups(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	// First, we'll make the set of channels that we want to seed the
 	// channel source with. Both channels will be fully populated in the
@@ -162,7 +165,7 @@ func TestFetchStaticChanBackups(t *testing.T) {
 	// With the channel source populated, we'll now attempt to create a set
 	// of backups for all the channels. This should succeed, as all items
 	// are populated within the channel source.
-	backups, err := FetchStaticChanBackups(chanSource, chanSource)
+	backups, err := FetchStaticChanBackups(ctx, chanSource, chanSource)
 	require.NoError(t, err, "unable to create chan back ups")
 
 	if len(backups) != numChans {
@@ -177,7 +180,7 @@ func TestFetchStaticChanBackups(t *testing.T) {
 	copy(n[:], randomChan2.IdentityPub.SerializeCompressed())
 	delete(chanSource.addrs, n)
 
-	_, err = FetchStaticChanBackups(chanSource, chanSource)
+	_, err = FetchStaticChanBackups(ctx, chanSource, chanSource)
 	if err == nil {
 		t.Fatalf("query with incomplete information should fail")
 	}
@@ -186,7 +189,7 @@ func TestFetchStaticChanBackups(t *testing.T) {
 	// source at all, then we'll fail as well.
 	chanSource = newMockChannelSource()
 	chanSource.failQuery = true
-	_, err = FetchStaticChanBackups(chanSource, chanSource)
+	_, err = FetchStaticChanBackups(ctx, chanSource, chanSource)
 	if err == nil {
 		t.Fatalf("query should fail")
 	}

chanbackup/pubsub.go

@@ -2,6 +2,7 @@ package chanbackup
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"net"
 	"os"
@@ -81,7 +82,8 @@ type ChannelNotifier interface {
 	// synchronization point to ensure that the chanbackup.SubSwapper does
 	// not miss any channel open or close events in the period between when
 	// it's created, and when it requests the channel subscription.
-	SubscribeChans(map[wire.OutPoint]struct{}) (*ChannelSubscription, error)
+	SubscribeChans(context.Context,
+		map[wire.OutPoint]struct{}) (*ChannelSubscription, error)
 }
 
 // SubSwapper subscribes to new updates to the open channel state, and then
@@ -119,16 +121,17 @@ type SubSwapper struct {
 // set of channels, and the required interfaces to be notified of new channel
 // updates, pack a multi backup, and swap the current best backup from its
 // storage location.
-func NewSubSwapper(startingChans []Single, chanNotifier ChannelNotifier,
-	keyRing keychain.KeyRing, backupSwapper Swapper) (*SubSwapper, error) {
+func NewSubSwapper(ctx context.Context, startingChans []Single,
+	chanNotifier ChannelNotifier, keyRing keychain.KeyRing,
+	backupSwapper Swapper) (*SubSwapper, error) {
 
 	// First, we'll subscribe to the latest set of channel updates given
 	// the set of channels we already know of.
 	knownChans := make(map[wire.OutPoint]struct{})
 	for _, chanBackup := range startingChans {
 		knownChans[chanBackup.FundingOutpoint] = struct{}{}
 	}
-	chanEvents, err := chanNotifier.SubscribeChans(knownChans)
+	chanEvents, err := chanNotifier.SubscribeChans(ctx, knownChans)
 	if err != nil {
 		return nil, err
 	}

chanbackup/pubsub_test.go

@@ -1,6 +1,7 @@
 package chanbackup
 
 import (
+	"context"
 	"fmt"
 	"testing"
 	"time"
@@ -62,8 +63,8 @@ func newMockChannelNotifier() *mockChannelNotifier {
 	}
 }
 
-func (m *mockChannelNotifier) SubscribeChans(chans map[wire.OutPoint]struct{}) (
-	*ChannelSubscription, error) {
+func (m *mockChannelNotifier) SubscribeChans(_ context.Context,
+	_ map[wire.OutPoint]struct{}) (*ChannelSubscription, error) {
 
 	if m.fail {
 		return nil, fmt.Errorf("fail")
@@ -80,6 +81,7 @@ func (m *mockChannelNotifier) SubscribeChans(chans map[wire.OutPoint]struct{}) (
 // channel subscription, then the entire sub-swapper will fail to start.
 func TestNewSubSwapperSubscribeFail(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	keyRing := &lnencrypt.MockKeyRing{}
 
@@ -88,10 +90,8 @@ func TestNewSubSwapperSubscribeFail(t *testing.T) {
 		fail: true,
 	}
 
-	_, err := NewSubSwapper(nil, &chanNotifier, keyRing, &swapper)
-	if err == nil {
-		t.Fatalf("expected fail due to lack of subscription")
-	}
+	_, err := NewSubSwapper(ctx, nil, &chanNotifier, keyRing, &swapper)
+	require.Error(t, err)
 }
 
 func assertExpectedBackupSwap(t *testing.T, swapper *mockSwapper,
@@ -158,7 +158,9 @@ func TestSubSwapperIdempotentStartStop(t *testing.T) {
 	var chanNotifier mockChannelNotifier
 
 	swapper := newMockSwapper(keyRing)
-	subSwapper, err := NewSubSwapper(nil, &chanNotifier, keyRing, swapper)
+	subSwapper, err := NewSubSwapper(
+		context.Background(), nil, &chanNotifier, keyRing, swapper,
+	)
 	require.NoError(t, err, "unable to init subSwapper")
 
 	if err := subSwapper.Start(); err != nil {
@@ -224,7 +226,8 @@ func TestSubSwapperUpdater(t *testing.T) {
 	// With our channel set created, we'll make a fresh sub swapper
 	// instance to begin our test.
 	subSwapper, err := NewSubSwapper(
-		initialChanSet, chanNotifier, keyRing, swapper,
+		context.Background(), initialChanSet, chanNotifier, keyRing,
+		swapper,
 	)
 	require.NoError(t, err, "unable to make swapper")
 	if err := subSwapper.Start(); err != nil {

channel_notifier.go

@@ -1,6 +1,7 @@
 package lnd
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/btcsuite/btcd/wire"
@@ -31,7 +32,8 @@ type channelNotifier struct {
 // the channel subscription.
 //
 // NOTE: This is part of the chanbackup.ChannelNotifier interface.
-func (c *channelNotifier) SubscribeChans(startingChans map[wire.OutPoint]struct{}) (
+func (c *channelNotifier) SubscribeChans(ctx context.Context,
+	startingChans map[wire.OutPoint]struct{}) (
 	*chanbackup.ChannelSubscription, error) {
 
 	ltndLog.Infof("Channel backup proxy channel notifier starting")
@@ -46,7 +48,7 @@ func (c *channelNotifier) SubscribeChans(startingChans map[wire.OutPoint]struct{
 	// confirmed channels.
 	sendChanOpenUpdate := func(newOrPendingChan *channeldb.OpenChannel) {
 		_, nodeAddrs, err := c.addrs.AddrsForNode(
-			newOrPendingChan.IdentityPub,
+			ctx, newOrPendingChan.IdentityPub,
 		)
 		if err != nil {
 			pub := newOrPendingChan.IdentityPub

channeldb/addr_source.go

@@ -1,6 +1,7 @@
 package channeldb
 
 import (
+	"context"
 	"errors"
 	"net"
 
@@ -13,7 +14,8 @@ type AddrSource interface {
 	// AddrsForNode returns all known addresses for the target node public
 	// key. The returned boolean must indicate if the given node is unknown
 	// to the backing source.
-	AddrsForNode(nodePub *btcec.PublicKey) (bool, []net.Addr, error)
+	AddrsForNode(ctx context.Context,
+		nodePub *btcec.PublicKey) (bool, []net.Addr, error)
 }
 
 // multiAddrSource is an implementation of AddrSource which gathers all the
@@ -38,8 +40,8 @@ func NewMultiAddrSource(sources ...AddrSource) AddrSource {
 // node.
 //
 // NOTE: this implements the AddrSource interface.
-func (c *multiAddrSource) AddrsForNode(nodePub *btcec.PublicKey) (bool,
-	[]net.Addr, error) {
+func (c *multiAddrSource) AddrsForNode(ctx context.Context,
+	nodePub *btcec.PublicKey) (bool, []net.Addr, error) {
 
 	if len(c.sources) == 0 {
 		return false, nil, errors.New("no address sources")
@@ -55,7 +57,7 @@ func (c *multiAddrSource) AddrsForNode(nodePub *btcec.PublicKey) (bool,
 	// Iterate over all the address sources and query each one for the
 	// addresses it has for the node in question.
 	for _, src := range c.sources {
-		isKnown, addrs, err := src.AddrsForNode(nodePub)
+		isKnown, addrs, err := src.AddrsForNode(ctx, nodePub)
 		if err != nil {
 			return false, nil, err
 		}

channeldb/addr_source_test.go

@@ -1,6 +1,7 @@
 package channeldb
 
 import (
+	"context"
 	"net"
 	"testing"
 
@@ -19,6 +20,7 @@ var (
 // deduplicates the results of a set of AddrSource implementations.
 func TestMultiAddrSource(t *testing.T) {
 	t.Parallel()
+	ctx := context.Background()
 
 	var pk1 = newTestPubKey(t)
 
@@ -35,12 +37,12 @@ func TestMultiAddrSource(t *testing.T) {
 		})
 
 		// Let source 1 know of 2 addresses (addr 1 and 2) for node 1.
-		src1.On("AddrsForNode", pk1).Return(
+		src1.On("AddrsForNode", ctx, pk1).Return(
 			true, []net.Addr{addr1, addr2}, nil,
 		).Once()
 
 		// Let source 2 know of 2 addresses (addr 2 and 3) for node 1.
-		src2.On("AddrsForNode", pk1).Return(
+		src2.On("AddrsForNode", ctx, pk1).Return(
 			true, []net.Addr{addr2, addr3}, nil,
 			[]net.Addr{addr2, addr3}, nil,
 		).Once()
@@ -51,7 +53,7 @@ func TestMultiAddrSource(t *testing.T) {
 
 		// Query it for the addresses known for node 1. The results
 		// should contain addr 1, 2 and 3.
-		known, addrs, err := multiSrc.AddrsForNode(pk1)
+		known, addrs, err := multiSrc.AddrsForNode(ctx, pk1)
 		require.NoError(t, err)
 		require.True(t, known)
 		require.ElementsMatch(t, addrs, []net.Addr{addr1, addr2, addr3})
@@ -70,18 +72,18 @@ func TestMultiAddrSource(t *testing.T) {
 		})
 
 		// Let source 1 know of address 1 for node 1.
-		src1.On("AddrsForNode", pk1).Return(
+		src1.On("AddrsForNode", ctx, pk1).Return(
 			true, []net.Addr{addr1}, nil,
 		).Once()
-		src2.On("AddrsForNode", pk1).Return(false, nil, nil).Once()
+		src2.On("AddrsForNode", ctx, pk1).Return(false, nil, nil).Once()
 
 		// Create a multi-addr source that consists of both source 1
 		// and 2.
 		multiSrc := NewMultiAddrSource(src1, src2)
 
 		// Query it for the addresses known for node 1. The results
 		// should contain addr 1.
-		known, addrs, err := multiSrc.AddrsForNode(pk1)
+		known, addrs, err := multiSrc.AddrsForNode(ctx, pk1)
 		require.NoError(t, err)
 		require.True(t, known)
 		require.ElementsMatch(t, addrs, []net.Addr{addr1})
@@ -103,13 +105,13 @@ func TestMultiAddrSource(t *testing.T) {
 		// and 2. Neither source known of node 1.
 		multiSrc := NewMultiAddrSource(src1, src2)
 
-		src1.On("AddrsForNode", pk1).Return(false, nil, nil).Once()
-		src2.On("AddrsForNode", pk1).Return(false, nil, nil).Once()
+		src1.On("AddrsForNode", ctx, pk1).Return(false, nil, nil).Once()
+		src2.On("AddrsForNode", ctx, pk1).Return(false, nil, nil).Once()
 
 		// Query it for the addresses known for node 1. It should return
 		// false to indicate that the node is unknown to all backing
 		// sources.
-		known, addrs, err := multiSrc.AddrsForNode(pk1)
+		known, addrs, err := multiSrc.AddrsForNode(ctx, pk1)
 		require.NoError(t, err)
 		require.False(t, known)
 		require.Empty(t, addrs)
@@ -127,10 +129,10 @@ func newMockAddrSource(t *testing.T) *mockAddrSource {
 	return &mockAddrSource{t: t}
 }
 
-func (m *mockAddrSource) AddrsForNode(pub *btcec.PublicKey) (bool, []net.Addr,
-	error) {
+func (m *mockAddrSource) AddrsForNode(ctx context.Context,
+	pub *btcec.PublicKey) (bool, []net.Addr, error) {
 
-	args := m.Called(pub)
+	args := m.Called(ctx, pub)
 	if args.Get(1) == nil {
 		return args.Bool(0), nil, args.Error(2)
 	}