Compatibility w/ AWS Fargate

.gitignore

@@ -1,6 +1,6 @@
 /pkg
 /build
-/bin
+# /bin
 *.swp
 *.iml
 tags

README.md

@@ -3,6 +3,8 @@
 ![GitHub go.mod Go version](https://img.shields.io/github/go-mod/go-version/aws/aws-lambda-runtime-interface-emulator)
 ![GitHub](https://img.shields.io/github/license/aws/aws-lambda-runtime-interface-emulator)
 
+## Patches
+-  [Fix already reserved panic on concurrent invokes](https://github.com/aws/aws-lambda-runtime-interface-emulator/pull/133)
 
 The Lambda Runtime Interface Emulator is a proxy for Lambda’s Runtime and Extensions APIs, which allows customers to
 locally test their Lambda function packaged as a container image. It is a lightweight web-server that converts
@@ -171,6 +173,11 @@ The rest of these Environment Variables can be set to match AWS Lambda's environ
 * `AWS_LAMBDA_FUNCTION_VERSION`
 * `AWS_LAMBDA_FUNCTION_NAME`
 * `AWS_LAMBDA_FUNCTION_MEMORY_SIZE`
+* `AWS_LAMBDA_SERVER_MAX_INVOCATIONS`
+* `API_ACCESS_KEY`
+
+Request headers:
+* `forward-response` to forward api response to server
 
 ## Level of support
 

cmd/aws-lambda-rie/handlers.go

@@ -72,7 +72,7 @@ func printEndReports(invokeId string, initDuration string, memorySize string, in
 		invokeId, invokeDuration, math.Ceil(invokeDuration), memorySize, memorySize)
 }
 
-func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs interop.Bootstrap) {
+func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs interop.Bootstrap, doneCallback func(invokeResp *ResponseWriterProxy)) {
 	log.Debugf("invoke: -> %s %s %v", r.Method, r.URL, r.Header)
 	bodyBytes, err := ioutil.ReadAll(r.Body)
 	if err != nil {
@@ -81,6 +81,18 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		return
 	}
 
+	// Fix CORS error: https://github.com/aws/aws-lambda-runtime-interface-emulator/pull/84/files#diff-2c1a36d379bd5ed6e893749912ac8473cc46dddf5765024d46b44da2eb56348d
+	if origin := r.Header.Get("Origin"); origin != "" {
+		w.Header().Set("Access-Control-Allow-Origin", GetenvWithDefault("ACCESS_CONTROL_ALLOW_ORIGIN", origin))
+		w.Header().Set("Access-Control-Allow-Methods", GetenvWithDefault("ACCESS_CONTROL_ALLOW_METHODS", "POST, OPTIONS"))
+		w.Header().Set("Access-Control-Allow-Headers", GetenvWithDefault("ACCESS_CONTROL_ALLOW_HEADERS", "*"))
+		w.Header().Set("Access-Control-Allow-Credentials", GetenvWithDefault("ACCESS_CONTROL_ALLOW_CREDENTIALS", "true"))
+	}
+	if r.Method == "OPTIONS" {
+		w.WriteHeader(200)
+		return
+	}
+
 	initDuration := ""
 	inv := GetenvWithDefault("AWS_LAMBDA_FUNCTION_TIMEOUT", "300")
 	timeoutDuration, _ := time.ParseDuration(inv + "s")
@@ -165,9 +177,11 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		case rapidcore.ErrInvokeTimeout:
 			printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)
 
-			w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
+			invokeResp.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
+			w.Write(invokeResp.Body)	
 			time.Sleep(100 * time.Millisecond)
 			//initDone = false
+			doneCallback(invokeResp) // Call done after printEndReports
 			return
 		}
 	}
@@ -178,6 +192,8 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		w.WriteHeader(invokeResp.StatusCode)
 	}
 	w.Write(invokeResp.Body)
