allocator

Author: tomerd

Sources/AWSLambdaRuntimeCore/LambdaContext.swift

@@ -33,9 +33,13 @@ extension Lambda {
         ///         Most importantly the `EventLoop` must never be blocked.
         public let eventLoop: EventLoop
 
-        internal init(logger: Logger, eventLoop: EventLoop) {
+        /// `ByteBufferAllocator` to allocate `ByteBuffer`
+        public let allocator: ByteBufferAllocator
+
+        internal init(logger: Logger, eventLoop: EventLoop, allocator: ByteBufferAllocator) {
             self.eventLoop = eventLoop
             self.logger = logger
+            self.allocator = allocator
         }
     }
 }
@@ -87,7 +91,8 @@ extension Lambda {
                       cognitoIdentity: String? = nil,
                       clientContext: String? = nil,
                       logger: Logger,
-                      eventLoop: EventLoop) {
+                      eventLoop: EventLoop,
+                      allocator: ByteBufferAllocator) {
             self.requestID = requestID
             self.traceID = traceID
             self.invokedFunctionARN = invokedFunctionARN
@@ -96,7 +101,7 @@ extension Lambda {
             self.deadline = deadline
             // utility
             self.eventLoop = eventLoop
-            self.allocator = ByteBufferAllocator()
+            self.allocator = allocator
             // mutate logger with context
             var logger = logger
             logger[metadataKey: "awsRequestID"] = .string(requestID)

Sources/AWSLambdaRuntimeCore/LambdaRunner.swift

@@ -21,12 +21,14 @@ extension Lambda {
     internal final class Runner {
         private let runtimeClient: RuntimeClient
         private let eventLoop: EventLoop
+        private let allocator: ByteBufferAllocator
 
         private var isGettingNextInvocation = false
 
         init(eventLoop: EventLoop, configuration: Configuration) {
             self.eventLoop = eventLoop
             self.runtimeClient = RuntimeClient(eventLoop: self.eventLoop, configuration: configuration.runtimeEngine)
+            self.allocator = ByteBufferAllocator()
         }
 
         /// Run the user provided initializer. This *must* only be called once.
@@ -36,7 +38,9 @@ extension Lambda {
             logger.debug("initializing lambda")
             // 1. create the handler from the factory
             // 2. report initialization error if one occured
-            let context = InitializationContext(logger: logger, eventLoop: self.eventLoop)
+            let context = InitializationContext(logger: logger,
+                                                eventLoop: self.eventLoop,
+                                                allocator: self.allocator)
             return factory(context)
                 // Hopping back to "our" EventLoop is importnant in case the factory returns a future
                 // that originated from a foreign EventLoop/EventLoopGroup.
@@ -61,7 +65,10 @@ extension Lambda {
             }.flatMap { invocation, event in
                 // 2. send invocation to handler
                 self.isGettingNextInvocation = false
-                let context = Context(logger: logger, eventLoop: self.eventLoop, invocation: invocation)
+                let context = Context(logger: logger,
+                                      eventLoop: self.eventLoop,
+                                      allocator: self.allocator,
+                                      invocation: invocation)
                 logger.debug("sending invocation to lambda handler \(handler)")
                 return handler.handle(context: context, event: event)
                     // Hopping back to "our" EventLoop is importnant in case the handler returns a future that
@@ -94,15 +101,16 @@ extension Lambda {
 }
 
 private extension Lambda.Context {
-    convenience init(logger: Logger, eventLoop: EventLoop, invocation: Lambda.Invocation) {
+    convenience init(logger: Logger, eventLoop: EventLoop, allocator: ByteBufferAllocator, invocation: Lambda.Invocation) {
         self.init(requestID: invocation.requestID,
                   traceID: invocation.traceID,
                   invokedFunctionARN: invocation.invokedFunctionARN,
                   deadline: DispatchWallTime(millisSinceEpoch: invocation.deadlineInMillisSinceEpoch),
                   cognitoIdentity: invocation.cognitoIdentity,
                   clientContext: invocation.clientContext,
                   logger: logger,
-                  eventLoop: eventLoop)
+                  eventLoop: eventLoop,
+                  allocator: allocator)
     }
 }
 

Sources/AWSLambdaTesting/Lambda+Testing.swift

@@ -102,7 +102,8 @@ extension Lambda {
                               invokedFunctionARN: config.invokedFunctionARN,
                               deadline: .now() + config.timeout,
                               logger: logger,
-                              eventLoop: eventLoop)
+                              eventLoop: eventLoop,
+                              allocator: ByteBufferAllocator())
 
         return try eventLoop.flatSubmit {
             handler.handle(context: context, event: event)

Tests/AWSLambdaRuntimeCoreTests/LambdaTest.swift

@@ -263,7 +263,8 @@ class LambdaTest: XCTestCase {
                                      cognitoIdentity: nil,
                                      clientContext: nil,
                                      logger: Logger(label: "test"),
-                                     eventLoop: MultiThreadedEventLoopGroup(numberOfThreads: 1).next())
+                                     eventLoop: MultiThreadedEventLoopGroup(numberOfThreads: 1).next(),
+                                     allocator: ByteBufferAllocator())
         XCTAssertGreaterThan(context.deadline, .now())
 
         let expiredContext = Lambda.Context(requestID: context.requestID,
@@ -273,7 +274,8 @@ class LambdaTest: XCTestCase {
                                             cognitoIdentity: context.cognitoIdentity,
                                             clientContext: context.clientContext,
                                             logger: context.logger,
-                                            eventLoop: context.eventLoop)
+                                            eventLoop: context.eventLoop,
+                                            allocator: context.allocator)
         XCTAssertLessThan(expiredContext.deadline, .now())
     }
 
@@ -285,7 +287,8 @@ class LambdaTest: XCTestCase {
                                      cognitoIdentity: nil,
                                      clientContext: nil,
                                      logger: Logger(label: "test"),
-                                     eventLoop: MultiThreadedEventLoopGroup(numberOfThreads: 1).next())
+                                     eventLoop: MultiThreadedEventLoopGroup(numberOfThreads: 1).next(),
+                                     allocator: ByteBufferAllocator())
         XCTAssertLessThanOrEqual(context.getRemainingTime(), .seconds(1))
         XCTAssertGreaterThan(context.getRemainingTime(), .milliseconds(800))
     }

Tests/AWSLambdaRuntimeTests/Lambda+CodeableTest.swift

@@ -72,7 +72,8 @@ class CodableLambdaTest: XCTestCase {
                        cognitoIdentity: nil,
                        clientContext: nil,
                        logger: Logger(label: "test"),
-                       eventLoop: self.eventLoopGroup.next())
+                       eventLoop: self.eventLoopGroup.next(),
+                       allocator: ByteBufferAllocator())
     }
 }