p2p/discover: improve discv5 NODES response packing

p2p/discover/v5_udp.go

@@ -24,7 +24,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"math"
 	"net"
 	"sync"
 	"time"
@@ -41,7 +40,6 @@ const (
 	lookupRequestLimit      = 3  // max requests against a single node during lookup
 	findnodeResultLimit     = 16 // applies in FINDNODE handler
 	totalNodesResponseLimit = 5  // applies in waitForNodes
-	nodesResponseItemLimit  = 3  // applies in sendNodes
 
 	respTimeoutV5 = 700 * time.Millisecond
 )
@@ -832,17 +830,27 @@ func packNodes(reqid []byte, nodes []*enode.Node) []*v5wire.Nodes {
 		return []*v5wire.Nodes{{ReqID: reqid, Total: 1}}
 	}
 
-	total := uint8(math.Ceil(float64(len(nodes)) / 3))
+	// This limit represents the available space for nodes in output packets. Maximum
+	// packet size is 1280, and out of this ~80 bytes will be taken up by the packet
+	// frame. So limiting to 1000 bytes here leaves 200 bytes for other fields of the
+	// NODES message, which is a lot.
+	const sizeLimit = 1000
+
 	var resp []*v5wire.Nodes
 	for len(nodes) > 0 {
-		p := &v5wire.Nodes{ReqID: reqid, Total: total}
-		items := min(nodesResponseItemLimit, len(nodes))
-		for i := 0; i < items; i++ {
-			p.Nodes = append(p.Nodes, nodes[i].Record())
+		p := &v5wire.Nodes{ReqID: reqid}
+		size := uint64(0)
+		for len(nodes) > 0 && size < sizeLimit {
+			r := nodes[0].Record()
+			size += r.Size()
+			p.Nodes = append(p.Nodes, r)
+			nodes = nodes[1:]
 		}
-		nodes = nodes[items:]
 		resp = append(resp, p)
 	}
+	for _, msg := range resp {
+		msg.Total = uint8(len(resp))
+	}
 	return resp
 }
 

p2p/discover/v5_udp_test.go

