Add invocation timeout handling (#6)

Co-authored-by: Dominik Schubert <[email protected]>

Makefile

@@ -1,5 +1,6 @@
 # LOCALSTACK CHANGES 2022-03-10: remove linker flags and add gc flags for delve debugger
 # LOCALSTACK CHANGES 2022-03-28: change compile src folder
+# LOCALSTACK CHANGES 2022-11-14: add --rm flag to compile-with-docker
 
 # RELEASE_BUILD_LINKER_FLAGS disables DWARF and symbol table generation to reduce binary size
 #RELEASE_BUILD_LINKER_FLAGS=-s -w
@@ -21,7 +22,7 @@ compile-lambda-linux-all:
 	make ARCH=arm64 compile-lambda-linux
 
 compile-with-docker:
-	docker run --env GOPROXY=direct -v $(shell pwd):/LambdaRuntimeLocal -w /LambdaRuntimeLocal golang:1.18 make ARCH=${ARCH} compile-lambda-linux
+	docker run --rm --env GOPROXY=direct -v $(shell pwd):/LambdaRuntimeLocal -w /LambdaRuntimeLocal golang:1.18 make ARCH=${ARCH} compile-lambda-linux
 
 compile-lambda-linux:
 	CGO_ENABLED=0 GOOS=linux GOARCH=${GO_ARCH_${ARCH}} go build -ldflags "${RELEASE_BUILD_LINKER_FLAGS}" -gcflags="all=-N -l" -o ${DESTINATION_${ARCH}} ./cmd/localstack

cmd/localstack/custom_interop.go

@@ -52,9 +52,9 @@ type InvokeRequest struct {
 
 type ErrorResponse struct {
 	ErrorMessage string   `json:"errorMessage"`
-	ErrorType    string   `json:"errorType"`
-	RequestId    string   `json:"requestId"`
-	StackTrace   []string `json:"stackTrace"`
+	ErrorType    string   `json:"errorType,omitempty"`
+	RequestId    string   `json:"requestId,omitempty"`
+	StackTrace   []string `json:"stackTrace,omitempty"`
 }
 
 func NewCustomInteropServer(lsOpts *LsOpts, delegate rapidcore.InteropServer, logCollector *LogCollector) (server *CustomInteropServer) {
@@ -99,12 +99,36 @@ func NewCustomInteropServer(lsOpts *LsOpts, delegate rapidcore.InteropServer, lo
 					CorrelationID:      "invokeCorrelationID",
 					NeedDebugLogs:      true,
 					InvokedFunctionArn: invokeR.InvokedFunctionArn,
+					//DeadlineNs:
 				})
+				timeout := int(server.delegate.GetInvokeTimeout().Seconds())
+				isErr := false
 				if err != nil {
-					log.Fatalln(err)
+					switch err {
+					case rapidcore.ErrInvokeTimeout:
+						log.Debugf("Got invoke timeout")
+						isErr = true
+						errorResponse := ErrorResponse{
+							ErrorMessage: fmt.Sprintf(
+								"%s %s Task timed out after %d.00 seconds",
+								time.Now().Format("2006-01-02T15:04:05Z"),
+								invokeR.InvokeId,
+								timeout,
+							),
+						}
+						jsonErrorResponse, err := json.Marshal(errorResponse)
+						if err != nil {
+							log.Fatalln("unable to marshall json timeout response")
+						}
+						_, err = invokeResp.Write(jsonErrorResponse)
+						if err != nil {
+							log.Fatalln("unable to write to response")
+						}
+					default:
+						log.Fatalln(err)
+					}
 				}
-				inv := GetEnvOrDie("AWS_LAMBDA_FUNCTION_TIMEOUT")
-				timeoutDuration, _ := time.ParseDuration(inv + "s")
+				timeoutDuration := time.Duration(timeout) * time.Second
 				memorySize := GetEnvOrDie("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
 				PrintEndReports(invokeR.InvokeId, "", memorySize, invokeStart, timeoutDuration, logCollector)
 
@@ -117,8 +141,7 @@ func NewCustomInteropServer(lsOpts *LsOpts, delegate rapidcore.InteropServer, lo
 				var errR map[string]any
 				marshalErr := json.Unmarshal(invokeResp.Body, &errR)
 
-				isErr := false
-				if marshalErr == nil {
+				if !isErr && marshalErr == nil {
 					_, isErr = errR["errorType"]
 				}
 

cmd/localstack/main.go

@@ -9,6 +9,7 @@ import (
 	_ "net/http/pprof"
 	"os"
 	"runtime/debug"
+	"strconv"
 )
 
 type LsOpts struct {
@@ -72,12 +73,18 @@ func main() {
 	// initialize all flows and start runtime API
 	go sandbox.Create()
 
+	// get timeout
+	invokeTimeoutEnv := GetEnvOrDie("AWS_LAMBDA_FUNCTION_TIMEOUT")
+	invokeTimeoutSeconds, err := strconv.Atoi(invokeTimeoutEnv)
+	if err != nil {
+		log.Fatalln(err)
+	}
 	// start runtime init
-	go InitHandler(sandbox, GetEnvOrDie("AWS_LAMBDA_FUNCTION_VERSION"), 30) // TODO: replace this with a custom init
+	go InitHandler(sandbox, GetEnvOrDie("AWS_LAMBDA_FUNCTION_VERSION"), int64(invokeTimeoutSeconds)) // TODO: replace this with a custom init
 
 	// TODO: make the tracing server optional
 	// start blocking with the tracing server
-	err := http.ListenAndServe("0.0.0.0:"+lsOpts.InitTracingPort, http.DefaultServeMux)
+	err = http.ListenAndServe("0.0.0.0:"+lsOpts.InitTracingPort, http.DefaultServeMux)
 	if err != nil {
 		log.Fatal("Failed to start debug server")
 	}