Fix RIE init duration reporting

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)

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 {

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()

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)
@@ -56,9 +61,13 @@ func GetenvWithDefault(key string, defaultValue string) string {
 	return envValue
 }
 
-func printEndReports(invokeId string, initDuration string, memorySize string, invokeStart time.Time, timeoutDuration time.Duration) {
+func printEndReports(invokeId string, initDuration string, memorySize string, invokeStart time.Time, invokeEnd time.Time, timeoutDuration time.Duration) {
 	// Calcuation invoke duration
-	invokeDuration := math.Min(float64(time.Now().Sub(invokeStart).Nanoseconds()),
+	elapsed := invokeEnd.Sub(invokeStart)
+	if elapsed < 0 {
+		elapsed = 0
+	}
+	invokeDuration := math.Min(float64(elapsed.Nanoseconds()),
 		float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond)
 
 	fmt.Println("END RequestId: " + invokeId)
@@ -74,6 +83,68 @@ 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)
+}
+
+// initReportGracePeriod bounds the post-invoke wait for init completion.
+// Rapid's reset path is capped at 2s; this is slightly larger so a normal
+// timeout-during-init still records Init Duration, while a stuck init omits
+// that field instead of hanging the HTTP handler.
+var initReportGracePeriod = 2500 * time.Millisecond
+
+func awaitInitCompletionWithin(sandbox Sandbox, timeout time.Duration) time.Time {
+	result := make(chan time.Time, 1)
+	go func() {
+		result <- sandbox.AwaitInitCompletion()
+	}()
+
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+
+	select {
+	case initEnd := <-result:
+		return initEnd
+	case <-timer.C:
+		log.Warn("Timed out waiting for init completion; omitting Init Duration from REPORT")
+		return time.Time{}
+	}
+}
+
+func printInvokeReport(sandbox Sandbox, invokeID string, initStart time.Time, invokeStart time.Time, memorySize string, timeoutDuration time.Duration) {
+	invokeEnd := time.Now()
+	if initStart.IsZero() {
+		printEndReports(invokeID, "", memorySize, invokeStart, invokeEnd, timeoutDuration)
+		return
+	}
+
+	initEnd := awaitInitCompletionWithin(sandbox, initReportGracePeriod)
+	initDuration := formatInitDuration(initStart, initEnd, timeoutDuration)
+	if initEnd.After(invokeStart) {
+		invokeStart = initEnd
+	}
+	printEndReports(invokeID, initDuration, memorySize, invokeStart, invokeEnd, 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 +161,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 +172,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,24 +255,26 @@ 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)
-
 			w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
+			if flusher, ok := w.(http.Flusher); ok {
+				flusher.Flush()
+			}
+			printInvokeReport(sandbox, invokePayload.ID, initStart, invokeStart, memorySize, timeoutDuration)
 			time.Sleep(100 * time.Millisecond)
 			//initDone = false
 			return
 		}
 	}
 
-	printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)
+	printInvokeReport(sandbox, invokePayload.ID, initStart, invokeStart, memorySize, timeoutDuration)
 
 	if invokeResp.StatusCode != 0 {
 		w.WriteHeader(invokeResp.StatusCode)
 	}
 	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 +312,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
 }