Issue 1639 - LambdaTestTool - Sync Invoke Endpoint

Tools/LambdaTestTool/src/Amazon.Lambda.TestTool.BlazorTester/Controllers/RuntimeApiController.cs

@@ -1,4 +1,6 @@
 using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.IO;
 using System.Text;
 using System.Threading.Tasks;
@@ -13,13 +15,14 @@ public class RuntimeApiController : ControllerBase
         private const string HEADER_BREAK = "-----------------------------------";
         private readonly IRuntimeApiDataStore _runtimeApiDataStore;
 
+        private static readonly ConcurrentDictionary<string, string> _registeredExtensions = new ();
+
         public RuntimeApiController(IRuntimeApiDataStore runtimeApiDataStore)
         {
             _runtimeApiDataStore = runtimeApiDataStore;
         }
 
         [HttpPost("/runtime/test-event")]
-        [HttpPost("/2015-03-31/functions/function/invocations")]
         public async Task<IActionResult> PostTestEvent()
         {
             using var reader = new StreamReader(Request.Body);
@@ -28,6 +31,101 @@ public async Task<IActionResult> PostTestEvent()
 
             return Accepted();
         }
+
+        [HttpPost("/2020-01-01/extension/register")]
+        public async Task<IActionResult> PostRegisterExtension()
+        {
+            using var reader = new StreamReader(Request.Body);
+            // Read and discard - do not need the body
+            await reader.ReadToEndAsync();
+
+            var extensionId = Guid.NewGuid().ToString();
+            var extensionName = Request.Headers["Lambda-Extension-Name"];
+            _registeredExtensions.TryAdd(extensionId, extensionName);
+
+            Response.Headers["Lambda-Extension-Identifier"] = extensionId;
+            return Ok();
+        }
+
+        [HttpGet("/2020-01-01/extension/event/next")]
+        public async Task<IActionResult> GetNextExtensionEvent()
+        {
+            var extensionId = Request.Headers["Lambda-Extension-Identifier"];
+            if (_registeredExtensions.ContainsKey(extensionId)) {
+                await Task.Delay(TimeSpan.FromSeconds(15));
+                Console.WriteLine(HEADER_BREAK);
+                Console.WriteLine($"Extension ID: {extensionId} - Returning NoContent");
+                return NoContent();
+            } else {
+                Console.WriteLine(HEADER_BREAK);
+                Console.Error.WriteLine($"Extension ID: {extensionId} - Was not found - Returning 404 - ");
+                return NotFound();
+            }
+        }
+
+        [HttpPost("/2015-03-31/functions/{functionName}/invocations")]
+        public async Task<IActionResult> PostTestInvokeEvent(string functionName)
+        {
+            using var reader = new StreamReader(Request.Body);
+            var testEvent = await reader.ReadToEndAsync();
+            var eventContainer = _runtimeApiDataStore.QueueEvent(testEvent);
+
+            // Need a task completion source so we can block until the event is executed.
+            var tcs = new TaskCompletionSource();
+
+            eventContainer.OnSuccess += () =>
+            {
+                try {
+                    tcs.SetResult();
+                } catch (InvalidOperationException) {
+                    // This can happen if both OnSuccess and OnError are called
+                    // Oddly, this does happen 1 time in 50 million requests
+                    // We can't check a variable because it's a race condition
+                }
+            };
+
+            eventContainer.OnError += () =>
+            {
+                try {
+                    tcs.SetResult();
+                } catch (InvalidOperationException) {
+                    // See note above
+                }
+            };
+
+            // Wait for our event to process
+            await tcs.Task;
+
+            if (eventContainer.ErrorResponse != null)
+            {
+                if (eventContainer.ErrorType == "Throttled")
+                {
+                    return Ok(new {
+                        StatusCode = 429,
+                        FunctionError = "Throttled",
+                        ExecutedVersion = "$LATEST",
+                        Payload = "{\"errorMessage\":\"Rate Exceeded.\"}"
+                    });
+                }
+                return Ok(new {
+                    StatusCode = 200,
+                    FunctionError = "Unhandled",
+                    ExecutedVersion = "$LATEST",
+                    Payload = "{\"errorMessage\":\"An error occurred.\",\"errorType\":\"Error\"}"
+                });
+            }
+
+            var response = new 
+            {
+                StatusCode = 200, // Accepted
+                // FunctionError = null, // TODO: Set this if there was an error
+                // LogResult = null, // TODO: Set this to the base64-encoded last 4 KB of log data produced by the function
+                Payload = eventContainer.Response, // Set this to the response from the function
+                // ExecutedVersion = null // TODO: Set this to the version of the function that was executed
+            };
+
+            return Ok(response);
+        }
 
         [HttpPost("/2018-06-01/runtime/init/error")]
         public IActionResult PostInitError([FromHeader(Name = "Lambda-Runtime-Function-Error-Type")] string errorType, [FromBody] string error)
