From 1171a17ec144337e0a357e9ca6067f313d6f79a1 Mon Sep 17 00:00:00 2001
From: Jayanth Sai Yarlagadda <72214836+JayYarlagadda@users.noreply.github.com>
Date: Sun, 2 Aug 2026 17:24:24 -0700
Subject: [PATCH 1/8] fix: report complete runtime init duration

---
 .../lambda/rapidcore/sandbox_emulator_api.go  |  5 ++
 internal/lambda/rapidcore/server.go           |  7 +++
 internal/lambda/rapidcore/server_test.go      | 39 +++++++++++-
 internal/lambda/rie/handlers.go               | 11 ++--
 internal/lambda/rie/handlers_test.go          | 59 +++++++++++++++++++
 5 files changed, 115 insertions(+), 6 deletions(-)
 create mode 100644 internal/lambda/rie/handlers_test.go

diff --git a/internal/lambda/rapidcore/sandbox_emulator_api.go b/internal/lambda/rapidcore/sandbox_emulator_api.go
index d5c1d47b..5a8efc55 100644
--- a/internal/lambda/rapidcore/sandbox_emulator_api.go
+++ b/internal/lambda/rapidcore/sandbox_emulator_api.go
@@ -12,6 +12,7 @@ import (
 // LambdaInvokeAPI are the methods used by the Runtime Interface Emulator
 type LambdaInvokeAPI interface {
 	Init(i *interop.Init, invokeTimeoutMs int64)
+	AwaitInitCompletion()
 	Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
 }
 
@@ -47,6 +48,10 @@ func (l *EmulatorAPI) Init(i *interop.Init, timeoutMs int64) {
 	}, timeoutMs)
 }
 
+func (l *EmulatorAPI) AwaitInitCompletion() {
+	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..8462ae39 100644
--- a/internal/lambda/rapidcore/server.go
+++ b/internal/lambda/rapidcore/server.go
@@ -100,6 +100,7 @@ type Server struct {
 	initContext             interop.InitContext
 	invoker                 interop.InvokeContext
 	initFailures            chan interop.InitFailure
+	initCompleted           chan struct{}
 	cachedInitErrorResponse *interop.ErrorInvokeResponse
 }
 
@@ -213,6 +214,7 @@ func (s *Server) Reserve(id string, traceID, lambdaSegmentID string) (*ReserveRe
 
 func (s *Server) awaitInitCompletion() {
 	initSuccess, initFailure := s.initContext.Wait()
+	close(s.initCompleted)
 	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()
@@ -516,6 +518,7 @@ func (s *Server) Init(i *interop.Init, invokeTimeoutMs int64) error {
 	s.SetInvokeTimeout(time.Duration(invokeTimeoutMs) * time.Millisecond)
 	s.setRapidPhase(phaseInitializing)
 	s.setInitFailuresChan()
+	s.initCompleted = make(chan struct{})
 	initCtx := s.sandboxContext.Init(i, invokeTimeoutMs)
 
 	s.initContext = initCtx
@@ -524,6 +527,10 @@ func (s *Server) Init(i *interop.Init, invokeTimeoutMs int64) error {
 	return nil
 }
 
+func (s *Server) AwaitInitCompletion() {
+	<-s.initCompleted
+}
+
 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..666feb0f 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,43 @@ 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{})
+	go func() {
+		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.ErrorIs(t, srv.AwaitInitialized(), ErrInitDoneFailed)
+}
+
 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..90f31838 100644
--- a/internal/lambda/rie/handlers.go
+++ b/internal/lambda/rie/handlers.go
@@ -28,6 +28,7 @@ import (
 
 type Sandbox interface {
 	Init(i *interop.Init, invokeTimeoutMs int64)
+	AwaitInitCompletion()
 	Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
 }
 
@@ -104,10 +105,11 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 
 	if !initDone {
 
-		initStart, initEnd := InitHandler(sandbox, functionVersion, timeout, bs)
+		initStart := InitHandler(sandbox, functionVersion, timeout, bs)
+		sandbox.AwaitInitCompletion()
 
 		// Calculate InitDuration
-		initTimeMS := math.Min(float64(initEnd.Sub(initStart).Nanoseconds()),
+		initTimeMS := math.Min(float64(time.Since(initStart).Nanoseconds()),
 			float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond)
 
 		initDuration = fmt.Sprintf("Init Duration: %.2f ms\t", initTimeMS)
@@ -214,7 +216,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 +254,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..81bf4c97
--- /dev/null
+++ b/internal/lambda/rie/handlers_test.go
@@ -0,0 +1,59 @@
+// 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/stretchr/testify/require"
+)
+
+type delayedInitSandbox struct {
+	delay time.Duration
+}
+
+func (s *delayedInitSandbox) Init(*interop.Init, int64) {}
+
+func (s *delayedInitSandbox) AwaitInitCompletion() {
+	time.Sleep(s.delay)
+}
+
+func (s *delayedInitSandbox) Invoke(http.ResponseWriter, *interop.Invoke) error {
+	return nil
+}
+
+func TestInvokeHandlerReportsRuntimeInitDuration(t *testing.T) {
+	initDone = false
+	t.Cleanup(func() { initDone = false })
+
+	request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil)
+	response := httptest.NewRecorder()
+	sandbox := &delayedInitSandbox{delay: 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))
+}

From 1668238a6f923fc9dbe8a4a0a5548092938c4423 Mon Sep 17 00:00:00 2001
From: Jayanth Sai Yarlagadda <72214836+JayYarlagadda@users.noreply.github.com>
Date: Sun, 2 Aug 2026 17:54:57 -0700
Subject: [PATCH 2/8] fix: preserve timeout while measuring init

---
 .../lambda/rapidcore/sandbox_emulator_api.go  |  7 +--
 internal/lambda/rapidcore/server.go           | 49 ++++++++++++++-----
 internal/lambda/rapidcore/server_test.go      |  8 ++-
 internal/lambda/rie/handlers.go               | 42 ++++++++++------
 internal/lambda/rie/handlers_test.go          | 22 +++++++--
 5 files changed, 93 insertions(+), 35 deletions(-)

diff --git a/internal/lambda/rapidcore/sandbox_emulator_api.go b/internal/lambda/rapidcore/sandbox_emulator_api.go
index 5a8efc55..a6efea1c 100644
--- a/internal/lambda/rapidcore/sandbox_emulator_api.go
+++ b/internal/lambda/rapidcore/sandbox_emulator_api.go
@@ -7,12 +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()
+	AwaitInitCompletion() time.Time
 	Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
 }
 
@@ -48,8 +49,8 @@ func (l *EmulatorAPI) Init(i *interop.Init, timeoutMs int64) {
 	}, timeoutMs)
 }
 
