Remove allocator from context (#73)

hummingbird-project · Aug 14, 2024 · b199fd6 · b199fd6
1 parent 3d694eb
commit b199fd6
Show file tree

Hide file tree

Showing 11 changed files with 19 additions and 27 deletions.
diff --git a/Sources/HummingbirdWSClient/WebSocketClient.swift b/Sources/HummingbirdWSClient/WebSocketClient.swift
@@ -43,11 +43,9 @@ public struct WebSocketClient {
     /// Basic context implementation of ``WebSocketContext``.
     /// Used by non-router web socket handle function
     public struct Context: WebSocketContext {
-        public let allocator: ByteBufferAllocator
         public let logger: Logger
 
-        package init(allocator: ByteBufferAllocator, logger: Logger) {
-            self.allocator = allocator
+        package init(logger: Logger) {
             self.logger = logger
         }
     }

diff --git a/Sources/HummingbirdWSClient/WebSocketClientChannel.swift b/Sources/HummingbirdWSClient/WebSocketClientChannel.swift
@@ -104,7 +104,7 @@ struct WebSocketClientChannel: ClientConnectionChannel {
                     autoPing: self.configuration.autoPing
                 ),
                 asyncChannel: webSocketChannel,
-                context: WebSocketClient.Context(allocator: webSocketChannel.channel.allocator, logger: logger),
+                context: WebSocketClient.Context(logger: logger),
                 handler: self.handler
             )
         case .notUpgraded:

diff --git a/Sources/HummingbirdWSCompression/PerMessageDeflateExtension.swift b/Sources/HummingbirdWSCompression/PerMessageDeflateExtension.swift
@@ -196,7 +196,11 @@ struct PerMessageDeflateExtension: WebSocketExtension {
                     unmaskedData.writeBytes([0, 0, 255, 255])
                     self.state = .idle
                 }
-                frame.data = try unmaskedData.decompressStream(with: self.decompressor, maxSize: maxSize, allocator: context.allocator)
+                frame.data = try unmaskedData.decompressStream(
+                    with: self.decompressor,
+                    maxSize: maxSize,
+                    allocator: ByteBufferAllocator()
+                )
                 frame.maskKey = nil
                 if resetStream, frame.fin {
                     try self.decompressor.resetStream()
@@ -241,7 +245,7 @@ struct PerMessageDeflateExtension: WebSocketExtension {
                     newFrame.rsv1 = true
                     self.sendState = .sendingMessage
                 }
-                newFrame.data = try newFrame.data.compressStream(with: self.compressor, flush: .sync, allocator: context.allocator)
+                newFrame.data = try newFrame.data.compressStream(with: self.compressor, flush: .sync, allocator: ByteBufferAllocator())
                 // if final frame then remove last four bytes 0x00 0x00 0xff 0xff
                 // (see  https://datatracker.ietf.org/doc/html/rfc7692#section-7.2.1)
                 if newFrame.fin {

diff --git a/Sources/HummingbirdWSCore/WebSocketContext.swift b/Sources/HummingbirdWSCore/WebSocketContext.swift
@@ -17,6 +17,5 @@ import NIOCore
 
 /// Protocol for WebSocket Data handling functions context parameter
 public protocol WebSocketContext: Sendable {
-    var allocator: ByteBufferAllocator { get }
     var logger: Logger { get }
 }
diff --git a/Sources/HummingbirdWSCore/WebSocketExtension.swift b/Sources/HummingbirdWSCore/WebSocketExtension.swift
@@ -20,11 +20,9 @@ import NIOWebSocket
 
 /// Basic context implementation of ``WebSocketContext``.
 public struct WebSocketExtensionContext {
-    public let allocator: ByteBufferAllocator
     public let logger: Logger
 
-    init(allocator: ByteBufferAllocator, logger: Logger) {
-        self.allocator = allocator
+    init(logger: Logger) {
         self.logger = logger
     }
 }

diff --git a/Sources/HummingbirdWSCore/WebSocketHandler.swift b/Sources/HummingbirdWSCore/WebSocketHandler.swift
@@ -79,7 +79,6 @@ package actor WebSocketHandler {
     let type: WebSocketType
     let configuration: Configuration
     let logger: Logger
-    let allocator: ByteBufferAllocator
     var pingData: ByteBuffer
     var pingTime: ContinuousClock.Instant = .now
     var closeState: CloseState
@@ -94,7 +93,6 @@ package actor WebSocketHandler {
         self.type = type
         self.configuration = configuration
         self.logger = context.logger
-        self.allocator = context.allocator
         self.pingData = ByteBufferAllocator().buffer(capacity: Self.pingDataSize)
         self.closeState = .open
     }
@@ -208,7 +206,7 @@ package actor WebSocketHandler {
             for ext in self.configuration.extensions {
                 frame = try await ext.processFrameToSend(
                     frame,
-                    context: WebSocketExtensionContext(allocator: self.allocator, logger: self.logger)
+                    context: WebSocketExtensionContext(logger: self.logger)
                 )
             }
         } catch {
@@ -281,7 +279,7 @@ package actor WebSocketHandler {
     ) async throws {
         switch self.closeState {
         case .open:
-            var buffer = self.allocator.buffer(capacity: 2 + (reason?.utf8.count ?? 0))
+            var buffer = ByteBufferAllocator().buffer(capacity: 2 + (reason?.utf8.count ?? 0))
             buffer.write(webSocketErrorCode: code)
             if let reason {
                 buffer.writeString(reason)
@@ -326,7 +324,7 @@ package actor WebSocketHandler {
                 .protocolError
             }
 
-            var buffer = self.allocator.buffer(capacity: 2)
+            var buffer = ByteBufferAllocator().buffer(capacity: 2)
             buffer.write(webSocketErrorCode: code)
 
             try await self.write(frame: .init(fin: true, opcode: .connectionClose, data: buffer))

diff --git a/Sources/HummingbirdWSCore/WebSocketInboundStream.swift b/Sources/HummingbirdWSCore/WebSocketInboundStream.swift
@@ -72,7 +72,7 @@ public final class WebSocketInboundStream: AsyncSequence, Sendable {
                         for ext in self.handler.configuration.extensions.reversed() {
                             frame = try await ext.processReceivedFrame(
                                 frame,
-                                context: WebSocketExtensionContext(allocator: self.handler.allocator, logger: self.handler.logger)
+                                context: WebSocketExtensionContext(logger: self.handler.logger)
                             )
                         }
                         return .init(from: frame)

diff --git a/Sources/HummingbirdWSCore/WebSocketOutboundWriter.swift b/Sources/HummingbirdWSCore/WebSocketOutboundWriter.swift
@@ -40,7 +40,7 @@ public struct WebSocketOutboundWriter: Sendable {
             try await self.handler.write(frame: .init(fin: true, opcode: .binary, data: buffer))
         case .text(let string):
             // send text based data
-            let buffer = self.handler.allocator.buffer(string: string)
+            let buffer = ByteBuffer(string: string)
             try await self.handler.write(frame: .init(fin: true, opcode: .text, data: buffer))
         case .pong:
             // send unexplained pong as a heartbeat
@@ -73,7 +73,7 @@ public struct WebSocketOutboundWriter: Sendable {
 
         /// Write string to WebSocket frame
         public mutating func callAsFunction(_ text: String) async throws {
-            let buffer = self.handler.allocator.buffer(string: text)
+            let buffer = ByteBuffer(string: text)
             try await self.write(buffer, opcode: self.opcode)
         }
 

diff --git a/Sources/HummingbirdWebSocket/WebSocketChannel.swift b/Sources/HummingbirdWebSocket/WebSocketChannel.swift
@@ -42,11 +42,9 @@ public struct HTTP1WebSocketUpgradeChannel: ServerChildChannel, HTTPChannelHandl
     /// Basic context implementation of ``WebSocketContext``.
     /// Used by non-router web socket handle function
     public struct Context: WebSocketContext {
-        public let allocator: ByteBufferAllocator
         public let logger: Logger
 
-        package init(allocator: ByteBufferAllocator, logger: Logger) {
-            self.allocator = allocator
+        package init(logger: Logger) {
             self.logger = logger
         }
     }
@@ -77,7 +75,7 @@ public struct HTTP1WebSocketUpgradeChannel: ServerChildChannel, HTTPChannelHandl
                             logger: logger
                         )
                         return (headers, { asyncChannel, logger in
-                            let context = Context(allocator: channel.allocator, logger: logger)
+                            let context = Context(logger: logger)
                             do {
                                 _ = try await WebSocketHandler.handle(
                                     type: .server,
@@ -126,7 +124,7 @@ public struct HTTP1WebSocketUpgradeChannel: ServerChildChannel, HTTPChannelHandl
                             logger: logger
                         )
                         return (headers, { asyncChannel, logger in
-                            let context = Context(allocator: channel.allocator, logger: logger)
+                            let context = Context(logger: logger)
                             do {
                                 _ = try await WebSocketHandler.handle(
                                     type: .server,

diff --git a/Sources/HummingbirdWebSocket/WebSocketRouter.swift b/Sources/HummingbirdWebSocket/WebSocketRouter.swift
@@ -37,9 +37,6 @@ public struct WebSocketRouterContext<Context: WebSocketContext>: WebSocketContex
     /// Logger attached to request context
     @inlinable
     public var logger: Logger { self.requestContext.logger }
-    /// ByteBuffer allocator attached to request context
-    @inlinable
-    public var allocator: ByteBufferAllocator { self.requestContext.allocator }
 }
 
 /// Reference to a WebSocket handler

diff --git a/Tests/HummingbirdWebSocketTests/WebSocketExtensionTests.swift b/Tests/HummingbirdWebSocketTests/WebSocketExtensionTests.swift
@@ -300,7 +300,7 @@ struct XorWebSocketExtension: WebSocketExtension {
     func shutdown() {}
 
     func xorFrame(_ frame: WebSocketFrame, context: WebSocketExtensionContext) -> WebSocketFrame {
-        var newBuffer = context.allocator.buffer(capacity: frame.data.readableBytes)
+        var newBuffer = ByteBufferAllocator().buffer(capacity: frame.data.readableBytes)
         for byte in frame.unmaskedData.readableBytesView {
             newBuffer.writeInteger(byte ^ self.value)
         }