Repository: mdlayher/netlink
Pull request #157 (2019-12-06 14:01:54)

Opportunistically decode attributes in AttributeDecoder
Hi Matt,

As discussed yesterday, this patch contributes two things:

- adds `AttributeDecoder.TypeFlags()` that returns the Type with an inverted `attrTypeMask`
- inverts the call chain between `UnmarshalAttributes()` and `ad.Next()`

This change is motivated by https://github.com/ti-mo/conntrack/issues/13, as I'd like to expose a `netfilter.AttributeDecoder/Encoder` and `conntrack.AttributeDecoder/Encoder` interface all the way up the call chain. If done right, decoding a whole netlink message can be done with only a single allocation of each layer's `Attribute` type. :scream: 

It doesn't save much on the benches in the repo, because it decodes all attributes anyway, but it allows for large efficiency gains in downstream libraries. Because attributes can now be decoded opportunistically, (eg. calling `ad.Next()` until you have the attributes you care about and calling it quits), a caller that uses AttributeDecoder will no longer implicitly decode the whole message. This potentially saves a lot of CPU time.

I also cut down on memory allocations, since AttributeDecoder now uses a single `netlink.Attribute` as a scratch buffer it repeatedly unmarshals into, completely removing allocations from the hot path.

This is illustrated by the benchmark results.

Before:
```
λ  netlink 38c649d ✓ go test -bench BenchmarkUnmarshalAttributes
goos: linux
goarch: amd64
pkg: github.com/mdlayher/netlink
BenchmarkUnmarshalAttributes/0-16  	354637596	         3.13 ns/op	       0 B/op	       0 allocs/op
BenchmarkUnmarshalAttributes/1-16  	 8755629	       134 ns/op	      33 B/op	       2 allocs/op
BenchmarkUnmarshalAttributes/8-16  	 1403520	       856 ns/op	     544 B/op	      12 allocs/op
BenchmarkUnmarshalAttributes/64-16 	  182522	      6417 ns/op	    8160 B/op	      71 allocs/op
BenchmarkUnmarshalAttributes/512-16         	   10000	    114465 ns/op	  294880 B/op	     522 allocs/op
```

After:
```
λ  netlink master ✓ go test -bench BenchmarkUnmarshalAttributes
goos: linux
goarch: amd64
pkg: github.com/mdlayher/netlink
BenchmarkUnmarshalAttributes/0-16  	13132034	        92.4 ns/op	     112 B/op	       1 allocs/op
BenchmarkUnmarshalAttributes/1-16  	 5177278	       234 ns/op	     145 B/op	       3 allocs/op
BenchmarkUnmarshalAttributes/8-16  	 1474484	       806 ns/op	     432 B/op	      10 allocs/op
BenchmarkUnmarshalAttributes/64-16 	  183102	      6580 ns/op	    6256 B/op	      66 allocs/op
BenchmarkUnmarshalAttributes/512-16         	   10000	    103788 ns/op	  278640 B/op	     514 allocs/op
```

The extra allocation visible in 0 and 1 is due to an added `make([]Attribute, 0, ad.Len())` in `UnmarshalAttributes` that pre-allocates the backing array of the output slice. This amortizes nicely over time.

The existing test suite was very helpful in weeding out any inconsistencies in behaviour, so for example I had to add an extra check in `UnmarshalAttributes` to make sure it explicitly returns a nil `[]Attribute` when no attributes where decoded.

Let me know what you think!