@@ -160,7 +160,7 @@ func TestUDPv5_findnodeHandling(t *testing.T) {
 	defer test.close()
 
 	// Create test nodes and insert them into the table.
-	nodes253 := nodesAtDistance(test.table.self().ID(), 253, 10)
+	nodes253 := nodesAtDistance(test.table.self().ID(), 253, 16)
 	nodes249 := nodesAtDistance(test.table.self().ID(), 249, 4)
 	nodes248 := nodesAtDistance(test.table.self().ID(), 248, 10)
 	fillTable(test.table, wrapNodes(nodes253))
@@ -185,15 +185,15 @@ func TestUDPv5_findnodeHandling(t *testing.T) {
 
 	// This request gets all the distance-253 nodes.
 	test.packetIn(&v5wire.Findnode{ReqID: []byte{4}, Distances: []uint{253}})
-	test.expectNodes([]byte{4}, 4, nodes253)
+	test.expectNodes([]byte{4}, 2, nodes253)
 
 	// This request gets all the distance-249 nodes and some more at 248 because
 	// the bucket at 249 is not full.
 	test.packetIn(&v5wire.Findnode{ReqID: []byte{5}, Distances: []uint{249, 248}})
 	var nodes []*enode.Node
 	nodes = append(nodes, nodes249...)
 	nodes = append(nodes, nodes248[:10]...)
-	test.expectNodes([]byte{5}, 5, nodes)
+	test.expectNodes([]byte{5}, 1, nodes)
 }
 
 func (test *udpV5Test) expectNodes(wantReqID []byte, wantTotal uint8, wantNodes []*enode.Node) {
@@ -207,9 +207,6 @@ func (test *udpV5Test) expectNodes(wantReqID []byte, wantTotal uint8, wantNodes
 			if !bytes.Equal(p.ReqID, wantReqID) {
 				test.t.Fatalf("wrong request ID %v in response, want %v", p.ReqID, wantReqID)
 			}
-			if len(p.Nodes) > 3 {
-				test.t.Fatalf("too many nodes in response")
-			}
 			if p.Total != wantTotal {
 				test.t.Fatalf("wrong total response count %d, want %d", p.Total, wantTotal)
 			}

p2p/enr/enr.go

@@ -96,6 +96,24 @@ type pair struct {
 	v rlp.RawValue
 }
 
+// Size returns the encoded size of the record.
+func (r *Record) Size() uint64 {
+	if r.raw != nil {
+		return uint64(len(r.raw))
+	}
+	return computeSize(r)
+}
+
+func computeSize(r *Record) uint64 {
+	size := uint64(rlp.IntSize(r.seq))
+	size += rlp.BytesSize(r.signature)
+	for _, p := range r.pairs {
+		size += rlp.StringSize(p.k)
+		size += uint64(len(p.v))
+	}
+	return rlp.ListSize(size)
+}
+
 // Seq returns the sequence number.
 func (r *Record) Seq() uint64 {
 	return r.seq

p2p/enr/enr_test.go

@@ -169,6 +169,32 @@ func TestDirty(t *testing.T) {
 	}
 }
 
+func TestSize(t *testing.T) {
+	var r Record
+
+	// Empty record size is 3 bytes.
+	// Unsigned records cannot be encoded, but they could, the encoding
+	// would be [ 0, 0 ] -> 0xC28080.
+	assert.Equal(t, uint64(3), r.Size())
+
+	// Add one attribute. The size increases to 5, the encoding
+	// would be [ 0, 0, "k", "v" ] -> 0xC58080C26B76.
+	r.Set(WithEntry("k", "v"))
+	assert.Equal(t, uint64(5), r.Size())
+
+	// Now add a signature.
+	nodeid := []byte{1, 2, 3, 4, 5, 6, 7, 8}
+	signTest(nodeid, &r)
+	assert.Equal(t, uint64(45), r.Size())
+	enc, _ := rlp.EncodeToBytes(&r)
+	if r.Size() != uint64(len(enc)) {
+		t.Error("Size() not equal encoded length", len(enc))
+	}
+	if r.Size() != computeSize(&r) {
+		t.Error("Size() not equal computed size", computeSize(&r))
+	}
+}
+
 func TestSeq(t *testing.T) {
 	var r Record
 
@@ -268,8 +294,11 @@ func TestSignEncodeAndDecodeRandom(t *testing.T) {
 	}
 
 	require.NoError(t, signTest([]byte{5}, &r))
-	_, err := rlp.EncodeToBytes(r)
+
+	enc, err := rlp.EncodeToBytes(r)
 	require.NoError(t, err)
+	require.Equal(t, uint64(len(enc)), r.Size())
+	require.Equal(t, uint64(len(enc)), computeSize(&r))
 
 	for k, v := range pairs {
 		desc := fmt.Sprintf("key %q", k)

rlp/raw.go

@@ -28,6 +28,38 @@ type RawValue []byte
 
 var rawValueType = reflect.TypeOf(RawValue{})
 
+// StringSize returns the encoded size of a string.
+func StringSize(s string) uint64 {
+	switch {
+	case len(s) == 0:
+		return 1
+	case len(s) == 1:
+		if s[0] < 0x7f {
+			return 1
+		} else {
+			return 2
+		}
+	default:
+		return uint64(headsize(uint64(len(s))) + len(s))
+	}
+}
+
+// BytesSize returns the encoded size of a byte slice.
+func BytesSize(b []byte) uint64 {
+	switch {
+	case len(b) == 0:
+		return 1
+	case len(b) == 1:
+		if b[0] < 0x7f {
+			return 1
+		} else {
+			return 2
+		}
+	default:
+		return uint64(headsize(uint64(len(b))) + len(b))
+	}
+}
+
 // ListSize returns the encoded size of an RLP list with the given
 // content size.
 func ListSize(contentSize uint64) uint64 {

rlp/raw_test.go

@@ -283,3 +283,32 @@ func TestAppendUint64Random(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestBytesSize(t *testing.T) {
+	tests := []struct {
+		v    []byte
+		size uint64
+	}{
+		{v: []byte{}, size: 1},
+		{v: []byte{0x1}, size: 1},
+		{v: []byte{0x7E}, size: 1},
+		{v: []byte{0x7F}, size: 2},
+		{v: []byte{0x80}, size: 2},
+		{v: []byte{0xFF}, size: 2},
+		{v: []byte{0xFF, 0xF0}, size: 3},
+		{v: make([]byte, 55), size: 56},
+		{v: make([]byte, 56), size: 58},
+	}
+
+	for _, test := range tests {
+		s := BytesSize(test.v)
+		if s != test.size {
+			t.Errorf("BytesSize(%#x) -> %d, want %d", test.v, s, test.size)
+		}
+		s = StringSize(string(test.v))
+		if s != test.size {
+			t.Errorf("StringSize(%#x) -> %d, want %d", test.v, s, test.size)
+		}
+
+	}
+}