channeldb: implement a multi-source AddrSource

In this commit, we implement a version of the AddrSource interface which
merges the results of a set of AddrSource implementations. This will
later be used to merge the results of the channel db and graph db.
This commit is contained in:
Elle Mouton 2024-10-22 14:15:35 +02:00
parent 083d3c9d7c
commit 9537026df9
No known key found for this signature in database
GPG key ID: D7D916376026F177
2 changed files with 172 additions and 0 deletions

View file

@ -13,3 +13,56 @@ type AddrSource interface {
// key.
AddrsForNode(nodePub *btcec.PublicKey) ([]net.Addr, error)
}
// multiAddrSource is an implementation of AddrSource which gathers all the
// known addresses for a given node from multiple backends and de-duplicates the
// results.
type multiAddrSource struct {
sources []AddrSource
}
// NewMultiAddrSource constructs a new AddrSource which will query all the
// provided sources for a node's addresses and will then de-duplicate the
// results.
func NewMultiAddrSource(sources ...AddrSource) AddrSource {
return &multiAddrSource{
sources: sources,
}
}
// AddrsForNode returns all known addresses for the target node public key. It
// queries all the address sources provided and de-duplicates the results.
//
// NOTE: this implements the AddrSource interface.
func (c *multiAddrSource) AddrsForNode(nodePub *btcec.PublicKey) ([]net.Addr,
error) {
// The multiple address sources will likely contain duplicate addresses,
// so we use a map here to de-dup them.
dedupedAddrs := make(map[string]net.Addr)
// 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 {
addrs, err := src.AddrsForNode(nodePub)
if err != nil {
return nil, err
}
for _, addr := range addrs {
dedupedAddrs[addr.String()] = addr
}
}
// Convert the map into a list we can return.
addrs := make([]net.Addr, 0, len(dedupedAddrs))
for _, addr := range dedupedAddrs {
addrs = append(addrs, addr)
}
return addrs, nil
}
// A compile-time check to ensure that multiAddrSource implements the AddrSource
// interface.
var _ AddrSource = (*multiAddrSource)(nil)

View file

@ -0,0 +1,119 @@
package channeldb
import (
"net"
"testing"
"github.com/btcsuite/btcd/btcec/v2"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
var (
addr1 = &net.TCPAddr{IP: (net.IP)([]byte{0x1}), Port: 1}
addr2 = &net.TCPAddr{IP: (net.IP)([]byte{0x2}), Port: 2}
addr3 = &net.TCPAddr{IP: (net.IP)([]byte{0x3}), Port: 3}
)
// TestMultiAddrSource tests that the multiAddrSource correctly merges and
// deduplicates the results of a set of AddrSource implementations.
func TestMultiAddrSource(t *testing.T) {
t.Parallel()
var pk1 = newTestPubKey(t)
t.Run("both sources have results", func(t *testing.T) {
t.Parallel()
var (
src1 = newMockAddrSource(t)
src2 = newMockAddrSource(t)
)
t.Cleanup(func() {
src1.AssertExpectations(t)
src2.AssertExpectations(t)
})
// Let source 1 know of 2 addresses (addr 1 and 2) for node 1.
src1.On("AddrsForNode", pk1).Return(
[]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(
[]net.Addr{addr2, addr3}, 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, 2 and 3.
addrs, err := multiSrc.AddrsForNode(pk1)
require.NoError(t, err)
require.ElementsMatch(t, addrs, []net.Addr{addr1, addr2, addr3})
})
t.Run("only once source has results", func(t *testing.T) {
t.Parallel()
var (
src1 = newMockAddrSource(t)
src2 = newMockAddrSource(t)
)
t.Cleanup(func() {
src1.AssertExpectations(t)
src2.AssertExpectations(t)
})
// Let source 1 know of address 1 for node 1.
src1.On("AddrsForNode", pk1).Return(
[]net.Addr{addr1}, nil,
).Once()
src2.On("AddrsForNode", pk1).Return(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.
addrs, err := multiSrc.AddrsForNode(pk1)
require.NoError(t, err)
require.ElementsMatch(t, addrs, []net.Addr{addr1})
})
}
type mockAddrSource struct {
t *testing.T
mock.Mock
}
var _ AddrSource = (*mockAddrSource)(nil)
func newMockAddrSource(t *testing.T) *mockAddrSource {
return &mockAddrSource{t: t}
}
func (m *mockAddrSource) AddrsForNode(pub *btcec.PublicKey) ([]net.Addr,
error) {
args := m.Called(pub)
if args.Get(0) == nil {
return nil, args.Error(1)
}
addrs, ok := args.Get(0).([]net.Addr)
require.True(m.t, ok)
return addrs, args.Error(1)
}
func newTestPubKey(t *testing.T) *btcec.PublicKey {
priv, err := btcec.NewPrivateKey()
require.NoError(t, err)
return priv.PubKey()
}