Patch:
diff --git a/attribute.go b/attribute.go
index b2ac945..61112ae 100644
--- a/attribute.go
+++ b/attribute.go
@@ -107,26 +107,27 @@ func MarshalAttributes(attrs []Attribute) ([]byte, error) {
 // It is recommend to use the AttributeDecoder type where possible instead of calling
 // UnmarshalAttributes and using package nlenc functions directly.
 func UnmarshalAttributes(b []byte) ([]Attribute, error) {
-	var attrs []Attribute
-	var i int
-	for {
-		if i > len(b) || len(b[i:]) == 0 {
-			break
-		}
 
-		var a Attribute
-		if err := (&a).unmarshal(b[i:]); err != nil {
-			return nil, err
-		}
+	ad, err := NewAttributeDecoder(b)
+	if err != nil {
+		return nil, err
+	}
 
-		if a.Length == 0 {
-			i += nlaHeaderLen
-			continue
-		}
+	// Return a nil slice when there are no attributes to decode.
+	if ad.Len() == 0 {
+		return nil, nil
+	}
 
-		i += nlaAlign(int(a.Length))
+	attrs := make([]Attribute, 0, ad.Len())
+
+	for ad.Next() {
+		if ad.attr().Length != 0 {
+			attrs = append(attrs, ad.attr())
+		}
+	}
 
-		attrs = append(attrs, a)
+	if err := ad.Err(); err != nil {
+		return nil, err
 	}
 
 	return attrs, nil
@@ -148,10 +149,14 @@ type AttributeDecoder struct {
 	// If not set, the native byte order will be used.
 	ByteOrder binary.ByteOrder
 
-	// The attributes being worked on, and the iterator index into the slice of
-	// attributes.
-	attrs []Attribute
-	i     int
+	// The current attribute being worked on.
+	a Attribute
+
+	// The slice of input bytes and its iterator index.
+	b []byte
+	i int
+
+	length int
 
 	// Any error encountered while decoding attributes.
 	err error
@@ -160,17 +165,21 @@ type AttributeDecoder struct {
 // NewAttributeDecoder creates an AttributeDecoder that unpacks Attributes
 // from b and prepares the decoder for iteration.
 func NewAttributeDecoder(b []byte) (*AttributeDecoder, error) {
-	attrs, err := UnmarshalAttributes(b)
-	if err != nil {
-		return nil, err
-	}
 
-	return &AttributeDecoder{
+	ad := &AttributeDecoder{
 		// By default, use native byte order.
 		ByteOrder: nlenc.NativeEndian(),
 
-		attrs: attrs,
-	}, nil
+		b: b,
+	}
+
+	var err error
+	ad.length, err = ad.available()
+	if err != nil {
+		return nil, err
+	}
+
+	return ad, nil
 }
 
 // Next advances the decoder to the next netlink attribute.  It returns false
@@ -181,34 +190,93 @@ func (ad *AttributeDecoder) Next() bool {
 		return false
 	}
 
-	ad.i++
+	// Exit if array pointer is at or beyond the end of the slice.
+	if ad.i >= len(ad.b) {
+		return false
+	}
+
+	if err := ad.a.unmarshal(ad.b[ad.i:]); err != nil {
+		ad.err = err
+		return false
+	}
+
+	// Advance the pointer by at least one header's length.
+	if int(ad.a.Length) < nlaHeaderLen {
+		ad.i += nlaHeaderLen
+	} else {
+		ad.i += nlaAlign(int(ad.a.Length))
+	}
 
-	// More attributes?
-	return len(ad.attrs) >= ad.i
+	return true
 }
 
 // Type returns the Attribute.Type field of the current netlink attribute
 // pointed to by the decoder.
 //
 // Type masks off the high bits of the netlink attribute type which may contain
-// the Nested and NetByteOrder flags. If you need direct access to these flags,
-// consider using UnmarshalAttributes instead.
+// the Nested and NetByteOrder flags. These can be obtained by calling TypeFlags.
 func (ad *AttributeDecoder) Type() uint16 {
 	// Mask off any flags stored in the high bits.
-	return ad.attr().Type & attrTypeMask
+	return ad.a.Type & attrTypeMask
+}
+
+// TypeFlags returns the two high bits of the Attribute.Type field of the current
+// netlink attribute pointed to by the decoder.
+//
+// These bits of the netlink attribute type are used for the Nested and NetByteOrder
+// flags, available as the Nested and NetByteOrder constants in this package.
+func (ad *AttributeDecoder) TypeFlags() uint16 {
+	return ad.a.Type & ^attrTypeMask
 }
 
 // Len returns the number of netlink attributes pointed to by the decoder.
-func (ad *AttributeDecoder) Len() int { return len(ad.attrs) }
+func (ad *AttributeDecoder) Len() int { return ad.length }
+
+// count scans the input slice to count the number of netlink attributes
+// that could be decoded by Next().
+func (ad *AttributeDecoder) available() (int, error) {
+
+	var i, count int
+	for {
+
+		// No more data to read.
+		if i >= len(ad.b) {
+			break
+		}
+
+		// Make sure there's at least a header's worth
+		// of data to read on each iteration.
+		if len(ad.b[i:]) < nlaHeaderLen {
+			return 0, errInvalidAttribute
+		}
+
+		// Extract the length of the attribute.
+		l := int(nlenc.Uint16(ad.b[i : i+2]))
+
+		// Ignore zero-length attributes.
+		if l != 0 {
+			count++
+		}
+
+		// Advance by at least a header's worth of bytes.
+		if l < nlaHeaderLen {
+			l = nlaHeaderLen
+		}
+
+		i += nlaAlign(l)
+	}
+
+	return count, nil
+}
 
 // attr returns the current Attribute pointed to by the decoder.
 func (ad *AttributeDecoder) attr() Attribute {
-	return ad.attrs[ad.i-1]
+	return ad.a
 }
 
 // data returns the Data field of the current Attribute pointed to by the decoder.
 func (ad *AttributeDecoder) data() []byte {
-	return ad.attr().Data
+	return ad.a.Data
 }
 
 // Err returns the first error encountered by the decoder.


Test patch:
diff --git a/attribute_test.go b/attribute_test.go
index eeecab2..01f9a35 100644
--- a/attribute_test.go
+++ b/attribute_test.go
@@ -754,6 +754,22 @@ func TestAttributeDecoderOK(t *testing.T) {
 				})
 			},
 		},
+		{
+			name: "typeflags",
+			attrs: []Attribute{{
+				Type: 0xffff,
+			}},
+			fn: func(ad *AttributeDecoder) {
+
+				if diff := cmp.Diff(ad.Type(), uint16(0x3fff)); diff != "" {
+					panicf("unexpected Type (-want +got):\n%s", diff)
+				}
+
+				if diff := cmp.Diff(ad.TypeFlags(), uint16(0xc000)); diff != "" {
+					panicf("unexpected TypeFlags (-want +got):\n%s", diff)
+				}
+			},
+		},
 	}
 
 	for _, tt := range tests {