fix: address review findings on the Telemetry API implementation

Four problems, all reachable in normal use:

Delivery used http.DefaultClient, which has no timeout. Batches that reach a
size limit are delivered on the goroutine producing the event, so a subscriber
that accepted a connection and never answered would stall the sandbox's event
pipeline. Delivery now uses a client that gives up.

platform.report was emitted with status success unconditionally, including from
the invoke-timeout path, so a timed-out invocation was reported as a successful
one -- the opposite of what an extension watching that field needs. The status
is now threaded through.

The report's memory metrics were serialized as JSON strings, because the memory
size arrives as one. The Telemetry API defines them as numbers, and a consumer
decoding into a typed struct rejects strings. They are parsed now, and the
report also carries initDurationMs on a cold start, as a real function's does.

Dispatch decided whether to buffer or deliver in one lock section and acted in
another. A Subscribe interleaving between the two would insert itself, replay a
copy of the buffer that did not yet hold the event, and leave the event buffered
for nobody -- precisely the ordering the buffer exists to handle. The decision
now happens under a single held lock.

Tests cover the timeout, the status, the numeric metrics and the cold-start
figure. The lost-event window is not covered: the two lock sections it needed to
interleave between were adjacent and 400 attempts never hit it, so that fix
rests on reading the code rather than on a failing test.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: lizthegrey

internal/lambda/rie/handlers.go

@@ -68,30 +68,52 @@ func GetenvWithDefault(key string, defaultValue string) string {
 	return envValue
 }
 
-func printEndReports(invokeId string, initDuration string, memorySize string, invokeStart time.Time, timeoutDuration time.Duration) {
+// reportRecord builds the platform.report record. The metrics are numbers, as the
+// Telemetry API defines them: a consumer decoding into a typed struct rejects the
+// strings these values arrive as.
+func reportRecord(invokeId string, status string, invokeDuration float64, memorySize string, initDurationMs float64) map[string]interface{} {
+	// The emulator cannot measure memory actually used, so it reports the
+	// configured size for both, as the printed REPORT line does.
+	memorySizeMB, err := strconv.Atoi(memorySize)
+	if err != nil {
+		log.Warnf("AWS_LAMBDA_FUNCTION_MEMORY_SIZE is %q, which is not a number", memorySize)
+	}
+
+	metrics := map[string]interface{}{
+		"durationMs":       invokeDuration,
+		"billedDurationMs": math.Ceil(invokeDuration),
+		"memorySizeMB":     memorySizeMB,
+		"maxMemoryUsedMB":  memorySizeMB,
+	}
+	// Present only on the report for an invocation that initialized the
+	// environment, as it is in telemetry from a real function.
+	if initDurationMs > 0 {
+		metrics["initDurationMs"] = initDurationMs
+	}
+
+	return map[string]interface{}{
+		"requestId": invokeId,
+		"status":    status,
+		"metrics":   metrics,
+	}
+}
+
+// printEndReports reports the end of an invocation. status is what the Telemetry
+// API calls it: "success", or "timeout" when the invocation ran out of time.
+func printEndReports(invokeId string, initDuration string, memorySize string, invokeStart time.Time, timeoutDuration time.Duration, status string, initDurationMs float64) {
 	// Calcuation invoke duration
 	invokeDuration := math.Min(float64(time.Now().Sub(invokeStart).Nanoseconds()),
 		float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond)
 
 	fmt.Println("END RequestId: " + invokeId)
 
 	if telemetryEvents != nil {
-		now := time.Now().Format(time.RFC3339)
 		// platform.end is deliberately not emitted: it is absent from telemetry
 		// captured off a real function under the current schema version.
 		telemetryEvents.Dispatch(standalonetelemetry.SandboxEvent{
-			Time: now,
-			Type: "platform.report",
-			PlatformEvent: map[string]interface{}{
-				"requestId": invokeId,
-				"status":    "success",
-				"metrics": map[string]interface{}{
-					"durationMs":       invokeDuration,
-					"billedDurationMs": math.Ceil(invokeDuration),
-					"memorySizeMB":     memorySize,
-					"maxMemoryUsedMB":  memorySize,
-				},
-			},
+			Time:          time.Now().Format(time.RFC3339),
+			Type:          "platform.report",
+			PlatformEvent: reportRecord(invokeId, status, invokeDuration, memorySize, initDurationMs),
 		})
 	}
 	// We set the Max Memory Used and Memory Size to be the same (whatever it is set to) since there is
@@ -123,6 +145,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 	}
 
 	initDuration := ""
+	initDurationMs := float64(0)
 	inv := GetenvWithDefault("AWS_LAMBDA_FUNCTION_TIMEOUT", "300")
 	timeoutDuration, _ := time.ParseDuration(inv + "s")
 	// Default
@@ -143,6 +166,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 			float64(timeoutDuration.Nanoseconds())) / float64(time.Millisecond)
 
 		initDuration = fmt.Sprintf("Init Duration: %.2f ms\t", initTimeMS)
+		initDurationMs = initTimeMS
 
 		// Set initDone so next invokes do not try to Init the function again
 		initDone = true