-func (l *EmulatorAPI) AwaitInitCompletion() {
-	l.server.AwaitInitCompletion()
+func (l *EmulatorAPI) AwaitInitCompletion() time.Time {
+	return l.server.AwaitInitCompletion()
 }
 
 // Invoke method is only used by the Runtime interface emulator
diff --git a/internal/lambda/rapidcore/server.go b/internal/lambda/rapidcore/server.go
index 8462ae39..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,7 +105,7 @@ type Server struct {
 	initContext             interop.InitContext
 	invoker                 interop.InvokeContext
 	initFailures            chan interop.InitFailure
-	initCompleted           chan struct{}
+	initCompletion          *initCompletion
 	cachedInitErrorResponse *interop.ErrorInvokeResponse
 }
 
@@ -212,19 +217,20 @@ func (s *Server) Reserve(id string, traceID, lambdaSegmentID string) (*ReserveRe
 	return resp, err
 }
 
-func (s *Server) awaitInitCompletion() {
-	initSuccess, initFailure := s.initContext.Wait()
-	close(s.initCompleted)
+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) {
@@ -502,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 {
@@ -514,21 +521,39 @@ 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()
-	s.initCompleted = make(chan struct{})
+	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() {
-	<-s.initCompleted
+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 {
diff --git a/internal/lambda/rapidcore/server_test.go b/internal/lambda/rapidcore/server_test.go
index 666feb0f..3f4c3e7f 100644
--- a/internal/lambda/rapidcore/server_test.go
+++ b/internal/lambda/rapidcore/server_test.go
@@ -150,8 +150,9 @@ func TestAwaitInitCompletionWaitsWithoutConsumingFailure(t *testing.T) {
 
 	srv.Init(&interop.Init{EnvironmentVariables: env.NewEnvironment()}, int64(time.Second/time.Millisecond))
 	initCompleted := make(chan struct{})
+	var completedAt time.Time
 	go func() {
-		srv.AwaitInitCompletion()
+		completedAt = srv.AwaitInitCompletion()
 		close(initCompleted)
 	}()
 
@@ -167,9 +168,14 @@ func TestAwaitInitCompletionWaitsWithoutConsumingFailure(t *testing.T) {
 	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 90f31838..8e4ad6bc 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,7 +29,7 @@ import (
 
 type Sandbox interface {
 	Init(i *interop.Init, invokeTimeoutMs int64)
-	AwaitInitCompletion()
+	AwaitInitCompletion() time.Time
 	Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
 }
 
@@ -45,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)
@@ -75,6 +79,21 @@ func printEndReports(invokeId string, initDuration string, memorySize string, in
 		invokeId, invokeDuration, math.Ceil(invokeDuration), memorySize, memorySize)
 }
 
+func formatInitDuration(sandbox Sandbox, initStart time.Time, timeoutDuration time.Duration) string {
+	if initStart.IsZero() {
+		return ""
+	}
+
+	initEnd := sandbox.AwaitInitCompletion()
+	if 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 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)
@@ -91,7 +110,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
@@ -103,20 +121,14 @@ 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")
 
+	var initStart time.Time
+	initMutex.Lock()
 	if !initDone {
-
-		initStart := InitHandler(sandbox, functionVersion, timeout, bs)
-		sandbox.AwaitInitCompletion()
-
-		// Calculate InitDuration
-		initTimeMS := math.Min(float64(time.Since(initStart).Nanoseconds()),
-			float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond)
-
-		initDuration = fmt.Sprintf("Init Duration: %.2f ms\t", initTimeMS)
-
+		initStart = InitHandler(sandbox, functionVersion, timeout, bs)
 		// Set initDone so next invokes do not try to Init the function again
 		initDone = true
 	}
+	initMutex.Unlock()
 
 	invokeStart := time.Now()
 	invokeID := r.Header.Get("X-Amzn-RequestId")
@@ -199,7 +211,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)
+			printEndReports(invokePayload.ID, formatInitDuration(sandbox, initStart, timeoutDuration), memorySize, invokeStart, timeoutDuration)
 
 			w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
 			time.Sleep(100 * time.Millisecond)
@@ -208,7 +220,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		}
 	}
 
-	printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)
+	printEndReports(invokePayload.ID, formatInitDuration(sandbox, initStart, timeoutDuration), memorySize, invokeStart, timeoutDuration)
 
 	if invokeResp.StatusCode != 0 {
 		w.WriteHeader(invokeResp.StatusCode)
diff --git a/internal/lambda/rie/handlers_test.go b/internal/lambda/rie/handlers_test.go
index 81bf4c97..95043324 100644
--- a/internal/lambda/rie/handlers_test.go
+++ b/internal/lambda/rie/handlers_test.go
@@ -18,22 +18,36 @@ import (
 )
 
 type delayedInitSandbox struct {
-	delay time.Duration
+	delay           time.Duration
+	invokeCalled    bool
+	initCompletedAt time.Time
 }
 
 func (s *delayedInitSandbox) Init(*interop.Init, int64) {}
 
-func (s *delayedInitSandbox) AwaitInitCompletion() {
-	time.Sleep(s.delay)
+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.delay)
+	s.initCompletedAt = time.Now()
 	return nil
 }
 
 func TestInvokeHandlerReportsRuntimeInitDuration(t *testing.T) {
+	initMutex.Lock()
 	initDone = false
-	t.Cleanup(func() { 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()

From 2c52c1e9dd14adf096762342caa658c3db27ced7 Mon Sep 17 00:00:00 2001
From: Jayanth Sai Yarlagadda <72214836+JayYarlagadda@users.noreply.github.com>
Date: Sun, 2 Aug 2026 20:54:53 -0700
Subject: [PATCH 3/8] fix: separate init and invoke durations

---
 internal/lambda/rie/handlers.go      | 37 +++++++++++++++----------
 internal/lambda/rie/handlers_test.go | 41 ++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/internal/lambda/rie/handlers.go b/internal/lambda/rie/handlers.go
index 8e4ad6bc..28e470a0 100644
--- a/internal/lambda/rie/handlers.go
+++ b/internal/lambda/rie/handlers.go
@@ -79,13 +79,21 @@ func printEndReports(invokeId string, initDuration string, memorySize string, in
 		invokeId, invokeDuration, math.Ceil(invokeDuration), memorySize, memorySize)
 }
 
-func formatInitDuration(sandbox Sandbox, initStart time.Time, timeoutDuration time.Duration) string {
-	if initStart.IsZero() {
-		return ""
+func startInitOnce(sandbox Sandbox, functionVersion string, timeout int64, bs interop.Bootstrap) time.Time {
+	initMutex.Lock()
+	defer initMutex.Unlock()
+
+	if initDone {
+		return time.Time{}
 	}
 
-	initEnd := sandbox.AwaitInitCompletion()
-	if initEnd.IsZero() {
+	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 ""
 	}
 
@@ -121,14 +129,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")
 
-	var initStart time.Time
-	initMutex.Lock()
-	if !initDone {
-		initStart = InitHandler(sandbox, functionVersion, timeout, bs)
-		// Set initDone so next invokes do not try to Init the function again
-		initDone = true
-	}
-	initMutex.Unlock()
+	initStart := startInitOnce(sandbox, functionVersion, timeout, bs)
 
 	invokeStart := time.Now()
 	invokeID := r.Header.Get("X-Amzn-RequestId")
@@ -211,7 +212,8 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 			w.WriteHeader(http.StatusGatewayTimeout)
 			return
 		case rapidcore.ErrInvokeTimeout:
-			printEndReports(invokePayload.ID, formatInitDuration(sandbox, initStart, timeoutDuration), memorySize, invokeStart, timeoutDuration)
+			initEnd := sandbox.AwaitInitCompletion()
+			printEndReports(invokePayload.ID, formatInitDuration(initStart, initEnd, timeoutDuration), memorySize, invokeStart, timeoutDuration)
 
 			w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
 			time.Sleep(100 * time.Millisecond)
@@ -220,7 +222,12 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		}
 	}
 
-	printEndReports(invokePayload.ID, formatInitDuration(sandbox, initStart, timeoutDuration), memorySize, invokeStart, timeoutDuration)
+	initEnd := sandbox.AwaitInitCompletion()
+	initDuration := formatInitDuration(initStart, initEnd, timeoutDuration)
+	if !initStart.IsZero() && initEnd.After(invokeStart) {
+		invokeStart = initEnd
+	}
+	printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)
 
 	if invokeResp.StatusCode != 0 {
 		w.WriteHeader(invokeResp.StatusCode)
diff --git a/internal/lambda/rie/handlers_test.go b/internal/lambda/rie/handlers_test.go
index 95043324..3aa55b96 100644
--- a/internal/lambda/rie/handlers_test.go
+++ b/internal/lambda/rie/handlers_test.go
@@ -39,6 +39,41 @@ func (s *delayedInitSandbox) Invoke(http.ResponseWriter, *interop.Invoke) error
 	return nil
 }
 
+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
@@ -70,4 +105,10 @@ func TestInvokeHandlerReportsRuntimeInitDuration(t *testing.T) {
 	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))
 }

From 9ea5783132a20178699419d9ab72169adb53c2cc Mon Sep 17 00:00:00 2001
From: Jayanth Sai Yarlagadda <72214836+JayYarlagadda@users.noreply.github.com>
Date: Sun, 2 Aug 2026 21:18:07 -0700
Subject: [PATCH 4/8] fix: omit init duration when init fails

---
 .../lambda/rapidcore/sandbox_emulator_api.go  |  4 +-
 internal/lambda/rapidcore/server.go           |  8 ++--
 internal/lambda/rapidcore/server_test.go      |  8 +++-
 internal/lambda/rie/handlers.go               | 17 +++++--
 internal/lambda/rie/handlers_test.go          | 47 +++++++++++++++++--
 5 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/internal/lambda/rapidcore/sandbox_emulator_api.go b/internal/lambda/rapidcore/sandbox_emulator_api.go
index a6efea1c..7d9d86cd 100644
--- a/internal/lambda/rapidcore/sandbox_emulator_api.go
+++ b/internal/lambda/rapidcore/sandbox_emulator_api.go
@@ -13,7 +13,7 @@ import (
 // LambdaInvokeAPI are the methods used by the Runtime Interface Emulator
 type LambdaInvokeAPI interface {
 	Init(i *interop.Init, invokeTimeoutMs int64)
-	AwaitInitCompletion() time.Time
+	AwaitInitCompletion() (time.Time, bool)
 	Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
 }
 
@@ -49,7 +49,7 @@ func (l *EmulatorAPI) Init(i *interop.Init, timeoutMs int64) {
 	}, timeoutMs)
 }
 
-func (l *EmulatorAPI) AwaitInitCompletion() time.Time {
+func (l *EmulatorAPI) AwaitInitCompletion() (time.Time, bool) {
 	return l.server.AwaitInitCompletion()
 }
 
diff --git a/internal/lambda/rapidcore/server.go b/internal/lambda/rapidcore/server.go
index f3c159ce..94608950 100644
--- a/internal/lambda/rapidcore/server.go
+++ b/internal/lambda/rapidcore/server.go
@@ -75,6 +75,7 @@ type InvokeContext struct {
 type initCompletion struct {
 	done        chan struct{}
 	completedAt time.Time
+	succeeded   bool
 }
 
 type Server struct {
@@ -220,6 +221,7 @@ func (s *Server) Reserve(id string, traceID, lambdaSegmentID string) (*ReserveRe
 func (s *Server) awaitInitCompletion(initContext interop.InitContext, initFailures chan interop.InitFailure, completion *initCompletion) {
 	initSuccess, initFailure := initContext.Wait()
 	completion.completedAt = time.Now()
+	completion.succeeded = initFailure == nil
 	close(completion.done)
 	if initFailure != nil {
 		// In standalone, we don't have to block rapid start() goroutine until init failure is consumed
@@ -547,13 +549,13 @@ func (s *Server) Init(i *interop.Init, invokeTimeoutMs int64) error {
 	return nil
 }
 
-func (s *Server) AwaitInitCompletion() time.Time {
+func (s *Server) AwaitInitCompletion() (time.Time, bool) {
 	completion := s.getInitCompletion()
 	if completion == nil {
-		return time.Time{}
+		return time.Time{}, false
 	}
 	<-completion.done
-	return completion.completedAt
+	return completion.completedAt, completion.succeeded
 }
 
 func (s *Server) FastInvoke(w http.ResponseWriter, i *interop.Invoke, direct bool) error {
diff --git a/internal/lambda/rapidcore/server_test.go b/internal/lambda/rapidcore/server_test.go
index 3f4c3e7f..c8fbff50 100644
--- a/internal/lambda/rapidcore/server_test.go
+++ b/internal/lambda/rapidcore/server_test.go
@@ -151,8 +151,9 @@ func TestAwaitInitCompletionWaitsWithoutConsumingFailure(t *testing.T) {
 	srv.Init(&interop.Init{EnvironmentVariables: env.NewEnvironment()}, int64(time.Second/time.Millisecond))
 	initCompleted := make(chan struct{})
 	var completedAt time.Time
+	var initSucceeded bool
 	go func() {
-		completedAt = srv.AwaitInitCompletion()
+		completedAt, initSucceeded = srv.AwaitInitCompletion()
 		close(initCompleted)
 	}()
 
@@ -169,11 +170,14 @@ func TestAwaitInitCompletionWaitsWithoutConsumingFailure(t *testing.T) {
 		require.Fail(t, "timed out waiting for init completion")
 	}
 	require.False(t, completedAt.IsZero())
+	require.False(t, initSucceeded)
 	require.ErrorIs(t, srv.AwaitInitialized(), ErrInitDoneFailed)
 }
 
 func TestAwaitInitCompletionBeforeInitReturnsZeroTime(t *testing.T) {
-	require.True(t, NewServer().AwaitInitCompletion().IsZero())
+	completedAt, initSucceeded := NewServer().AwaitInitCompletion()
+	require.True(t, completedAt.IsZero())
+	require.False(t, initSucceeded)
 }
 
 func TestInitErrorBeforeReserve(t *testing.T) {
diff --git a/internal/lambda/rie/handlers.go b/internal/lambda/rie/handlers.go
index 28e470a0..db7c2dc2 100644
--- a/internal/lambda/rie/handlers.go
+++ b/internal/lambda/rie/handlers.go
@@ -29,7 +29,7 @@ import (
 
 type Sandbox interface {
 	Init(i *interop.Init, invokeTimeoutMs int64)
-	AwaitInitCompletion() time.Time
+	AwaitInitCompletion() (time.Time, bool)
 	Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
 }
 
@@ -212,8 +212,12 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 			w.WriteHeader(http.StatusGatewayTimeout)
 			return
 		case rapidcore.ErrInvokeTimeout:
-			initEnd := sandbox.AwaitInitCompletion()
-			printEndReports(invokePayload.ID, formatInitDuration(initStart, initEnd, timeoutDuration), memorySize, invokeStart, timeoutDuration)
+			initEnd, initSucceeded := sandbox.AwaitInitCompletion()
+			initDuration := ""
+			if initSucceeded {
+				initDuration = formatInitDuration(initStart, initEnd, timeoutDuration)
+			}
+			printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)
 
 			w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
 			time.Sleep(100 * time.Millisecond)
@@ -222,8 +226,11 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		}
 	}
 
-	initEnd := sandbox.AwaitInitCompletion()
-	initDuration := formatInitDuration(initStart, initEnd, timeoutDuration)
+	initEnd, initSucceeded := sandbox.AwaitInitCompletion()
+	initDuration := ""
+	if initSucceeded {
+		initDuration = formatInitDuration(initStart, initEnd, timeoutDuration)
+	}
 	if !initStart.IsZero() && initEnd.After(invokeStart) {
 		invokeStart = initEnd
 	}
diff --git a/internal/lambda/rie/handlers_test.go b/internal/lambda/rie/handlers_test.go
index 3aa55b96..9d836a65 100644
--- a/internal/lambda/rie/handlers_test.go
+++ b/internal/lambda/rie/handlers_test.go
@@ -14,6 +14,7 @@ import (
 	"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"
 )
 
@@ -21,22 +22,24 @@ type delayedInitSandbox struct {
 	delay           time.Duration
 	invokeCalled    bool
 	initCompletedAt time.Time
+	initSucceeded   bool
+	invokeErr       error
 }
 
 func (s *delayedInitSandbox) Init(*interop.Init, int64) {}
 
-func (s *delayedInitSandbox) AwaitInitCompletion() time.Time {
+func (s *delayedInitSandbox) AwaitInitCompletion() (time.Time, bool) {
 	if !s.invokeCalled {
 		panic("AwaitInitCompletion called before Invoke")
 	}
-	return s.initCompletedAt
+	return s.initCompletedAt, s.initSucceeded
 }
 
 func (s *delayedInitSandbox) Invoke(http.ResponseWriter, *interop.Invoke) error {
 	s.invokeCalled = true
 	time.Sleep(s.delay)
 	s.initCompletedAt = time.Now()
-	return nil
+	return s.invokeErr
 }
 
 type panicInitSandbox struct {
@@ -86,7 +89,7 @@ func TestInvokeHandlerReportsRuntimeInitDuration(t *testing.T) {
 
 	request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil)
 	response := httptest.NewRecorder()
-	sandbox := &delayedInitSandbox{delay: 50 * time.Millisecond}
+	sandbox := &delayedInitSandbox{delay: 50 * time.Millisecond, initSucceeded: true}
 
 	reader, writer, err := os.Pipe()
 	require.NoError(t, err)
@@ -112,3 +115,39 @@ func TestInvokeHandlerReportsRuntimeInitDuration(t *testing.T) {
 	require.NoError(t, err)
 	require.Less(t, durationMilliseconds, float64(40))
 }
+
+func TestInvokeHandlerOmitsInitDurationWhenInitTimesOut(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{
+		delay:         10 * time.Millisecond,
+		invokeErr:     rapidcore.ErrInvokeTimeout,
+		initSucceeded: false,
+	}
+
+	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:")
+	require.Contains(t, string(output), "Duration:")
+}

From 8bc27a7ce0dab73108ae051c48fd08b9c35718cd Mon Sep 17 00:00:00 2001
From: Jayanth Sai Yarlagadda <72214836+JayYarlagadda@users.noreply.github.com>
Date: Mon, 3 Aug 2026 00:42:23 -0700
Subject: [PATCH 5/8] fix: preserve cold-start timeout report

---
 .../lambda/rapidcore/sandbox_emulator_api.go  |  4 +--
 internal/lambda/rapidcore/server.go           |  8 +++---
 internal/lambda/rapidcore/server_test.go      |  8 ++----
 internal/lambda/rie/handlers.go               | 16 ++++--------
 internal/lambda/rie/handlers_test.go          | 25 +++++++++++--------
 5 files changed, 27 insertions(+), 34 deletions(-)

diff --git a/internal/lambda/rapidcore/sandbox_emulator_api.go b/internal/lambda/rapidcore/sandbox_emulator_api.go
index 7d9d86cd..a6efea1c 100644
--- a/internal/lambda/rapidcore/sandbox_emulator_api.go
+++ b/internal/lambda/rapidcore/sandbox_emulator_api.go
@@ -13,7 +13,7 @@ import (
 // LambdaInvokeAPI are the methods used by the Runtime Interface Emulator
 type LambdaInvokeAPI interface {
 	Init(i *interop.Init, invokeTimeoutMs int64)
-	AwaitInitCompletion() (time.Time, bool)
+	AwaitInitCompletion() time.Time
 	Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
 }
 
@@ -49,7 +49,7 @@ func (l *EmulatorAPI) Init(i *interop.Init, timeoutMs int64) {
 	}, timeoutMs)
 }
 
-func (l *EmulatorAPI) AwaitInitCompletion() (time.Time, bool) {
+func (l *EmulatorAPI) AwaitInitCompletion() time.Time {
 	return l.server.AwaitInitCompletion()
 }
 
diff --git a/internal/lambda/rapidcore/server.go b/internal/lambda/rapidcore/server.go
index 94608950..f3c159ce 100644
--- a/internal/lambda/rapidcore/server.go
+++ b/internal/lambda/rapidcore/server.go
@@ -75,7 +75,6 @@ type InvokeContext struct {
 type initCompletion struct {
 	done        chan struct{}
 	completedAt time.Time
-	succeeded   bool
 }
 
 type Server struct {
@@ -221,7 +220,6 @@ func (s *Server) Reserve(id string, traceID, lambdaSegmentID string) (*ReserveRe
 func (s *Server) awaitInitCompletion(initContext interop.InitContext, initFailures chan interop.InitFailure, completion *initCompletion) {
 	initSuccess, initFailure := initContext.Wait()
 	completion.completedAt = time.Now()
-	completion.succeeded = initFailure == nil
 	close(completion.done)
 	if initFailure != nil {
 		// In standalone, we don't have to block rapid start() goroutine until init failure is consumed
@@ -549,13 +547,13 @@ func (s *Server) Init(i *interop.Init, invokeTimeoutMs int64) error {
 	return nil
 }
 
-func (s *Server) AwaitInitCompletion() (time.Time, bool) {
+func (s *Server) AwaitInitCompletion() time.Time {
 	completion := s.getInitCompletion()
 	if completion == nil {
-		return time.Time{}, false
+		return time.Time{}
 	}
 	<-completion.done
-	return completion.completedAt, completion.succeeded
+	return completion.completedAt
 }
 
 func (s *Server) FastInvoke(w http.ResponseWriter, i *interop.Invoke, direct bool) error {
diff --git a/internal/lambda/rapidcore/server_test.go b/internal/lambda/rapidcore/server_test.go
index c8fbff50..3f4c3e7f 100644
--- a/internal/lambda/rapidcore/server_test.go
+++ b/internal/lambda/rapidcore/server_test.go
@@ -151,9 +151,8 @@ func TestAwaitInitCompletionWaitsWithoutConsumingFailure(t *testing.T) {
 	srv.Init(&interop.Init{EnvironmentVariables: env.NewEnvironment()}, int64(time.Second/time.Millisecond))
 	initCompleted := make(chan struct{})
 	var completedAt time.Time
-	var initSucceeded bool
 	go func() {
-		completedAt, initSucceeded = srv.AwaitInitCompletion()
+		completedAt = srv.AwaitInitCompletion()
 		close(initCompleted)
 	}()
 
@@ -170,14 +169,11 @@ func TestAwaitInitCompletionWaitsWithoutConsumingFailure(t *testing.T) {
 		require.Fail(t, "timed out waiting for init completion")
 	}
 	require.False(t, completedAt.IsZero())
-	require.False(t, initSucceeded)
 	require.ErrorIs(t, srv.AwaitInitialized(), ErrInitDoneFailed)
 }
 
 func TestAwaitInitCompletionBeforeInitReturnsZeroTime(t *testing.T) {
-	completedAt, initSucceeded := NewServer().AwaitInitCompletion()
-	require.True(t, completedAt.IsZero())
-	require.False(t, initSucceeded)
+	require.True(t, NewServer().AwaitInitCompletion().IsZero())
 }
 
 func TestInitErrorBeforeReserve(t *testing.T) {
diff --git a/internal/lambda/rie/handlers.go b/internal/lambda/rie/handlers.go
index db7c2dc2..45913309 100644
--- a/internal/lambda/rie/handlers.go
+++ b/internal/lambda/rie/handlers.go
@@ -29,7 +29,7 @@ import (
 
 type Sandbox interface {
 	Init(i *interop.Init, invokeTimeoutMs int64)
-	AwaitInitCompletion() (time.Time, bool)
+	AwaitInitCompletion() time.Time
 	Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error
 }
 
@@ -212,11 +212,8 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 			w.WriteHeader(http.StatusGatewayTimeout)
 			return
 		case rapidcore.ErrInvokeTimeout:
-			initEnd, initSucceeded := sandbox.AwaitInitCompletion()
-			initDuration := ""
-			if initSucceeded {
-				initDuration = formatInitDuration(initStart, initEnd, timeoutDuration)
-			}
+			initEnd := sandbox.AwaitInitCompletion()
+			initDuration := formatInitDuration(initStart, initEnd, timeoutDuration)
 			printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)
 
 			w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
@@ -226,11 +223,8 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		}
 	}
 
-	initEnd, initSucceeded := sandbox.AwaitInitCompletion()
-	initDuration := ""
-	if initSucceeded {
-		initDuration = formatInitDuration(initStart, initEnd, timeoutDuration)
-	}
+	initEnd := sandbox.AwaitInitCompletion()
+	initDuration := formatInitDuration(initStart, initEnd, timeoutDuration)
 	if !initStart.IsZero() && initEnd.After(invokeStart) {
 		invokeStart = initEnd
 	}
diff --git a/internal/lambda/rie/handlers_test.go b/internal/lambda/rie/handlers_test.go
index 9d836a65..963d45c6 100644
--- a/internal/lambda/rie/handlers_test.go
+++ b/internal/lambda/rie/handlers_test.go
@@ -22,17 +22,16 @@ type delayedInitSandbox struct {
 	delay           time.Duration
 	invokeCalled    bool
 	initCompletedAt time.Time
-	initSucceeded   bool
 	invokeErr       error
 }
 
 func (s *delayedInitSandbox) Init(*interop.Init, int64) {}
 
-func (s *delayedInitSandbox) AwaitInitCompletion() (time.Time, bool) {
+func (s *delayedInitSandbox) AwaitInitCompletion() time.Time {
 	if !s.invokeCalled {
 		panic("AwaitInitCompletion called before Invoke")
 	}
-	return s.initCompletedAt, s.initSucceeded
+	return s.initCompletedAt
 }
 
 func (s *delayedInitSandbox) Invoke(http.ResponseWriter, *interop.Invoke) error {
@@ -89,7 +88,7 @@ func TestInvokeHandlerReportsRuntimeInitDuration(t *testing.T) {
 
 	request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil)
 	response := httptest.NewRecorder()
-	sandbox := &delayedInitSandbox{delay: 50 * time.Millisecond, initSucceeded: true}
+	sandbox := &delayedInitSandbox{delay: 50 * time.Millisecond}
 
 	reader, writer, err := os.Pipe()
 	require.NoError(t, err)
@@ -116,7 +115,7 @@ func TestInvokeHandlerReportsRuntimeInitDuration(t *testing.T) {
 	require.Less(t, durationMilliseconds, float64(40))
 }
 
-func TestInvokeHandlerOmitsInitDurationWhenInitTimesOut(t *testing.T) {
+func TestInvokeHandlerReportsInitDurationWhenInitTimesOut(t *testing.T) {
 	initMutex.Lock()
 	initDone = false
 	initMutex.Unlock()
@@ -130,9 +129,8 @@ func TestInvokeHandlerOmitsInitDurationWhenInitTimesOut(t *testing.T) {
 	request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil)
 	response := httptest.NewRecorder()
 	sandbox := &delayedInitSandbox{
-		delay:         10 * time.Millisecond,
-		invokeErr:     rapidcore.ErrInvokeTimeout,
-		initSucceeded: false,
+		delay:     10 * time.Millisecond,
+		invokeErr: rapidcore.ErrInvokeTimeout,
 	}
 
 	reader, writer, err := os.Pipe()
@@ -148,6 +146,13 @@ func TestInvokeHandlerOmitsInitDurationWhenInitTimesOut(t *testing.T) {
 	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:")
-	require.Contains(t, string(output), "Duration:")
+	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)
 }

From fe7881b7b7ef32869f9631b6de61bec9fa4224a7 Mon Sep 17 00:00:00 2001
From: Jayanth Sai Yarlagadda <72214836+JayYarlagadda@users.noreply.github.com>
Date: Mon, 3 Aug 2026 01:30:38 -0700
Subject: [PATCH 6/8] fix: avoid overlapping timeout durations

---
 internal/lambda/rie/handlers.go      |   4 +-
 internal/lambda/rie/handlers_test.go | 100 +++++++++++++++++++++++++--
 2 files changed, 99 insertions(+), 5 deletions(-)

diff --git a/internal/lambda/rie/handlers.go b/internal/lambda/rie/handlers.go
index 45913309..0476eff0 100644
--- a/internal/lambda/rie/handlers.go
+++ b/internal/lambda/rie/handlers.go
@@ -101,7 +101,6 @@ func formatInitDuration(initStart time.Time, initEnd time.Time, timeoutDuration
 		float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond)
 	return fmt.Sprintf("Init Duration: %.2f ms\t", initTimeMS)
 }
-
 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)
@@ -214,6 +213,9 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		case rapidcore.ErrInvokeTimeout:
 			initEnd := sandbox.AwaitInitCompletion()
 			initDuration := formatInitDuration(initStart, initEnd, timeoutDuration)
+			if !initStart.IsZero() && initEnd.After(invokeStart) {
+				invokeStart = initEnd
+			}
 			printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)
 
 			w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
diff --git a/internal/lambda/rie/handlers_test.go b/internal/lambda/rie/handlers_test.go
index 963d45c6..c89cd4ab 100644
--- a/internal/lambda/rie/handlers_test.go
+++ b/internal/lambda/rie/handlers_test.go
@@ -19,7 +19,8 @@ import (
 )
 
 type delayedInitSandbox struct {
-	delay           time.Duration
+	initDelay       time.Duration
+	invokeDelay     time.Duration
 	invokeCalled    bool
 	initCompletedAt time.Time
 	invokeErr       error
@@ -36,8 +37,9 @@ func (s *delayedInitSandbox) AwaitInitCompletion() time.Time {
 
 func (s *delayedInitSandbox) Invoke(http.ResponseWriter, *interop.Invoke) error {
 	s.invokeCalled = true
-	time.Sleep(s.delay)
+	time.Sleep(s.initDelay)
 	s.initCompletedAt = time.Now()
+	time.Sleep(s.invokeDelay)
 	return s.invokeErr
 }
 
@@ -88,7 +90,7 @@ func TestInvokeHandlerReportsRuntimeInitDuration(t *testing.T) {
 
 	request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil)
 	response := httptest.NewRecorder()
-	sandbox := &delayedInitSandbox{delay: 50 * time.Millisecond}
+	sandbox := &delayedInitSandbox{initDelay: 50 * time.Millisecond}
 
 	reader, writer, err := os.Pipe()
 	require.NoError(t, err)
@@ -129,7 +131,7 @@ func TestInvokeHandlerReportsInitDurationWhenInitTimesOut(t *testing.T) {
 	request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil)
 	response := httptest.NewRecorder()
 	sandbox := &delayedInitSandbox{
-		delay:     10 * time.Millisecond,
+		initDelay: 50 * time.Millisecond,
 		invokeErr: rapidcore.ErrInvokeTimeout,
 	}
 
@@ -155,4 +157,94 @@ func TestInvokeHandlerReportsInitDurationWhenInitTimesOut(t *testing.T) {
 
 	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))
+	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:   50 * time.Millisecond,
+		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())
+	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(40))
+	require.Less(t, durationMilliseconds, float64(90))
+	require.LessOrEqual(t, initDurationMilliseconds+durationMilliseconds, float64(1020))
+}
+
+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))
 }

From 8c18e3591cdf9b3961e41744539f2ae7f4f8f8ad Mon Sep 17 00:00:00 2001
From: Jayanth Sai Yarlagadda <72214836+JayYarlagadda@users.noreply.github.com>
Date: Mon, 3 Aug 2026 11:05:16 -0700
Subject: [PATCH 7/8] refactor: centralize invoke report timing

---
 internal/lambda/rie/handlers.go      | 24 ++++++++++++------------
 internal/lambda/rie/handlers_test.go |  4 ++--
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/internal/lambda/rie/handlers.go b/internal/lambda/rie/handlers.go
index 0476eff0..12525fcf 100644
--- a/internal/lambda/rie/handlers.go
+++ b/internal/lambda/rie/handlers.go
@@ -101,6 +101,16 @@ func formatInitDuration(initStart time.Time, initEnd time.Time, timeoutDuration
 		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)
@@ -211,12 +221,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 			w.WriteHeader(http.StatusGatewayTimeout)
 			return
 		case rapidcore.ErrInvokeTimeout:
-			initEnd := sandbox.AwaitInitCompletion()
-			initDuration := formatInitDuration(initStart, initEnd, timeoutDuration)
-			if !initStart.IsZero() && initEnd.After(invokeStart) {
-				invokeStart = initEnd
-			}
-			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)
@@ -225,12 +230,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		}
 	}
 
-	initEnd := sandbox.AwaitInitCompletion()
-	initDuration := formatInitDuration(initStart, initEnd, timeoutDuration)
-	if !initStart.IsZero() && initEnd.After(invokeStart) {
-		invokeStart = initEnd
-	}
-	printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)
+	printInvokeReport(sandbox, invokePayload.ID, initStart, invokeStart, memorySize, timeoutDuration)
 
 	if invokeResp.StatusCode != 0 {
 		w.WriteHeader(invokeResp.StatusCode)
diff --git a/internal/lambda/rie/handlers_test.go b/internal/lambda/rie/handlers_test.go
index c89cd4ab..5e89425b 100644
--- a/internal/lambda/rie/handlers_test.go
+++ b/internal/lambda/rie/handlers_test.go
@@ -159,7 +159,7 @@ func TestInvokeHandlerReportsInitDurationWhenInitTimesOut(t *testing.T) {
 	require.Len(t, matches, 2)
 	durationMilliseconds, err := strconv.ParseFloat(matches[1], 64)
 	require.NoError(t, err)
-	require.Less(t, durationMilliseconds, float64(40))
+	require.Less(t, durationMilliseconds, initDurationMilliseconds)
 	require.LessOrEqual(t, initDurationMilliseconds+durationMilliseconds, float64(1020))
 }
 
@@ -206,7 +206,7 @@ func TestInvokeHandlerSeparatesInitFromTimedOutInvocation(t *testing.T) {
 	durationMilliseconds, err := strconv.ParseFloat(durationMatches[1], 64)
 	require.NoError(t, err)
 	require.GreaterOrEqual(t, durationMilliseconds, float64(40))
-	require.Less(t, durationMilliseconds, float64(90))
+	require.Less(t, durationMilliseconds, initDurationMilliseconds*1.5)
 	require.LessOrEqual(t, initDurationMilliseconds+durationMilliseconds, float64(1020))
 }
 

From 427bcfb3793c08f340c214fcf94261836749f0ef Mon Sep 17 00:00:00 2001
From: Jayanth Sai Yarlagadda <72214836+JayYarlagadda@users.noreply.github.com>
Date: Mon, 3 Aug 2026 12:06:55 -0700
Subject: [PATCH 8/8] test: assert timeout intervals against elapsed time

---
 internal/lambda/rie/handlers_test.go | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/internal/lambda/rie/handlers_test.go b/internal/lambda/rie/handlers_test.go
index 5e89425b..773b5b76 100644
--- a/internal/lambda/rie/handlers_test.go
+++ b/internal/lambda/rie/handlers_test.go
@@ -177,8 +177,8 @@ func TestInvokeHandlerSeparatesInitFromTimedOutInvocation(t *testing.T) {
 	request := httptest.NewRequest(http.MethodPost, "/2015-03-31/functions/function/invocations", nil)
 	response := httptest.NewRecorder()
 	sandbox := &delayedInitSandbox{
-		initDelay:   50 * time.Millisecond,
-		invokeDelay: 50 * time.Millisecond,
+		initDelay:   200 * time.Millisecond,
+		invokeDelay: 200 * time.Millisecond,
 		invokeErr:   rapidcore.ErrInvokeTimeout,
 	}
 
@@ -188,7 +188,9 @@ func TestInvokeHandlerSeparatesInitFromTimedOutInvocation(t *testing.T) {
 	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)
@@ -205,9 +207,8 @@ func TestInvokeHandlerSeparatesInitFromTimedOutInvocation(t *testing.T) {
 	require.Len(t, durationMatches, 2)
 	durationMilliseconds, err := strconv.ParseFloat(durationMatches[1], 64)
 	require.NoError(t, err)
-	require.GreaterOrEqual(t, durationMilliseconds, float64(40))
-	require.Less(t, durationMilliseconds, initDurationMilliseconds*1.5)
-	require.LessOrEqual(t, initDurationMilliseconds+durationMilliseconds, float64(1020))
+	require.GreaterOrEqual(t, durationMilliseconds, float64(190))
+	require.LessOrEqual(t, initDurationMilliseconds+durationMilliseconds, elapsedMilliseconds)
 }
 
 func TestInvokeHandlerReportsWarmTimeoutWithoutInitDuration(t *testing.T) {