Decrease allocations in encode loop. (#86)

While investigating performance of HTTP/2 I noticed a substantial number
of allocations occuring in a hot loop in header compression. Specifically
I set up a simple harness that encoded headers 100 times, and then ran
that harness 10k times. This harness ended up performing 17 million
allocations, which is frankly a bit too high!

The fact that the number of allocations was in the millions suggested that
we were allocating some resources per header encode. That's not good:
assuming the target `ByteBuffer` is big enough we should not need to
perform new allocations to encode headers.

The allocations came from two places. Firstly, we were allocating in the
`HeaderTableStorage.indices(matching:)` function. This was because that
operation was returning a lazy collection, which required multiple heap
allocations for the closure contexts. Those needed to be eliminated.

Secondly, we were triggering a CoW of the `ByteBuffer` passed in to
`encode()`. This is the result of a complex bit of control flow that
fundamentally meant that the buffer was held in two places in that
method.

Removing these two sources of allocations dropped the count in my test
from 17 million to 176k, a decrease of 100x. We also saw a speed boost
of nearly 66% in terms of runtime of the microbenchmark.

- Removed `HeaderTableStorage.indices(matching:)`
- Implemented `HeaderTableStorage.closestMatch(name:value:), containing
   the repeated logic of the body of the two loops over the header
   indices.
- Used swap() and an empty byte buffer to ensure that encode does not
   cause a CoW operation.

Much faster and more memory efficient header encoding.

Sources/NIOHPACK/DynamicHeaderTable.swift

@@ -105,25 +105,12 @@ struct DynamicHeaderTable {
         // If we have a value, locate the index of the lowest header which contains that
         // value, but if no value matches, return the index of the lowest header with a
         // matching name alone.
-        var firstNameMatch: Int? = nil
-        for index in self.storage.indices(matching: name) {
-            if firstNameMatch == nil {
-                // record the first (most recent) index with a matching header name,
-                // in case there's no value match.
-                firstNameMatch = index
-            }
-
-            if self.storage.view(of: self.storage[index].value).matches(value) {
-                // this entry has both the name and the value we're seeking
-                return (index, true)
-            }
-        }
-
-        // no value matches -- but did we find a name?
-        if let index = firstNameMatch {
+        switch self.storage.closestMatch(name: name, value: value) {
+        case .full(let index):
+            return (index, true)
+        case .partial(let index):
             return (index, false)
-        } else {
-            // no matches at all
+        case .none:
             return nil
         }
     }

Sources/NIOHPACK/HPACKEncoder.swift

@@ -43,7 +43,7 @@ public struct HPACKEncoder {
     var headerIndexTable: IndexedHeaderTable
 
     private var state: EncoderState
-    private var buffer: ByteBuffer!
+    private var buffer: ByteBuffer
 
     /// Whether to use Huffman encoding.
     public let useHuffmanEncoding: Bool
@@ -110,6 +110,10 @@ public struct HPACKEncoder {
         self.headerIndexTable = IndexedHeaderTable(allocator: allocator, maxDynamicTableSize: maxDynamicTableSize)
         self.useHuffmanEncoding = useHuffmanEncoding
         self.state = .idle
+
+        // In my ideal world this allocation would be of size 0, but we need to have it be a little bit bigger to support the incremental encoding
+        // mode. I want to remove it: see https://github.com/apple/swift-nio-http2/issues/85.
+        self.buffer = allocator.buffer(capacity: 128)
     }
 
     /// Sets up the encoder to begin encoding a new header block.
@@ -120,9 +124,9 @@ public struct HPACKEncoder {
         if case .encoding = self.state {
             throw NIOHPACKErrors.EncoderAlreadyActive()
         }
-        
-        // create a buffer
-        self.buffer = allocator.buffer(capacity: 128)   // arbitrary size
+
+        self.buffer.clear()
+
         switch self.state {
         case .idle:
             self.state = .encoding
@@ -147,17 +151,23 @@ public struct HPACKEncoder {
         }
 
         self.state = .idle
-        defer { self.buffer = nil }
         return self.buffer
     }
 
     /// A one-shot encoder that writes to a provided buffer.
+    ///
+    /// In general this encoding mechanism is more efficient than the incremental one.
     public mutating func encode(headers: HPACKHeaders, to buffer: inout ByteBuffer) throws {
         if case .encoding = self.state {
             throw NIOHPACKErrors.EncoderAlreadyActive()
         }
-
-        self.buffer = buffer
+
+        swap(&self.buffer, &buffer)
+        defer {
+            swap(&self.buffer, &buffer)
+            self.state = .idle
+        }
+
         switch self.state {
         case .idle:
             self.state = .encoding
@@ -173,14 +183,8 @@ public struct HPACKEncoder {
         default:
             break
         }
-
-        do {
-            try self.append(headers: headers)
-            buffer = try self.endEncoding()
-        } catch {
-            _ = try? self.endEncoding()
-            throw error
-        }
+
+        try self.append(headers: headers)
     }
 
     /// Appends() headers in the default fashion: indexed if possible, literal+indexable if not.

Sources/NIOHPACK/HeaderTables.swift

@@ -118,11 +118,36 @@ struct HeaderTableStorage {
             self.buffer.equalCaseInsensitiveASCII(view: name, at: $0.name)
         }
     }
-
-    func indices<C: Collection>(matching name: C) -> LazyMatchingIndicesSequence where C.Element == UInt8 {
-        return self.headers.enumerated().lazy.filter {
-            self.buffer.equalCaseInsensitiveASCII(view: name, at: $1.name)
-        }.map { (idx, header) in idx }
+
+    enum MatchType {
+        case full(Int)
+        case partial(Int)
+        case none
+    }
+
+    func closestMatch<Name: Collection, Value: Collection>(name: Name, value: Value) -> MatchType where Name.Element == UInt8, Value.Element == UInt8 {
+        var partialIndex: Int? = nil
+
+        for (index, header) in self.headers.enumerated() {
+            // Check if the header name matches.
+            guard self.buffer.equalCaseInsensitiveASCII(view: name, at: header.name) else {
+                continue
+            }
+
+            if partialIndex == nil {
+                partialIndex = index
+            }
+
+            if self.view(of: header.value).matches(value) {
+                return .full(index)
+            }
+        }
+
+        if let partial = partialIndex {
+            return .partial(partial)
+        } else {
+            return .none
+        }
     }
 
     func firstIndex<C : Collection>(matching name: C) -> Int? where C.Element == UInt8 {

Sources/NIOHPACK/IndexedHeaderTable.swift

@@ -85,15 +85,14 @@ public struct IndexedHeaderTable {
         }
 
         var firstHeaderIndex: Int? = nil
-        for index in self.staticTable.indices(matching: name) {
-            // we've found a name, at least
-            if firstHeaderIndex == nil {
-                firstHeaderIndex = index
-            }
-
-            if self.staticTable.view(of: self.staticTable[index].value).matches(value) {
-                return (index, true)
-            }
+
+        switch self.staticTable.closestMatch(name: name, value: value) {
+        case .full(let index):
+            return (index, true)
+        case .partial(let index):
+            firstHeaderIndex = index
+        case .none:
+            break
         }
 
         // no complete match: search the dynamic table now