fix: separate init and invoke durations

Author: JayYarlagadda

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)

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