diff --git a/internal/lambda/rapidcore/sandbox_emulator_api.go b/internal/lambda/rapidcore/sandbox_emulator_api.go index d5c1d47b..a6efea1c 100644 --- a/internal/lambda/rapidcore/sandbox_emulator_api.go +++ b/internal/lambda/rapidcore/sandbox_emulator_api.go @@ -7,11 +7,13 @@ import ( "github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop" "net/http" + "time" ) // LambdaInvokeAPI are the methods used by the Runtime Interface Emulator type LambdaInvokeAPI interface { Init(i *interop.Init, invokeTimeoutMs int64) + AwaitInitCompletion() time.Time Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error } @@ -47,6 +49,10 @@ func (l *EmulatorAPI) Init(i *interop.Init, timeoutMs int64) { }, timeoutMs) } +func (l *EmulatorAPI) AwaitInitCompletion() time.Time { + return l.server.AwaitInitCompletion() +} + // Invoke method is only used by the Runtime interface emulator func (l *EmulatorAPI) Invoke(w http.ResponseWriter, i *interop.Invoke) error { return l.server.Invoke(w, i) diff --git a/internal/lambda/rapidcore/server.go b/internal/lambda/rapidcore/server.go index b82d8276..f3c159ce 100644 --- a/internal/lambda/rapidcore/server.go +++ b/internal/lambda/rapidcore/server.go @@ -72,6 +72,11 @@ type InvokeContext struct { Direct bool } +type initCompletion struct { + done chan struct{} + completedAt time.Time +} + type Server struct { InternalStateGetter interop.InternalStateGetter @@ -100,6 +105,7 @@ type Server struct { initContext interop.InitContext invoker interop.InvokeContext initFailures chan interop.InitFailure + initCompletion *initCompletion cachedInitErrorResponse *interop.ErrorInvokeResponse } @@ -211,18 +217,20 @@ func (s *Server) Reserve(id string, traceID, lambdaSegmentID string) (*ReserveRe return resp, err } -func (s *Server) awaitInitCompletion() { - initSuccess, initFailure := s.initContext.Wait() +func (s *Server) awaitInitCompletion(initContext interop.InitContext, initFailures chan interop.InitFailure, completion *initCompletion) { + initSuccess, initFailure := initContext.Wait() + completion.completedAt = time.Now() + close(completion.done) if initFailure != nil { // In standalone, we don't have to block rapid start() goroutine until init failure is consumed // because there is no channel back to the invoker until an invoke arrives via a Reserve() initFailure.Ack <- struct{}{} - s.initFailures <- *initFailure + initFailures <- *initFailure } else { initSuccess.Ack <- struct{}{} } // always closing the channel makes this method idempotent - close(s.initFailures) + close(initFailures) } func (s *Server) setReplyStream(w http.ResponseWriter, direct bool) (string, error) { @@ -500,10 +508,11 @@ func deadlineNsFromTimeoutMs(timeoutMs int64) int64 { return mono + timeoutMs*1000*1000 } -func (s *Server) setInitFailuresChan() { +func (s *Server) setInitFailuresChan() chan interop.InitFailure { s.mutex.Lock() defer s.mutex.Unlock() s.initFailures = make(chan interop.InitFailure) + return s.initFailures } func (s *Server) getInitFailuresChan() chan interop.InitFailure { @@ -512,18 +521,41 @@ func (s *Server) getInitFailuresChan() chan interop.InitFailure { return s.initFailures } +func (s *Server) setInitCompletion() *initCompletion { + s.mutex.Lock() + defer s.mutex.Unlock() + s.initCompletion = &initCompletion{done: make(chan struct{})} + return s.initCompletion +} + +func (s *Server) getInitCompletion() *initCompletion { + s.mutex.Lock() + defer s.mutex.Unlock() + return s.initCompletion +} + func (s *Server) Init(i *interop.Init, invokeTimeoutMs int64) error { s.SetInvokeTimeout(time.Duration(invokeTimeoutMs) * time.Millisecond) s.setRapidPhase(phaseInitializing) - s.setInitFailuresChan() + initFailures := s.setInitFailuresChan() + completion := s.setInitCompletion() initCtx := s.sandboxContext.Init(i, invokeTimeoutMs) s.initContext = initCtx - go s.awaitInitCompletion() + go s.awaitInitCompletion(initCtx, initFailures, completion) return nil } +func (s *Server) AwaitInitCompletion() time.Time { + completion := s.getInitCompletion() + if completion == nil { + return time.Time{} + } + <-completion.done + return completion.completedAt +} + func (s *Server) FastInvoke(w http.ResponseWriter, i *interop.Invoke, direct bool) error { invokeID, err := s.setReplyStream(w, direct) if err != nil { diff --git a/internal/lambda/rapidcore/server_test.go b/internal/lambda/rapidcore/server_test.go index 27402d4f..3f4c3e7f 100644 --- a/internal/lambda/rapidcore/server_test.go +++ b/internal/lambda/rapidcore/server_test.go @@ -12,10 +12,10 @@ import ( "testing" "time" - "github.com/stretchr/testify/require" "github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/core/statejson" "github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop" "github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/rapidcore/env" + "github.com/stretchr/testify/require" ) func waitForChanWithTimeout(channel <-chan error, timeout time.Duration) error { @@ -133,6 +133,49 @@ func TestInitSuccess(t *testing.T) { require.NoError(t, err) } +func TestAwaitInitCompletionWaitsWithoutConsumingFailure(t *testing.T) { + srv := NewServer() + srv.SetInternalStateGetter(func() statejson.InternalStateDescription { return statejson.InternalStateDescription{} }) + + releaseRuntimeInit := make(chan struct{}) + initHandler := func(successResp chan<- interop.InitSuccess, failureResp chan<- interop.InitFailure) { + <-releaseRuntimeInit + sendInitFailureResponse(failureResp, interop.InitFailure{}) + } + srv.SetSandboxContext(&SandboxContext{&mockRapidCtx{ + initHandler, + func() (interop.InvokeSuccess, *interop.InvokeFailure) { return interop.InvokeSuccess{}, nil }, + func() (interop.ResetSuccess, *interop.ResetFailure) { return interop.ResetSuccess{}, nil }, + }, "handler", "runtimeAPIhost:999", "test-token"}) + + srv.Init(&interop.Init{EnvironmentVariables: env.NewEnvironment()}, int64(time.Second/time.Millisecond)) + initCompleted := make(chan struct{}) + var completedAt time.Time + go func() { + completedAt = srv.AwaitInitCompletion() + close(initCompleted) + }() + + select { + case <-initCompleted: + require.Fail(t, "init completion returned before runtime initialization finished") + case <-time.After(10 * time.Millisecond): + } + + close(releaseRuntimeInit) + select { + case <-initCompleted: + case <-time.After(time.Second): + require.Fail(t, "timed out waiting for init completion") + } + require.False(t, completedAt.IsZero()) + require.ErrorIs(t, srv.AwaitInitialized(), ErrInitDoneFailed) +} + +func TestAwaitInitCompletionBeforeInitReturnsZeroTime(t *testing.T) { + require.True(t, NewServer().AwaitInitCompletion().IsZero()) +} + func TestInitErrorBeforeReserve(t *testing.T) { // Rapid thread sending init failure should not be blocked even if reserve hasn't arrived srv := NewServer() diff --git a/internal/lambda/rie/handlers.go b/internal/lambda/rie/handlers.go index 1abe1541..12525fcf 100644 --- a/internal/lambda/rie/handlers.go +++ b/internal/lambda/rie/handlers.go @@ -14,6 +14,7 @@ import ( "os" "strconv" "strings" + "sync" "time" "github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/core/statejson" @@ -28,6 +29,7 @@ import ( type Sandbox interface { Init(i *interop.Init, invokeTimeoutMs int64) + AwaitInitCompletion() time.Time Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error } @@ -44,7 +46,10 @@ type InteropServer interface { Restore(restore *interop.Restore) error } -var initDone bool +var ( + initDone bool + initMutex sync.Mutex +) func GetenvWithDefault(key string, defaultValue string) string { envValue := os.Getenv(key) @@ -74,6 +79,38 @@ func printEndReports(invokeId string, initDuration string, memorySize string, in invokeId, invokeDuration, math.Ceil(invokeDuration), memorySize, memorySize) } +func startInitOnce(sandbox Sandbox, functionVersion string, timeout int64, bs interop.Bootstrap) time.Time { + initMutex.Lock() + defer initMutex.Unlock() + + if initDone { + return time.Time{} + } + + initStart := InitHandler(sandbox, functionVersion, timeout, bs) + initDone = true + return initStart +} + +func formatInitDuration(initStart time.Time, initEnd time.Time, timeoutDuration time.Duration) string { + if initStart.IsZero() || initEnd.IsZero() { + return "" + } + + initTimeMS := math.Min(float64(initEnd.Sub(initStart).Nanoseconds()), + float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond) + return fmt.Sprintf("Init Duration: %.2f ms\t", initTimeMS) +} + +func printInvokeReport(sandbox Sandbox, invokeID string, initStart time.Time, invokeStart time.Time, memorySize string, timeoutDuration time.Duration) { + initEnd := sandbox.AwaitInitCompletion() + initDuration := formatInitDuration(initStart, initEnd, timeoutDuration) + if !initStart.IsZero() && initEnd.After(invokeStart) { + invokeStart = initEnd + } + printEndReports(invokeID, initDuration, memorySize, invokeStart, timeoutDuration) +} + func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs interop.Bootstrap) { log.Debugf("invoke: -> %s %s %v", r.Method, r.URL, r.Header) bodyBytes, err := ioutil.ReadAll(r.Body) @@ -90,7 +127,6 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i return } - initDuration := "" inv := GetenvWithDefault("AWS_LAMBDA_FUNCTION_TIMEOUT", "300") timeoutDuration, _ := time.ParseDuration(inv + "s") // Default @@ -102,19 +138,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i functionVersion := GetenvWithDefault("AWS_LAMBDA_FUNCTION_VERSION", "$LATEST") memorySize := GetenvWithDefault("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "3008") - if !initDone { - - initStart, initEnd := InitHandler(sandbox, functionVersion, timeout, bs) - - // Calculate InitDuration - initTimeMS := math.Min(float64(initEnd.Sub(initStart).Nanoseconds()), - float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond) - - initDuration = fmt.Sprintf("Init Duration: %.2f ms\t", initTimeMS) - - // Set initDone so next invokes do not try to Init the function again - initDone = true - } + initStart := startInitOnce(sandbox, functionVersion, timeout, bs) invokeStart := time.Now() invokeID := r.Header.Get("X-Amzn-RequestId") @@ -197,7 +221,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i w.WriteHeader(http.StatusGatewayTimeout) return case rapidcore.ErrInvokeTimeout: - printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration) + printInvokeReport(sandbox, invokePayload.ID, initStart, invokeStart, memorySize, timeoutDuration) w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout))) time.Sleep(100 * time.Millisecond) @@ -206,7 +230,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i } } - printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration) + printInvokeReport(sandbox, invokePayload.ID, initStart, invokeStart, memorySize, timeoutDuration) if invokeResp.StatusCode != 0 { w.WriteHeader(invokeResp.StatusCode) @@ -214,7 +238,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i w.Write(invokeResp.Body) } -func InitHandler(sandbox Sandbox, functionVersion string, timeout int64, bs interop.Bootstrap) (time.Time, time.Time) { +func InitHandler(sandbox Sandbox, functionVersion string, timeout int64, bs interop.Bootstrap) time.Time { additionalFunctionEnvironmentVariables := map[string]string{} // Add default Env Vars if they were not defined. This is a required otherwise 1p Python2.7, Python3.6, and @@ -252,6 +276,5 @@ func InitHandler(sandbox Sandbox, functionVersion string, timeout int64, bs inte Bootstrap: bs, EnvironmentVariables: env.NewEnvironment(), }, timeout*1000) - initEnd := time.Now() - return initStart, initEnd + return initStart } diff --git a/internal/lambda/rie/handlers_test.go b/internal/lambda/rie/handlers_test.go new file mode 100644 index 00000000..773b5b76 --- /dev/null +++ b/internal/lambda/rie/handlers_test.go @@ -0,0 +1,251 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package rie + +import ( + "io" + "net/http" + "net/http/httptest" + "os" + "regexp" + "strconv" + "testing" + "time" + + "github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop" + "github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/rapidcore" + "github.com/stretchr/testify/require" +) + +type delayedInitSandbox struct { + initDelay time.Duration + invokeDelay time.Duration + invokeCalled bool + initCompletedAt time.Time + invokeErr error +} + +func (s *delayedInitSandbox) Init(*interop.Init, int64) {} + +func (s *delayedInitSandbox) AwaitInitCompletion() time.Time { + if !s.invokeCalled { + panic("AwaitInitCompletion called before Invoke") + } + return s.initCompletedAt +} + +func (s *delayedInitSandbox) Invoke(http.ResponseWriter, *interop.Invoke) error { + s.invokeCalled = true + time.Sleep(s.initDelay) + s.initCompletedAt = time.Now() + time.Sleep(s.invokeDelay) + return s.invokeErr +} + +type panicInitSandbox struct { + delayedInitSandbox +} + +func (s *panicInitSandbox) Init(*interop.Init, int64) { + panic("init failed") +} + +func TestStartInitOnceReleasesLockAfterPanic(t *testing.T) { + initMutex.Lock() + initDone = false + initMutex.Unlock() + t.Cleanup(func() { + initMutex.Lock() + initDone = false + initMutex.Unlock() + }) + + func() { + defer func() { require.Equal(t, "init failed", recover()) }() + startInitOnce(&panicInitSandbox{}, "$LATEST", 1, nil) + }() + + initStarted := make(chan struct{}) + go func() { + startInitOnce(&delayedInitSandbox{}, "$LATEST", 1, nil) + close(initStarted) + }() + select { + case <-initStarted: + case <-time.After(time.Second): + require.Fail(t, "init mutex remained locked after panic") + } +} + +func TestInvokeHandlerReportsRuntimeInitDuration(t *testing.T) { + initMutex.Lock() + initDone = false + initMutex.Unlock() + t.Cleanup(func() { + initMutex.Lock() + initDone = false + initMutex.Unlock() + }) + + request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil) + response := httptest.NewRecorder() + sandbox := &delayedInitSandbox{initDelay: 50 * time.Millisecond} + + reader, writer, err := os.Pipe() + require.NoError(t, err) + originalStdout := os.Stdout + os.Stdout = writer + t.Cleanup(func() { os.Stdout = originalStdout }) + + InvokeHandler(response, request, sandbox, nil) + require.NoError(t, writer.Close()) + output, err := io.ReadAll(reader) + require.NoError(t, err) + require.NoError(t, reader.Close()) + + matches := regexp.MustCompile(`Init Duration: ([0-9.]+) ms`).FindStringSubmatch(string(output)) + require.Len(t, matches, 2) + durationMilliseconds, err := strconv.ParseFloat(matches[1], 64) + require.NoError(t, err) + require.GreaterOrEqual(t, durationMilliseconds, float64(40)) + + matches = regexp.MustCompile(`\tDuration: ([0-9.]+) ms`).FindStringSubmatch(string(output)) + require.Len(t, matches, 2) + durationMilliseconds, err = strconv.ParseFloat(matches[1], 64) + require.NoError(t, err) + require.Less(t, durationMilliseconds, float64(40)) +} + +func TestInvokeHandlerReportsInitDurationWhenInitTimesOut(t *testing.T) { + initMutex.Lock() + initDone = false + initMutex.Unlock() + t.Cleanup(func() { + initMutex.Lock() + initDone = false + initMutex.Unlock() + }) + t.Setenv("AWS_LAMBDA_FUNCTION_TIMEOUT", "1") + + request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil) + response := httptest.NewRecorder() + sandbox := &delayedInitSandbox{ + initDelay: 50 * time.Millisecond, + invokeErr: rapidcore.ErrInvokeTimeout, + } + + reader, writer, err := os.Pipe() + require.NoError(t, err) + originalStdout := os.Stdout + os.Stdout = writer + t.Cleanup(func() { os.Stdout = originalStdout }) + + InvokeHandler(response, request, sandbox, nil) + require.NoError(t, writer.Close()) + output, err := io.ReadAll(reader) + require.NoError(t, err) + require.NoError(t, reader.Close()) + + require.Equal(t, "Task timed out after 1.00 seconds", response.Body.String()) + matches := regexp.MustCompile(`Init Duration: ([0-9.]+) ms`).FindStringSubmatch(string(output)) + require.Len(t, matches, 2) + initDurationMilliseconds, err := strconv.ParseFloat(matches[1], 64) + require.NoError(t, err) + require.Greater(t, initDurationMilliseconds, float64(0)) + require.LessOrEqual(t, initDurationMilliseconds, float64(1000)) + + matches = regexp.MustCompile(`\tDuration: ([0-9.]+) ms`).FindStringSubmatch(string(output)) + require.Len(t, matches, 2) + durationMilliseconds, err := strconv.ParseFloat(matches[1], 64) + require.NoError(t, err) + require.Less(t, durationMilliseconds, initDurationMilliseconds) + require.LessOrEqual(t, initDurationMilliseconds+durationMilliseconds, float64(1020)) +} + +func TestInvokeHandlerSeparatesInitFromTimedOutInvocation(t *testing.T) { + initMutex.Lock() + initDone = false + initMutex.Unlock() + t.Cleanup(func() { + initMutex.Lock() + initDone = false + initMutex.Unlock() + }) + t.Setenv("AWS_LAMBDA_FUNCTION_TIMEOUT", "1") + + request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil) + response := httptest.NewRecorder() + sandbox := &delayedInitSandbox{ + initDelay: 200 * time.Millisecond, + invokeDelay: 200 * time.Millisecond, + invokeErr: rapidcore.ErrInvokeTimeout, + } + + reader, writer, err := os.Pipe() + require.NoError(t, err) + originalStdout := os.Stdout + os.Stdout = writer + t.Cleanup(func() { os.Stdout = originalStdout }) + + start := time.Now() + InvokeHandler(response, request, sandbox, nil) + elapsedMilliseconds := float64(time.Since(start)) / float64(time.Millisecond) + require.NoError(t, writer.Close()) + output, err := io.ReadAll(reader) + require.NoError(t, err) + require.NoError(t, reader.Close()) + + require.Equal(t, "Task timed out after 1.00 seconds", response.Body.String()) + initMatches := regexp.MustCompile(`Init Duration: ([0-9.]+) ms`).FindStringSubmatch(string(output)) + require.Len(t, initMatches, 2) + initDurationMilliseconds, err := strconv.ParseFloat(initMatches[1], 64) + require.NoError(t, err) + require.GreaterOrEqual(t, initDurationMilliseconds, float64(40)) + + durationMatches := regexp.MustCompile(`\tDuration: ([0-9.]+) ms`).FindStringSubmatch(string(output)) + require.Len(t, durationMatches, 2) + durationMilliseconds, err := strconv.ParseFloat(durationMatches[1], 64) + require.NoError(t, err) + require.GreaterOrEqual(t, durationMilliseconds, float64(190)) + require.LessOrEqual(t, initDurationMilliseconds+durationMilliseconds, elapsedMilliseconds) +} + +func TestInvokeHandlerReportsWarmTimeoutWithoutInitDuration(t *testing.T) { + initMutex.Lock() + initDone = true + initMutex.Unlock() + t.Cleanup(func() { + initMutex.Lock() + initDone = false + initMutex.Unlock() + }) + t.Setenv("AWS_LAMBDA_FUNCTION_TIMEOUT", "1") + + request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil) + response := httptest.NewRecorder() + sandbox := &delayedInitSandbox{ + invokeDelay: 50 * time.Millisecond, + invokeErr: rapidcore.ErrInvokeTimeout, + } + + reader, writer, err := os.Pipe() + require.NoError(t, err) + originalStdout := os.Stdout + os.Stdout = writer + t.Cleanup(func() { os.Stdout = originalStdout }) + + InvokeHandler(response, request, sandbox, nil) + require.NoError(t, writer.Close()) + output, err := io.ReadAll(reader) + require.NoError(t, err) + require.NoError(t, reader.Close()) + + require.Equal(t, "Task timed out after 1.00 seconds", response.Body.String()) + require.NotContains(t, string(output), "Init Duration:") + durationMatches := regexp.MustCompile(`\tDuration: ([0-9.]+) ms`).FindStringSubmatch(string(output)) + require.Len(t, durationMatches, 2) + durationMilliseconds, err := strconv.ParseFloat(durationMatches[1], 64) + require.NoError(t, err) + require.GreaterOrEqual(t, durationMilliseconds, float64(40)) +}