@@ -50,6 +148,8 @@ public async Task GetNextInvocation()
             Console.WriteLine(HEADER_BREAK);
             Console.WriteLine($"Next invocation returned: {activeEvent.AwsRequestId}");
 
+            // Set Deadline to 5 minutes from now in unix epoch ms
+            Response.Headers["Lambda-Runtime-Deadline-Ms"] = DateTimeOffset.UtcNow.AddMinutes(5).ToUnixTimeMilliseconds().ToString();
             Response.Headers["Lambda-Runtime-Aws-Request-Id"] = activeEvent.AwsRequestId;
             Response.Headers["Lambda-Runtime-Trace-Id"] = Guid.NewGuid().ToString();
             Response.Headers["Lambda-Runtime-Invoked-Function-Arn"] = activeEvent.FunctionArn;

Tools/LambdaTestTool/src/Amazon.Lambda.TestTool.BlazorTester/Pages/Runtime.razor

@@ -13,7 +13,7 @@
 
 <p>
     For Lambda functions written as executable assemblies, i.e. custom runtimes functions and top level statement functions, this page can be used for testing the functions locally. 
-    Set the <b>AWS_LAMBDA_RUNTIME_API</b> environment variable to <b>@httpContextAccessor.HttpContext.Request.Host</b> while debugging executable assembly 
+    Set the <b>AWS_LAMBDA_RUNTIME_API</b> environment variable to <b>@httpContextAccessor?.HttpContext?.Request?.Host</b> while debugging executable assembly 
     Lambda function. More information can be found in the <a href="/documentation">documentation</a>.
 </p>
 
@@ -75,6 +75,13 @@
                               <pre class="form-control" style="@Constants.ResponseErrorStyleSizeConstraint" >@RuntimeApiModel.ActiveEvent.ErrorResponse</pre>
                           </p>
                       }