@@ -229,7 +253,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, initDuration, memorySize, invokeStart, timeoutDuration, "timeout", initDurationMs)
 
 			w.Write([]byte(fmt.Sprintf("Task timed out after %d.00 seconds", timeout)))
 			time.Sleep(100 * time.Millisecond)
@@ -238,7 +262,7 @@ func InvokeHandler(w http.ResponseWriter, r *http.Request, sandbox Sandbox, bs i
 		}
 	}
 
-	printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration)
+	printEndReports(invokePayload.ID, initDuration, memorySize, invokeStart, timeoutDuration, "success", initDurationMs)
 
 	if invokeResp.StatusCode != 0 {
 		w.WriteHeader(invokeResp.StatusCode)

internal/lambda/rie/handlers_report_test.go

@@ -0,0 +1,66 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package rie
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// The Telemetry API defines these metrics as numbers. A consumer decoding into a
+// typed struct rejects the strings the memory size arrives as, so the record has
+// to carry numbers however the value reached us.
+func TestReportRecordMetricsAreNumbers(t *testing.T) {
+	encoded, err := json.Marshal(reportRecord("abc", "success", 12.5, "512", 0))
+	require.NoError(t, err)
+
+	var decoded struct {
+		RequestID string `json:"requestId"`
+		Status    string `json:"status"`
+		Metrics   struct {
+			DurationMs       float64 `json:"durationMs"`
+			BilledDurationMs float64 `json:"billedDurationMs"`
+			MemorySizeMB     int     `json:"memorySizeMB"`
+			MaxMemoryUsedMB  int     `json:"maxMemoryUsedMB"`
+		} `json:"metrics"`
+	}
+	require.NoError(t, json.Unmarshal(encoded, &decoded), "a typed consumer must be able to decode this")
+
+	assert.Equal(t, "abc", decoded.RequestID)
+	assert.Equal(t, "success", decoded.Status)
+	assert.Equal(t, 12.5, decoded.Metrics.DurationMs)
+	assert.Equal(t, float64(13), decoded.Metrics.BilledDurationMs)
+	assert.Equal(t, 512, decoded.Metrics.MemorySizeMB)
+}
+
+// An invocation that ran out of time is not a successful one. Extensions alarm on
+// this field, so reporting a timeout as a success would hide exactly the failure
+// they are watching for.
+func TestReportRecordCarriesTheInvocationStatus(t *testing.T) {
+	for _, status := range []string{"success", "timeout"} {
+		t.Run(status, func(t *testing.T) {
+			assert.Equal(t, status, reportRecord("abc", status, 1, "128", 0)["status"])
+		})
+	}
+}
+
+// A memory size that isn't a number should not make the record undecodable.
+func TestReportRecordSurvivesAnUnparseableMemorySize(t *testing.T) {
+	record := reportRecord("abc", "success", 1, "not-a-number", 0)
+	metrics := record["metrics"].(map[string]interface{})
+	assert.Equal(t, 0, metrics["memorySizeMB"])
+}
+
+// A cold start reports how long initialization took; a warm invocation has no
+// such figure and must not report a zero one.
+func TestReportRecordIncludesInitDurationOnlyOnColdStart(t *testing.T) {
+	cold := reportRecord("abc", "success", 1, "128", 120.5)["metrics"].(map[string]interface{})
+	assert.Equal(t, 120.5, cold["initDurationMs"])
+
+	warm := reportRecord("abc", "success", 1, "128", 0)["metrics"].(map[string]interface{})
+	assert.NotContains(t, warm, "initDurationMs")
+}

internal/lambda/rie/http.go

@@ -6,9 +6,9 @@ package rie
 import (
 	"net/http"
 
-	log "github.com/sirupsen/logrus"
 	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/interop"
 	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda/rapidcore"
+	log "github.com/sirupsen/logrus"
 )
 
 func startHTTPServer(ipport string, sandbox *rapidcore.SandboxBuilder, bs interop.Bootstrap) {

internal/lambda/rie/telemetry_subscription.go

@@ -172,28 +172,30 @@ func (s *TelemetrySubscriptionService) Subscribe(agentName string, body io.Reade
 // soon as initialization finishes, which stops new subscriptions but must leave
 // delivery to existing ones running for the life of the environment.
 func (s *TelemetrySubscriptionService) Dispatch(event standalonetelemetry.SandboxEvent) {
-	s.lock.Lock()
-	subscriptions := make([]*subscription, 0, len(s.subscriptions))
-	for _, sub := range s.subscriptions {
-		subscriptions = append(subscriptions, sub)
-	}
-	s.lock.Unlock()
-
 	record := telemetryRecord{Time: event.Time, Type: event.Type}
 	if event.PlatformEvent != nil {
 		record.Record = event.PlatformEvent
 	} else {
 		record.Record = event.LogMessage
 	}
 
-	if len(subscriptions) == 0 {
-		s.lock.Lock()
+	// Buffering and delivery are chosen under a single held lock. Deciding in one
+	// lock section and acting in another lets a Subscribe interleave: it would
+	// insert itself and replay a copy of the buffer that does not yet hold this
+	// event, and the event would then be buffered for nobody.
+	s.lock.Lock()
+	if len(s.subscriptions) == 0 {
 		if len(s.earlyEvents) < maxEarlyEvents {
 			s.earlyEvents = append(s.earlyEvents, record)
 		}
 		s.lock.Unlock()
 		return
 	}
+	subscriptions := make([]*subscription, 0, len(s.subscriptions))
+	for _, sub := range s.subscriptions {
+		subscriptions = append(subscriptions, sub)
+	}
+	s.lock.Unlock()
 
 	for _, sub := range subscriptions {
 		if sub.wants(event.Type) {
@@ -262,7 +264,7 @@ func (sub *subscription) deliver(batch []telemetryRecord) {
 		log.WithError(err).Warn("Telemetry API: could not encode a batch")
 		return
 	}
-	response, err := http.Post(sub.destination, "application/json", bytes.NewReader(body))
+	response, err := deliveryClient.Post(sub.destination, "application/json", bytes.NewReader(body))
 	if err != nil {
 		log.WithError(err).Warnf("Telemetry API: could not deliver to %s", sub.destination)
 		return
@@ -340,6 +342,12 @@ func (s *TelemetrySubscriptionService) GetServiceClosedErrorType() string {
 
 const telemetryEndpointPath = "/2022-07-01/telemetry"
 
+// Batches that reach a size limit are delivered on the goroutine producing the
+// event, so a subscriber that accepts a connection and never answers would stall
+// the sandbox's event pipeline. The default client has no timeout; this one gives
+// up instead.
+var deliveryClient = &http.Client{Timeout: 5 * time.Second}
+
 // resolveDestination rewrites the hostname extensions are told to use. Inside a
 // real execution environment "sandbox" resolves to the host running the runtime;
 // in the emulator everything shares one network namespace.

internal/lambda/rie/telemetry_subscription_test.go

@@ -284,3 +284,84 @@ func TestSandboxHostIsRewritten(t *testing.T) {
 		})
 	}
 }
+
+// Exercises Subscribe and Dispatch concurrently, for the race detector and to
+// confirm the event still arrives.
+//
+// This does not reproduce the lost-event window that motivated deciding
+// buffer-versus-deliver under a single lock: the two lock sections it needed to
+// interleave between were adjacent, and 400 attempts never hit it. The
+// correctness of that arrangement rests on reading the code, not on this test.
+func TestConcurrentSubscribeAndDispatch(t *testing.T) {
+	for attempt := 0; attempt < 100; attempt++ {
+		service := NewTelemetrySubscriptionService()
+		extension := newReceiver(t)
+
+		var wg sync.WaitGroup
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			service.Dispatch(platformEvent("platform.initStart", map[string]interface{}{"phase": "init"}))
+		}()
+		go func() {
+			defer wg.Done()
+			subscribe(t, service, "ext", extension.server.URL, []string{"platform"})
+		}()
+		wg.Wait()
+
+		// Whichever order they interleaved in, the event is either delivered or
+		// still buffered for the next subscriber. It must not be stranded.
+		deadline := time.Now().Add(2 * time.Second)
+		var seen bool
+		for time.Now().Before(deadline) && !seen {
+			for _, event := range extension.received() {
+				if event["type"] == "platform.initStart" {
+					seen = true
+				}
+			}
+			if !seen {
+				time.Sleep(10 * time.Millisecond)
+			}
+		}
+		if !seen {
+			t.Fatalf("attempt %d: initStart reached no subscriber", attempt)
+		}
+	}
+}
+
+// A subscriber that accepts a connection and never answers must not stall the
+// sandbox: delivery happens on the event-producing goroutine when a batch fills.
+func TestDeliveryToAHungSubscriberDoesNotBlockForever(t *testing.T) {
+	blocked := make(chan struct{})
+
+	hung := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		<-blocked // never answers while the test runs
+	}))
+	// Close waits for handlers in flight, so the handler has to be released
+	// first. Defers run last-in-first-out, hence this order.
+	defer hung.Close()
+	defer close(blocked)
+
+	service := NewTelemetrySubscriptionService()
+	body, _ := json.Marshal(map[string]interface{}{
+		"types":       []string{"function"},
+		"buffering":   map[string]int{"timeoutMs": 60000, "maxItems": 1, "maxBytes": 262144},
+		"destination": map[string]string{"protocol": "HTTP", "URI": hung.URL},
+	})
+	_, status, _, err := service.Subscribe("ext", strings.NewReader(string(body)), nil, "")
+	require.NoError(t, err)
+	require.Equal(t, http.StatusOK, status)
+
+	// maxItems of 1 means this dispatch delivers synchronously.
+	done := make(chan struct{})
+	go func() {
+		service.Dispatch(logEvent("function", "hello"))
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(deliveryClient.Timeout + 5*time.Second):
+		t.Fatal("Dispatch never returned; a hung subscriber can stall the event pipeline")
+	}
+}