+
+	doneCallback(invokeResp) // Call done after printEndReports
 }
 
 func InitHandler(sandbox Sandbox, functionVersion string, timeout int64, bs interop.Bootstrap) (time.Time, time.Time) {

cmd/aws-lambda-rie/http.go

@@ -4,7 +4,12 @@
 package main
 
 import (
+	"os"
+	"sync"
+	"bytes"
+	"strconv"
 	"net/http"
+	"encoding/json"
 
 	log "github.com/sirupsen/logrus"
 	"go.amzn.com/lambda/interop"
@@ -16,15 +21,85 @@ func startHTTPServer(ipport string, sandbox *rapidcore.SandboxBuilder, bs intero
 		Addr: ipport,
 	}
 
+	log.Warnf("Listening on %s", ipport)
+
+	maxInvocations := -1 // -1 means unlimited invocations
+	// Get max invocations from environment variable
+	maxInvocationsStr := os.Getenv("AWS_LAMBDA_SERVER_MAX_INVOCATIONS")
+	if maxInvocationsStr != "" {
+		if maxInvocationsInt, err := strconv.Atoi(maxInvocationsStr); err == nil {
+			maxInvocations = maxInvocationsInt
+		} else {
+			log.Panicf("Invalid value for AWS_LAMBDA_SERVER_MAX_INVOCATIONS: %s", maxInvocationsStr)
+		}
+	}
+
+	// Channel to signal server shutdown
+	shutdownChan := make(chan struct{})
+
 	// Pass a channel
 	http.HandleFunc("/2015-03-31/functions/function/invocations", func(w http.ResponseWriter, r *http.Request) {
-		InvokeHandler(w, r, sandbox.LambdaInvokeAPI(), bs)
+		InvokeHandler(w, r, sandbox.LambdaInvokeAPI(), bs, func(invokeResp *ResponseWriterProxy){
+
+			// Forward response if "forward-response" header exists
+			if forwardURL := r.Header.Get("forward-response"); forwardURL != "" {
+				// Create a wait group to wait for the API request to finish
+				var wg sync.WaitGroup
+				wg.Add(1)
+				go func() {
+					// Marshal the payload to JSON (you can use any other serialization format if needed)
+					apiPayloadJSON, err := json.Marshal(string(invokeResp.Body))
+					if err != nil {
+						log.Errorf("Failed to json marshal API payload: %s", err)
+						return
+					}
+
+					// Create an API request to the URL in the "forward-response" header
+					client := &http.Client{}
+					req, err := http.NewRequest("POST", forwardURL, bytes.NewReader(apiPayloadJSON))
+					// Add request headers
+					req.Header.Add("Authorization", "Token "+os.Getenv("API_ACCESS_KEY"))
+					req.Header.Add("Content-Type", "application/json")
+					// Send the request
+					resp, err := client.Do(req)
+					if err != nil {
+						log.Errorf("Failed to forward response: %s", err)
+						return
+					}
+					defer resp.Body.Close()
+
+					if resp.StatusCode == 200 {
+						log.Printf("Forwarded response was successful")
+					} else {
+						log.Errorf("Forwarding response failed with status code: %d", resp.StatusCode)
+					}
+
+					defer wg.Done() // Defer the Done() call to mark the API request as completed
+				}()
+
+				wg.Wait() // Wait for the API request to finish before proceeding
+			}
+
+			// Shutdown the server if the maximum number of invocations is reached
+			maxInvocations--
+			if maxInvocations == 0 {
+				close(shutdownChan)
+			}
+		})
 	})
 
-	// go routine (main thread waits)
-	if err := srv.ListenAndServe(); err != nil {
+
+	// go routine to handle server shutdown (main thread waits)
+	go func() {
+		<-shutdownChan
+		log.Printf("Maximum invocations (%s) reached. Shutting down the server", maxInvocationsStr)
+		if err := srv.Shutdown(nil); err != nil {
+			log.Panic(err)
+		}
+	}()
+
+	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
 		log.Panic(err)
 	}
 
-	log.Warnf("Listening on %s", ipport)
 }

lambda/rapidcore/server.go

@@ -1,6 +1,12 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+/**
+ * Required patches:
+ * - Fix already reserved panic on concurrent invokes
+ *		https://github.com/aws/aws-lambda-runtime-interface-emulator/pull/133/files
+ */
+
 package rapidcore
 
 import (
@@ -114,59 +120,35 @@ type Server struct {
 var _ interop.Server = (*Server)(nil)
 
 func (s *Server) setRapidPhase(phase rapidPhase) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	s.rapidPhase = phase
 }
 
 func (s *Server) getRapidPhase() rapidPhase {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	return s.rapidPhase
 }
 
 func (s *Server) setRuntimeState(state runtimeState) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	s.runtimeState = state
 }
 
 func (s *Server) getRuntimeState() runtimeState {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	return s.runtimeState
 }
 
 func (s *Server) SetInvokeTimeout(timeout time.Duration) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	s.invokeTimeout = timeout
 }
 
 func (s *Server) GetInvokeTimeout() time.Duration {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	return s.invokeTimeout
 }
 
 func (s *Server) GetInvokeContext() *InvokeContext {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	ctx := *s.invokeCtx
 	return &ctx
 }
 
 func (s *Server) setNewInvokeContext(invokeID string, traceID, lambdaSegmentID string) (*ReserveResponse, error) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	if s.invokeCtx != nil {
 		return nil, ErrAlreadyReserved
 	}
