From 845894ba64ba4a6576fa4c711e532d9edf6973da Mon Sep 17 00:00:00 2001
From: martincostello <martin@martincostello.com>
Date: Mon, 7 Oct 2024 10:36:27 +0100
Subject: [PATCH] Remove runtime responses once written

Fix memory leak from build up of `ResponseContext` instances by removing them from the dictionary once used to write the `LambdaTestResponse` to the associated channel.
---
 src/AwsLambdaTestServer/RuntimeHandler.cs | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/AwsLambdaTestServer/RuntimeHandler.cs b/src/AwsLambdaTestServer/RuntimeHandler.cs
index 20fb1d96..1a13e3e7 100644
--- a/src/AwsLambdaTestServer/RuntimeHandler.cs
+++ b/src/AwsLambdaTestServer/RuntimeHandler.cs
@@ -330,7 +330,11 @@ private async Task CompleteRequestChannelAsync(
         bool isSuccessful,
         CancellationToken cancellationToken)
     {
-        var context = _responses[awsRequestId];
+        if (!_responses.TryRemove(awsRequestId, out var context))
+        {
+            throw new InvalidOperationException($"Failed to complete channel for AWS request Id {awsRequestId}.");
+        }
+
         context.DurationTimer!.Stop();
 
         Logger?.LogInformation(