+                      else if (RuntimeApiModel.ActiveEvent.Response == null)
+                      {
+                            <p>
+                                <b>Response:</b>
+                                <pre class="form-control" style="@Constants.ResponseErrorStyleSizeConstraint" >No Response Received</pre>
+                            </p>
+                      }
                       else
                       {
                           <p>

Tools/LambdaTestTool/src/Amazon.Lambda.TestTool.BlazorTester/Services/RuntimeApiDataStore.cs

@@ -3,6 +3,7 @@
 using System.Collections.ObjectModel;
 using System.Linq;
 using System.Threading;
+using System.Threading.Tasks;
 
 namespace Amazon.Lambda.TestTool.BlazorTester.Services
 {
@@ -47,6 +48,23 @@ public EventContainer QueueEvent(string eventBody)
             {
                 _queuedEvents.Add(evnt);
             }
+
+            // Start a task that will cancel the event after a timeout or dispatch
+            Task.Run(async () => {
+                // Wait for the event to be dispatched or timeout
+                try {
+                    await evnt.DispatchedTCS.Task.WaitAsync(evnt.TimedOutCTS.Token);
+                } catch (TaskCanceledException) {
+                    // If the event was cancelled then it timed out
+                    lock(_lock) {
+                        _queuedEvents.Remove(evnt);
+                    }
+                    evnt.Cancel("Lambda throttled response error");
+                    return;
+                }
+
+                // If the event was not cancelled then it was dispatched, and we're good
+            });
 
             RaiseStateChanged();
             return evnt;
@@ -57,22 +75,28 @@ public bool TryActivateEvent(out IEventContainer activeEvent)
             activeEvent = null;
             lock(_lock)
             {
-                if (!_queuedEvents.Any())
-                {
-                    return false;
-                }
+                while (true) {
+                    if (!_queuedEvents.Any())
+                    {
+                        return false;
+                    }
 
-                var evnt = _queuedEvents[0];
-                _queuedEvents.RemoveAt(0);
-                evnt.EventStatus = IEventContainer.Status.Executing;
-                if (ActiveEvent != null)
-                {
-                    _executedEvents.Add(ActiveEvent as EventContainer);
+                    var evnt = _queuedEvents[0];
+                    _queuedEvents.RemoveAt(0);
+                    if (evnt.MarkExecuting()) {
+                        if (ActiveEvent != null)
+                        {
+                            _executedEvents.Add(ActiveEvent as EventContainer);
+                        }
+                        ActiveEvent = evnt;
+                        activeEvent = ActiveEvent;
+                        RaiseStateChanged();
+                        return true;
+                    }
+
+                    // If we get here there was an event but it was already failed (timed out)
+                    // Loop around and check if there are more events
                 }
-                ActiveEvent = evnt;
-                activeEvent = ActiveEvent;
-                RaiseStateChanged();
-                return true;
             }
         }
 
@@ -213,7 +237,10 @@ public enum Status {Queued, Executing, Success, Failure}
 
     public class EventContainer : IEventContainer
     {
-
+        public Action OnSuccess { get; set; }
+
+        public Action OnError { get; set; }
+
         private const string defaultFunctionArn = "arn:aws:lambda:us-west-2:123412341234:function:Function";
         public string AwsRequestId { get; }
         public string EventJson { get; }
@@ -223,13 +250,32 @@ public class EventContainer : IEventContainer
 
         public string Response { get; private set; }
 
+        public bool MarkExecuting()
+        {
+            lock(_statusLock){
+                if (EventStatus == IEventContainer.Status.Queued)
+                {
+                    EventStatus = IEventContainer.Status.Executing;
+
+                    // Mark that the event has been dispatched
+                    DispatchedTCS.TrySetResult();
+                    return true;
+                }
+            }
+            return false;
+        }
+
         public DateTime LastUpdated { get; private set; }
 
+        public Task TimeoutTask { get; set; }
+        public TaskCompletionSource DispatchedTCS { get; private set; } = new ();
+        public CancellationTokenSource TimedOutCTS { get; private set; } = new (1000);
+        private readonly object _statusLock = new();
         private IEventContainer.Status _status = IEventContainer.Status.Queued;
         public IEventContainer.Status EventStatus
         {
             get => _status;
-            set
+            private set
             {
                 _status = value;
                 LastUpdated = DateTime.Now;
@@ -253,18 +299,38 @@ public string FunctionArn
         public void ReportSuccessResponse(string response)
         {
             LastUpdated = DateTime.Now;
-            this.Response = response;
-            this.EventStatus = IEventContainer.Status.Success;
-            _dataStore.RaiseStateChanged();
+            lock (_statusLock) {
+                if (this.EventStatus == IEventContainer.Status.Executing) {
+                    this.EventStatus = IEventContainer.Status.Success;
+                    this.Response = response;
+                    OnSuccess?.Invoke();
+                    _dataStore.RaiseStateChanged();
+                }
+            }
         }
 
         public void ReportErrorResponse(string errorType, string errorBody)
         {
             LastUpdated = DateTime.Now;
-            this.ErrorType = errorType;
-            this.ErrorResponse = errorBody;
-            this.EventStatus = IEventContainer.Status.Failure;
-            _dataStore.RaiseStateChanged();
+            lock(_statusLock) {
+                if (EventStatus == IEventContainer.Status.Queued || this.EventStatus == IEventContainer.Status.Executing) {
+                    this.ErrorType = errorType;
+                    this.ErrorResponse = errorBody;
+                    this.EventStatus = IEventContainer.Status.Failure;
+                    OnError?.Invoke();
+                    _dataStore.RaiseStateChanged();
+                }
+            }
+        }
+
+        public void Cancel(string errorBody)
+        {
+            lock(_statusLock) {
+                if (EventStatus == IEventContainer.Status.Queued)
+                {
+                    ReportErrorResponse("Throttled", errorBody);
+                }
+            }
         }
     }
 }

Tools/LambdaTestTool/src/Amazon.Lambda.TestTool.BlazorTester/Startup.cs

@@ -80,7 +80,8 @@ public void ConfigureServices(IServiceCollection services)
             services.AddControllers();
             services.AddRazorPages();
             services.AddServerSideBlazor()
-                    .AddHubOptions(options => options.MaximumReceiveMessageSize = null);
+                    .AddHubOptions(options => options.MaximumReceiveMessageSize = null)
+                    .AddHubOptions(options => options.ClientTimeoutInterval = TimeSpan.FromMinutes(10));
             services.AddHttpContextAccessor();
 
             services.AddBlazoredModal();

Tools/LambdaTestTool/tests/Amazon.Lambda.TestTool.BlazorTester.Tests/RuntimeApiControllerTests.cs

@@ -151,7 +151,8 @@ public static async Task<TestSession> CreateSessionAsync(string host, int port)
                 var uriString = Utils.DetermineLaunchUrl(session.Host, session.Port, Constants.DEFAULT_HOST);
                 session.Client = new HttpClient()
                 {
-                    BaseAddress = new Uri(uriString)
+                    BaseAddress = new Uri(uriString),
+                    Timeout = TimeSpan.FromMinutes(15)
                 };
 
                 session.Store = session.WebHost.Services.GetService<IRuntimeApiDataStore>();