lnd.xprv/autopilot/prefattach_test.go
Johan T. Halseth 5b1e72a019
autopilot/interface+agent: return NodeScore from NodeScores
Since NodeScores no longer returns fully populated AttachmentDirectives,
we make this explicit by defining a new type NodeScore that includes a
subset of what the AttachmentDirective does.
2019-01-08 10:23:50 +01:00

502 lines
12 KiB
Go

package autopilot
import (
"bytes"
"io/ioutil"
"os"
"testing"
"time"
prand "math/rand"
"github.com/btcsuite/btcd/btcec"
"github.com/btcsuite/btcutil"
"github.com/lightningnetwork/lnd/channeldb"
)
type genGraphFunc func() (testGraph, func(), error)
type testGraph interface {
ChannelGraph
addRandChannel(*btcec.PublicKey, *btcec.PublicKey,
btcutil.Amount) (*ChannelEdge, *ChannelEdge, error)
addRandNode() (*btcec.PublicKey, error)
}
func newDiskChanGraph() (testGraph, func(), error) {
// First, create a temporary directory to be used for the duration of
// this test.
tempDirName, err := ioutil.TempDir("", "channeldb")
if err != nil {
return nil, nil, err
}
// Next, create channeldb for the first time.
cdb, err := channeldb.Open(tempDirName)
if err != nil {
return nil, nil, err
}
cleanUp := func() {
cdb.Close()
os.RemoveAll(tempDirName)
}
return &databaseChannelGraph{
db: cdb.ChannelGraph(),
}, cleanUp, nil
}
var _ testGraph = (*databaseChannelGraph)(nil)
func newMemChanGraph() (testGraph, func(), error) {
return newMemChannelGraph(), nil, nil
}
var _ testGraph = (*memChannelGraph)(nil)
var chanGraphs = []struct {
name string
genFunc genGraphFunc
}{
{
name: "disk_graph",
genFunc: newDiskChanGraph,
},
{
name: "mem_graph",
genFunc: newMemChanGraph,
},
}
// TestPrefAttachmentSelectEmptyGraph ensures that when passed an
// empty graph, the NodeSores function always returns a score of 0.
func TestPrefAttachmentSelectEmptyGraph(t *testing.T) {
const (
maxChanSize = btcutil.Amount(btcutil.SatoshiPerBitcoin)
)
prefAttach := NewPrefAttachment()
// Create a random public key, which we will query to get a score for.
pub, err := randKey()
if err != nil {
t.Fatalf("unable to generate key: %v", err)
}
nodes := map[NodeID]struct{}{
NewNodeID(pub): {},
}
for _, graph := range chanGraphs {
success := t.Run(graph.name, func(t1 *testing.T) {
graph, cleanup, err := graph.genFunc()
if err != nil {
t1.Fatalf("unable to create graph: %v", err)
}
if cleanup != nil {
defer cleanup()
}
// With the necessary state initialized, we'll now
// attempt to get the score for this one node.
const walletFunds = btcutil.SatoshiPerBitcoin
scores, err := prefAttach.NodeScores(graph, nil,
walletFunds, nodes)
if err != nil {
t1.Fatalf("unable to select attachment "+
"directives: %v", err)
}
// Since the graph is empty, we expect the score to be
// 0, giving an empty return map.
if len(scores) != 0 {
t1.Fatalf("expected empty score map, "+
"instead got %v ", len(scores))
}
})
if !success {
break
}
}
}
// completeGraph is a helper method that adds numNodes fully connected nodes to
// the graph.
func completeGraph(t *testing.T, g testGraph, numNodes int) {
const chanCapacity = btcutil.SatoshiPerBitcoin
nodes := make(map[int]*btcec.PublicKey)
for i := 0; i < numNodes; i++ {
for j := i + 1; j < numNodes; j++ {
node1 := nodes[i]
node2 := nodes[j]
edge1, edge2, err := g.addRandChannel(
node1, node2, chanCapacity)
if err != nil {
t.Fatalf("unable to generate channel: %v", err)
}
if node1 == nil {
pubKeyBytes := edge1.Peer.PubKey()
nodes[i], err = btcec.ParsePubKey(
pubKeyBytes[:], btcec.S256(),
)
if err != nil {
t.Fatalf("unable to parse pubkey: %v",
err)
}
}
if node2 == nil {
pubKeyBytes := edge2.Peer.PubKey()
nodes[j], err = btcec.ParsePubKey(
pubKeyBytes[:], btcec.S256(),
)
if err != nil {
t.Fatalf("unable to parse pubkey: %v",
err)
}
}
}
}
}
// TestPrefAttachmentSelectTwoVertexes ensures that when passed a
// graph with only two eligible vertexes, then both are given the same score,
// and the funds are appropriately allocated across each peer.
func TestPrefAttachmentSelectTwoVertexes(t *testing.T) {
t.Parallel()
prand.Seed(time.Now().Unix())
const (
maxChanSize = btcutil.Amount(btcutil.SatoshiPerBitcoin)
)
for _, graph := range chanGraphs {
success := t.Run(graph.name, func(t1 *testing.T) {
graph, cleanup, err := graph.genFunc()
if err != nil {
t1.Fatalf("unable to create graph: %v", err)
}
if cleanup != nil {
defer cleanup()
}
prefAttach := NewPrefAttachment()
// For this set, we'll load the memory graph with two
// nodes, and a random channel connecting them.
const chanCapacity = btcutil.SatoshiPerBitcoin
edge1, edge2, err := graph.addRandChannel(nil, nil, chanCapacity)
if err != nil {
t1.Fatalf("unable to generate channel: %v", err)
}
// We also add a third, non-connected node to the graph.
_, err = graph.addRandNode()
if err != nil {
t1.Fatalf("unable to add random node: %v", err)
}
// Get the score for all nodes found in the graph at
// this point.
nodes := make(map[NodeID]struct{})
if err := graph.ForEachNode(func(n Node) error {
nodes[n.PubKey()] = struct{}{}
return nil
}); err != nil {
t1.Fatalf("unable to traverse graph: %v", err)
}
if len(nodes) != 3 {
t1.Fatalf("expected 2 nodes, found %d", len(nodes))
}
// With the necessary state initialized, we'll now
// attempt to get our candidates channel score given
// the current state of the graph.
candidates, err := prefAttach.NodeScores(graph, nil,
maxChanSize, nodes)
if err != nil {
t1.Fatalf("unable to select attachment "+
"directives: %v", err)
}
// We expect two candidates, since one of the nodes
// doesn't have any channels.
if len(candidates) != 2 {
t1.Fatalf("2 nodes should be scored, "+
"instead %v were", len(candidates))
}
// The candidates should be amongst the two edges
// created above.
for nodeID, candidate := range candidates {
edge1Pub := edge1.Peer.PubKey()
edge2Pub := edge2.Peer.PubKey()
switch {
case bytes.Equal(nodeID[:], edge1Pub[:]):
case bytes.Equal(nodeID[:], edge2Pub[:]):
default:
t1.Fatalf("attached to unknown node: %x",
nodeID[:])
}
// Since each of the nodes has 1 channel, out
// of only one channel in the graph, we expect
// their score to be 0.5.
expScore := float64(0.5)
if candidate.Score != expScore {
t1.Fatalf("expected candidate score "+
"to be %v, instead was %v",
expScore, candidate.Score)
}
}
})
if !success {
break
}
}
}
// TestPrefAttachmentSelectGreedyAllocation tests that if upon
// returning node scores, the NodeScores method will attempt to greedily
// allocate all funds to each vertex (up to the max channel size).
func TestPrefAttachmentSelectGreedyAllocation(t *testing.T) {
t.Parallel()
prand.Seed(time.Now().Unix())
const (
maxChanSize = btcutil.Amount(btcutil.SatoshiPerBitcoin)
)
for _, graph := range chanGraphs {
success := t.Run(graph.name, func(t1 *testing.T) {
graph, cleanup, err := graph.genFunc()
if err != nil {
t1.Fatalf("unable to create graph: %v", err)
}
if cleanup != nil {
defer cleanup()
}
prefAttach := NewPrefAttachment()
const chanCapacity = btcutil.SatoshiPerBitcoin
// Next, we'll add 3 nodes to the graph, creating an
// "open triangle topology".
edge1, _, err := graph.addRandChannel(nil, nil,
chanCapacity)
if err != nil {
t1.Fatalf("unable to create channel: %v", err)
}
peerPubBytes := edge1.Peer.PubKey()
peerPub, err := btcec.ParsePubKey(
peerPubBytes[:], btcec.S256(),
)
if err != nil {
t.Fatalf("unable to parse pubkey: %v", err)
}
_, _, err = graph.addRandChannel(
peerPub, nil, chanCapacity,
)
if err != nil {
t1.Fatalf("unable to create channel: %v", err)
}
// At this point, there should be three nodes in the
// graph, with node node having two edges.
numNodes := 0
twoChans := false
nodes := make(map[NodeID]struct{})
if err := graph.ForEachNode(func(n Node) error {
numNodes++
nodes[n.PubKey()] = struct{}{}
numChans := 0
err := n.ForEachChannel(func(c ChannelEdge) error {
numChans++
return nil
})
if err != nil {
return err
}
twoChans = twoChans || (numChans == 2)
return nil
}); err != nil {
t1.Fatalf("unable to traverse graph: %v", err)
}
if numNodes != 3 {
t1.Fatalf("expected 3 nodes, instead have: %v",
numNodes)
}
if !twoChans {
t1.Fatalf("expected node to have two channels")
}
// We'll now begin our test, modeling the available
// wallet balance to be 5.5 BTC. We're shooting for a
// 50/50 allocation, and have 3 BTC in channels. As a
// result, the heuristic should try to greedily
// allocate funds to channels.
scores, err := prefAttach.NodeScores(graph, nil,
maxChanSize, nodes)
if err != nil {
t1.Fatalf("unable to select attachment "+
"directives: %v", err)
}
if len(scores) != len(nodes) {
t1.Fatalf("all nodes should be scored, "+
"instead %v were", len(scores))
}
// The candidates should have a non-zero score, and
// have the max chan size funds recommended channel
// size.
for _, candidate := range scores {
if candidate.Score == 0 {
t1.Fatalf("Expected non-zero score")
}
}
// Imagine a few channels are being opened, and there's
// only 0.5 BTC left. That should leave us with channel
// candidates of that size.
const remBalance = btcutil.SatoshiPerBitcoin * 0.5
scores, err = prefAttach.NodeScores(graph, nil,
remBalance, nodes)
if err != nil {
t1.Fatalf("unable to select attachment "+
"directives: %v", err)
}
if len(scores) != len(nodes) {
t1.Fatalf("all nodes should be scored, "+
"instead %v were", len(scores))
}
// Check that the recommended channel sizes are now the
// remaining channel balance.
for _, candidate := range scores {
if candidate.Score == 0 {
t1.Fatalf("Expected non-zero score")
}
}
})
if !success {
break
}
}
}
// TestPrefAttachmentSelectSkipNodes ensures that if a node was
// already selected as a channel counterparty, then that node will get a score
// of zero during scoring.
func TestPrefAttachmentSelectSkipNodes(t *testing.T) {
t.Parallel()
prand.Seed(time.Now().Unix())
const (
maxChanSize = btcutil.Amount(btcutil.SatoshiPerBitcoin)
)
for _, graph := range chanGraphs {
success := t.Run(graph.name, func(t1 *testing.T) {
graph, cleanup, err := graph.genFunc()
if err != nil {
t1.Fatalf("unable to create graph: %v", err)
}
if cleanup != nil {
defer cleanup()
}
prefAttach := NewPrefAttachment()
// Next, we'll create a simple topology of two nodes,
// with a single channel connecting them.
const chanCapacity = btcutil.SatoshiPerBitcoin
_, _, err = graph.addRandChannel(nil, nil,
chanCapacity)
if err != nil {
t1.Fatalf("unable to create channel: %v", err)
}
nodes := make(map[NodeID]struct{})
if err := graph.ForEachNode(func(n Node) error {
nodes[n.PubKey()] = struct{}{}
return nil
}); err != nil {
t1.Fatalf("unable to traverse graph: %v", err)
}
if len(nodes) != 2 {
t1.Fatalf("expected 2 nodes, found %d", len(nodes))
}
// With our graph created, we'll now get the scores for
// all nodes in the graph.
scores, err := prefAttach.NodeScores(graph, nil,
maxChanSize, nodes)
if err != nil {
t1.Fatalf("unable to select attachment "+
"directives: %v", err)
}
if len(scores) != len(nodes) {
t1.Fatalf("all nodes should be scored, "+
"instead %v were", len(scores))
}
// THey should all have a score, and a maxChanSize
// channel size recommendation.
for _, candidate := range scores {
if candidate.Score == 0 {
t1.Fatalf("Expected non-zero score")
}
}
// We'll simulate a channel update by adding the nodes
// to our set of channels.
var chans []Channel
for _, candidate := range scores {
chans = append(chans,
Channel{
Node: candidate.NodeID,
},
)
}
// If we attempt to make a call to the NodeScores
// function, without providing any new information,
// then all nodes should have a score of zero, since we
// already got channels to them.
scores, err = prefAttach.NodeScores(graph, chans,
maxChanSize, nodes)
if err != nil {
t1.Fatalf("unable to select attachment "+
"directives: %v", err)
}
// Since all should be given a score of 0, the map
// should be empty.
if len(scores) != 0 {
t1.Fatalf("expected empty score map, "+
"instead got %v ", len(scores))
}
})
if !success {
break
}
}
}