chore: sync from internal (rie-private-to-public)

.github/workflows/validate-branch-into-main.yaml

@@ -10,8 +10,9 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Check source branch
+        env:
+          SOURCE_BRANCH: ${{ github.head_ref }}
         run: |
-          SOURCE_BRANCH="${{ github.head_ref }}"
           if [[ "$SOURCE_BRANCH" != "develop" ]]; then
             echo "Error: Only pull requests from develop branch are allowed into main"
             echo "Current source branch ($SOURCE_BRANCH)."

internal/lambda-managed-instances/aws-lambda-rie/internal/app_test.go

@@ -79,7 +79,7 @@ func TestApp_ServeHTTP(t *testing.T) {
 			defer mockApp.AssertExpectations(t)
 			mockApp.On("Init", mock.Anything, mock.Anything, mock.Anything).Return(tt.initResponse)
 			if tt.initResponse == nil {
-				mockApp.On("Invoke", mock.Anything, mock.Anything, mock.Anything).Return(tt.invokeErr, tt.responseSent)
+				mockApp.On("Invoke", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(tt.invokeErr, tt.responseSent, false)
 			}
 
 			initMsg := intmodel.InitRequestMessage{

internal/lambda-managed-instances/aws-lambda-rie/internal/run.go

@@ -75,7 +75,7 @@ func Run(supv supvmodel.ProcessSupervisor, args []string, fileUtil utils.FileUti
 	}
 
 	rieApp := NewHTTPHandler(raptorApp, initMsg)
-	s, err := raptor.StartServer(raptorApp, rieApp, &raptor.TCPAddress{AddrPort: rieAddr})
+	s, err := raptor.StartServer(raptorApp, rieApp, &raptor.TCPAddress{AddrPort: rieAddr}, false)
 	if err != nil {
 		return nil, nil, nil, fmt.Errorf("could not start RIE server: %w", err)
 	}

internal/lambda-managed-instances/interop/sandbox_model.go

@@ -253,6 +253,7 @@ type EEStaticData struct {
 	XRayTracingMode    intmodel.XrayTracingMode
 	ArtefactType       intmodel.ArtefactType
 	RuntimeVersion     string
+	RuntimeRelease     string
 	AmiId              string
 	AvailabilityZoneId string
 }

internal/lambda-managed-instances/interop/service_log_values.go

@@ -27,6 +27,10 @@ const (
 	ShutdownWaitAllProcessesDuration          = "WaitCustomerProcessesExitDuration"
 	ShutdownRuntimeServerDuration             = "StopRuntimeServerDuration"
 
+	RuntimeNextCountMetric = "RuntimeNextCount"
+
+	RuntimeWorkerCountMetric = "RuntimeWorkerCount"
+
 	ClientErrorMetric           = "ClientError"
 	ClientErrorReasonTemplate   = "ClientErrorReason-%s"
 	CustomerErrorMetric         = "CustomerError"

internal/lambda-managed-instances/invoke/invoke_router.go

@@ -195,6 +195,12 @@ func (ir *InvokeRouter) GetRuntimePoolCounts() RuntimePoolCounts {
 func (ir *InvokeRouter) ReserveIdleRuntime(ctx context.Context, invokeID interop.InvokeID, timeout time.Duration) (interop.ReserveIdleRuntimeResponse, model.AppError) {
 	logging.Debug(ctx, "InvokeRouter: reserving idle runtime")
 
+	if _, exists := ir.runningInvokes.Get(invokeID); exists {
+		logging.Warn(ctx, "InvokeRouter: reservation collides with in-flight invoke")
+		return interop.ReserveIdleRuntimeFailureResponse{ErrorType: model.ErrorDuplicatedInvokeId},
+			model.NewClientError(ErrInvokeIdAlreadyExists, model.ErrorSeverityError, model.ErrorDuplicatedInvokeId)
+	}
+
 	err := ir.runtimePool.Reserve(invokeID, timeout, func() {
 		if ir.runtimePool.ExpireReservation(invokeID) {
 			logging.Info(ctx, "InvokeRouter: reservation expired")

internal/lambda-managed-instances/invoke/invoke_router_test.go

@@ -399,6 +399,29 @@ func TestReserveIdleRuntime_DuplicateInvokeID(t *testing.T) {
 	assert.Equal(t, model.ErrorDuplicatedInvokeId, appErr.ErrorType())
 }
 
+func TestReserveIdleRuntime_DuplicateAgainstRunningInvoke(t *testing.T) {
+	t.Parallel()
+
+	mocks, router := createMocksAndInitRouter()
+
+	_, err := router.RuntimeNext(mocks.ctx, mocks.runtimeNextRequest)
+	require.NoError(t, err)
+
+	const inFlightID = "reserve-vs-running"
+	router.runningInvokes.Set(inFlightID, &mocks.runnningInvoke)
+
+	resp, appErr := router.ReserveIdleRuntime(mocks.ctx, inFlightID, 100*time.Millisecond)
+
+	require.NotNil(t, appErr)
+	failResp, ok := resp.(interop.ReserveIdleRuntimeFailureResponse)
+	assert.True(t, ok, "expected ReserveIdleRuntimeFailureResponse")
+	assert.Equal(t, model.ErrorDuplicatedInvokeId, failResp.ErrorType)
+	assert.Equal(t, model.ErrorDuplicatedInvokeId, appErr.ErrorType())
+
+	assert.Equal(t, 0, router.runtimePool.ReservedCount(), "no reservation should have been recorded")
+	assert.Equal(t, 1, router.GetRuntimePoolCounts().Idle, "idle runtime should still be available")
+}
+
 func TestReserveIdleRuntime_Expiration(t *testing.T) {
 	t.Parallel()
 

internal/lambda-managed-instances/invoke/reservation_metrics.go

@@ -13,11 +13,13 @@ import (
 )
 
 const (
-	ReserveSuccessMetric = "ReserveSuccess"
-	ReserveFailedMetric  = "ReserveFailed"
+	ReserveSuccessMetric       = "ReserveSuccess"
+	ReserveFailedMetric        = "ReserveFailed"
+	ReserveParseDurationMetric = "ReserveParseDuration"
+	ReserveLogicDurationMetric = "ReserveLogicDuration"
 )
 
-func ReserveServiceLog(logger servicelogs.Logger, opStart time.Time, invokeID string, appErr model.AppError) {
+func ReserveServiceLog(logger servicelogs.Logger, opStart time.Time, invokeID string, appErr model.AppError, parseDuration, reserveDuration time.Duration) {
 	props := []servicelogs.Property{
 		{Name: "invoke_id", Value: invokeID},
 	}
@@ -43,6 +45,11 @@ func ReserveServiceLog(logger servicelogs.Logger, opStart time.Time, invokeID st
 		servicelogs.Counter(ReserveFailedMetric, failed),
 		servicelogs.Counter(interop.ClientErrorMetric, clientErrCnt),
 		servicelogs.Counter(interop.NonCustomerErrorMetric, nonCustomerErrCnt),
+		servicelogs.Timer(ReserveParseDurationMetric, parseDuration),
+	}
+
+	if reserveDuration > 0 {
+		metrics = append(metrics, servicelogs.Timer(ReserveLogicDurationMetric, reserveDuration))
 	}
 
 	if appErr != nil {

internal/lambda-managed-instances/invoke/runtime_pool.go

@@ -5,6 +5,7 @@ package invoke
 
 import (
 	"errors"
+	"log/slog"
 	"sync"
 	"time"
 
@@ -44,9 +45,14 @@ func (p *RuntimePool) Add(runtime runningInvoke) error {
 }
 
 func (p *RuntimePool) Reserve(invokeID interop.InvokeID, timeout time.Duration, onExpire func()) error {
+	lockStart := time.Now()
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
+	if lockWait := time.Since(lockStart); lockWait > time.Millisecond {
+		slog.Warn("RuntimePool.Reserve lock contention", "wait_us", lockWait.Microseconds())
+	}
+
 	if _, exists := p.reserved[invokeID]; exists {
 		return ErrInvokeIdAlreadyExists
 	}

internal/lambda-managed-instances/invoke/runtime_response_request.go

@@ -9,6 +9,7 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
+	"time"
 
 	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/interop"
 	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/logging"
@@ -23,6 +24,7 @@ const (
 
 type runtimeResponse struct {
 	request *http.Request
+	rc      *http.ResponseController
 
 	parsingErr model.AppError
 
@@ -31,14 +33,15 @@ type runtimeResponse struct {
 	responseMode string
 }
 
-func NewRuntimeResponse(ctx context.Context, request *http.Request, invokeID interop.InvokeID) runtimeResponse {
+func NewRuntimeResponse(ctx context.Context, request *http.Request, writer http.ResponseWriter, invokeID interop.InvokeID) runtimeResponse {
 	contentType := request.Header.Get(RuntimeContentTypeHeader)
 	if contentType == "" {
 
 		contentType = "application/octet-stream"
 	}
 	resp := runtimeResponse{
 		request:     request,
+		rc:          http.NewResponseController(writer),
 		contentType: contentType,
 		invokeID:    invokeID,
 	}
@@ -74,6 +77,12 @@ func (r *runtimeResponse) BodyReader() io.Reader {
 	return r.request.Body
 }
 
+func (r *runtimeResponse) Cancel() {
+	if err := r.rc.SetReadDeadline(time.Unix(0, 0)); err != nil {
+		slog.Warn("Cancel: SetReadDeadline failed", "err", err)
+	}
+}
+
 func (r *runtimeResponse) ResponseMode() string {
 	return r.responseMode
 }

internal/lambda-managed-instances/invoke/runtime_response_request_test.go

@@ -5,6 +5,7 @@ package invoke
 
 import (
 	"net/http"
+	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -66,7 +67,7 @@ func TestRuntimeResponse_TrailerError(t *testing.T) {
 			req.Trailer = make(http.Header)
 			req.Trailer.Set(FunctionErrorTypeTrailer, tc.errorTypeTrailer)
 			req.Trailer.Set(FunctionErrorBodyTrailer, tc.errorBodyTrailer)
-			resp := NewRuntimeResponse(req.Context(), req, "test-invoke-id")
+			resp := NewRuntimeResponse(req.Context(), req, httptest.NewRecorder(), "test-invoke-id")
 
 			actualTrailerError := resp.TrailerError()
 

internal/lambda-managed-instances/invoke/runtime_response_sender.go

@@ -33,9 +33,15 @@ func sendInvokeToRuntime(ctx context.Context, initData interop.InitStaticDataPro
 	runtimeReq.Header().Set(RuntimeRequestIdHeader, invokeReq.InvokeID())
 	runtimeReq.Header().Set(RuntimeDeadlineHeader, strconv.FormatInt(invokeReq.Deadline().UnixMilli(), 10))
 	runtimeReq.Header().Set(RuntimeFunctionArnHeader, initData.FunctionARN())
-	runtimeReq.Header().Set(RuntimeTraceIdHeader, traceId)
-	runtimeReq.Header().Set(RuntimeClientContextHeader, invokeReq.ClientContext())
-	runtimeReq.Header().Set(RuntimeCognitoIdentifyHeader, buildCognitoIdentifyHeader(invokeReq))
+	if traceId != "" {
+		runtimeReq.Header().Set(RuntimeTraceIdHeader, traceId)
+	}
+	if cc := invokeReq.ClientContext(); cc != "" {
+		runtimeReq.Header().Set(RuntimeClientContextHeader, cc)
+	}
+	if cogId := buildCognitoIdentifyHeader(invokeReq); cogId != "" {
+		runtimeReq.Header().Set(RuntimeCognitoIdentifyHeader, cogId)
+	}
 	runtimeReq.WriteHeader(http.StatusOK)
 
 	timedReader := &utils.TimedReader{

internal/lambda-managed-instances/invoke/runtime_response_sender_test.go

@@ -134,3 +134,36 @@ func TestSendResponseFailure_CtxCancelled(t *testing.T) {
 
 	checkResponseSenderExpectations(t, mocks)
 }
+
+func TestSendResponse_OmitsEmptyOptionalHeaders(t *testing.T) {
+	t.Parallel()
+
+	mocks := createMocksAndRuntimeResponder()
+	buildInitDataMocks(&mocks.initData)
+
+	mocks.invokeReq.On("ContentType").Return("application/json")
+	mocks.invokeReq.On("InvokeID").Return("123456")
+	mocks.invokeReq.On("Deadline").Return(time.Now().Add(time.Second))
+	mocks.invokeReq.On("ClientContext").Return("")
+	mocks.invokeReq.On("CognitoId").Return("")
+	mocks.invokeReq.On("CognitoPoolId").Return("")
+	mocks.invokeReq.On("BodyReader").Return(mocks.reader)
+
+	recorder := httptest.NewRecorder()
+	mocks.runtimeReq = recorder
+
+	_, _, _, err := sendInvokeToRuntime(mocks.ctx, &mocks.initData, &mocks.invokeReq, mocks.runtimeReq, "")
+	assert.NoError(t, err)
+
+	headers := recorder.Header()
+	assert.Empty(t, headers.Get(RuntimeTraceIdHeader))
+	assert.Empty(t, headers.Get(RuntimeClientContextHeader))
+	assert.Empty(t, headers.Get(RuntimeCognitoIdentifyHeader))
+
+	assert.NotEmpty(t, headers.Get(RuntimeRequestIdHeader))
+	assert.NotEmpty(t, headers.Get(RuntimeDeadlineHeader))
+	assert.NotEmpty(t, headers.Get(RuntimeFunctionArnHeader))
+	assert.NotEmpty(t, headers.Get(RuntimeContentTypeHeader))
+
+	checkResponseSenderExpectations(t, mocks)
+}

internal/lambda-managed-instances/rapi/handler/invocationresponse.go

@@ -28,7 +28,7 @@ func (h *invocationResponseHandler) ServeHTTP(writer http.ResponseWriter, reques
 	ctx := logging.WithInvokeID(request.Context(), invokeID)
 
 	logging.Debug(ctx, "Received Runtime Response")
-	resp := invoke.NewRuntimeResponse(ctx, request, invokeID)
+	resp := invoke.NewRuntimeResponse(ctx, request, writer, invokeID)
 
 	err := h.runtimeRespHandler.RuntimeResponse(ctx, &resp)
 	if err == nil {

internal/lambda-managed-instances/rapid/handlers.go

@@ -346,6 +346,7 @@ func doInitRuntime(
 		return customerErr
 	}
 
+	execCtx.initExecutionData.StaticData.RuntimeRelease = appctx.GetRuntimeRelease(execCtx.appCtx)
 	telemetry.SendInitRuntimeDoneLogEvent(execCtx.eventsAPI, execCtx.appCtx, phase, nil)
 
 	return nil