@@ -233,9 +215,6 @@ func (s *Server) awaitInitCompletion() {
 }
 
 func (s *Server) setReplyStream(w http.ResponseWriter, direct bool) (string, error) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	if s.invokeCtx == nil {
 		return "", ErrNotReserved
 	}
@@ -255,9 +234,6 @@ func (s *Server) setReplyStream(w http.ResponseWriter, direct bool) (string, err
 
 // Release closes the invocation, making server ready for reserve again
 func (s *Server) Release() error {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	if s.invokeCtx == nil {
 		return ErrNotReserved
 	}
@@ -274,9 +250,6 @@ func (s *Server) Release() error {
 
 // GetCurrentInvokeID
 func (s *Server) GetCurrentInvokeID() string {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
 	if s.invokeCtx == nil {
 		return ""
 	}
@@ -357,8 +330,6 @@ func (s *Server) sendResponseUnsafe(invokeID string, additionalHeaders map[strin
 
 func (s *Server) SendResponse(invokeID string, headers map[string]string, reader io.Reader, trailers http.Header, request *interop.CancellableRequest) error {
 	s.setRuntimeState(runtimeInvokeResponseSent)
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
 	runtimeCalledResponse := true
 	return s.sendResponseUnsafe(invokeID, headers, http.StatusOK, reader, trailers, request, runtimeCalledResponse)
 }
@@ -379,8 +350,6 @@ func (s *Server) SendInitErrorResponse(invokeID string, resp *interop.ErrorRespo
 func (s *Server) SendErrorResponse(invokeID string, resp *interop.ErrorResponse) error {
 	log.Debugf("Sending Error Response: %s", resp.ErrorType)
 	s.setRuntimeState(runtimeInvokeError)
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
 	additionalHeaders := map[string]string{contentTypeHeader: resp.ContentType, errorTypeHeader: resp.ErrorType}
 	if functionResponseMode := resp.FunctionResponseMode; functionResponseMode != "" {
 		additionalHeaders[functionResponseModeHeader] = functionResponseMode
@@ -485,14 +454,10 @@ func deadlineNsFromTimeoutMs(timeoutMs int64) int64 {
 }
 
 func (s *Server) setInitFailuresChan() {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
 	s.initFailures = make(chan interop.InitFailure)
 }
 
 func (s *Server) getInitFailuresChan() chan interop.InitFailure {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
 	return s.initFailures
 }
 
@@ -580,14 +545,10 @@ func (s *Server) FastInvoke(w http.ResponseWriter, i *interop.Invoke, direct boo
 }
 
 func (s *Server) setCachedInitErrorResponse(errResp *interop.ErrorResponse) {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
 	s.cachedInitErrorResponse = errResp
 }
 
 func (s *Server) getCachedInitErrorResponse() *interop.ErrorResponse {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
 	return s.cachedInitErrorResponse
 }
 
@@ -600,8 +561,6 @@ func (s *Server) trySendDefaultErrorResponse(resp *interop.ErrorResponse) {
 }
 
 func (s *Server) CurrentToken() *interop.Token {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
 	if s.invokeCtx == nil {
 		return nil
 	}
@@ -612,6 +571,9 @@ func (s *Server) CurrentToken() *interop.Token {
 // Invoke is used by the Runtime Interface Emulator (Rapid Local)
 // https://github.com/aws/aws-lambda-runtime-interface-emulator
 func (s *Server) Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
 	resetCtx, resetCancel := context.WithCancel(context.Background())
 	defer resetCancel()
 
@@ -749,7 +711,7 @@ func (s *Server) awaitInitialized() (initCompletionResponse, error) {
 // since it can be called twice when a caller wants to wait until init is complete
 func (s *Server) AwaitInitialized() error {
 	if _, err := s.awaitInitialized(); err != nil {
-		if releaseErr := s.Release(); err != nil {
+		if releaseErr := s.Release(); releaseErr != nil {
 			log.Infof("Error releasing after init failure %s: %s", err, releaseErr)
 		}
 		s.setRuntimeState(runtimeInitFailed)