refactor: centralize invoke report timing

Author: JayYarlagadda

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)

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