Changes for making extension work with lambda managed instance environment

CONTRIBUTING.md

@@ -58,6 +58,7 @@ First of all, thanks for contributing!. Before contributing please read the [COD
 1. Change the layer_name variable in zip.sh to avoid replacing the prod.
 1. Run the following command to publish the layer:
    `sh zip.sh`
+1. The zip file generated by previous step should have `extensions` folder in it, which should consist of the binary for the extension.
 1. Run the following command to verify that the layer version is published across regions:
    `sh verify_layer_versions.sh`
 

README.md

@@ -27,6 +27,9 @@ All the logs that are not sent to Sumo Logic during the Execution phase of the A
 
 If you would like to always send logs during the execution phase however, you can add extra execution time via a sleep function at the end of lambda code, which will give your extension time to run and send logs to Sumo Logic. We recommend setting this to two seconds.
 
+# Managed Instance Runtime Support
+This Lambda extension from version v1.4.0 also supports [managed instance](https://docs.aws.amazon.com/lambda/latest/dg/lambda-managed-instances.html) runtime.
+
 # Using Lambda extension in custom container images
 
 Follow the instruction in [docs](https://help.sumologic.com/03Send-Data/Collect-from-Other-Data-Sources/Collect_AWS_Lambda_Logs_using_an_Extension#For_AWS_Lambda_Functions_Created_Using_Container_Images:)

lambda-extensions/lambdaapi/extensionapiclient.go

@@ -46,15 +46,20 @@ const (
 )
 
 var (
-	lambdaEvents = []EventType{"INVOKE", "SHUTDOWN"}
+	lambdaEvents                = []EventType{"INVOKE", "SHUTDOWN"}
+	managedInstanceLambdaEvents = []EventType{"SHUTDOWN"}
 )
 
 // RegisterExtension is to register extension to Run Time API client. Call the following method on initialization as early as possible,
 // otherwise you may get a timeout error. Runtime initialization will start after all extensions are registered.
-func (client *Client) RegisterExtension(ctx context.Context) (*RegisterResponse, error) {
+func (client *Client) RegisterExtension(ctx context.Context, isManagedInstance bool) (*RegisterResponse, error) {
 	URL := client.baseURL + extensionURL + "register"
+	events := lambdaEvents
+	if isManagedInstance {
+		events = managedInstanceLambdaEvents
+	}
 	reqBody, err := json.Marshal(map[string]interface{}{
-		"events": lambdaEvents,
+		"events": events,
 	})
 	if err != nil {
 		return nil, err

lambda-extensions/lambdaapi/extensionapiclient_test.go

@@ -34,11 +34,11 @@ func TestRegisterExtension(t *testing.T) {
 	client := NewClient(srv.URL[7:], extensionName)
 
 	// Without Context
-	response, err := client.RegisterExtension(context.TODO())
+	response, err := client.RegisterExtension(context.TODO(), false)
 	commonAsserts(t, client, response, err)
 
 	// With Context
-	response, err = client.RegisterExtension(context.Background())
+	response, err = client.RegisterExtension(context.Background(), false)
 	commonAsserts(t, client, response, err)
 }
 
@@ -147,3 +147,146 @@ func TestExitError(t *testing.T) {
 	response, err = client.ExitError(context.Background(), "EXIT ERROR")
 	commonAsserts(t, client, response, err)
 }
+
+// TestRegisterExtension_ManagedInstanceMode tests extension registration in managed instance mode
+// In ManagedInstance mode, only SHUTDOWN events are registered (not INVOKE)
+func TestRegisterExtension_ManagedInstanceMode(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assertEqual(t, r.Method, http.MethodPost, "Method is not POST")
+		assertNotEmpty(t, r.Header.Get(extensionNameHeader), "Extension Name Header not present")
+
+		reqBytes, err := ioutil.ReadAll(r.Body)
+		assertNoError(t, err, "Received error while reading request")
+		defer func() {
+			if err := r.Body.Close(); err != nil {
+				log.Printf("failed to close body: %v", err)
+			}
+		}()
+		assertNotEmpty(t, reqBytes, "Received error in request")
+
+		// Verify the request body contains only SHUTDOWN event for managed instance mode
+		var reqBody map[string]interface{}
+		err = json.Unmarshal(reqBytes, &reqBody)
+		assertNoError(t, err, "Failed to unmarshal request body")
+
+		events, ok := reqBody["events"].([]interface{})
+		if !ok {
+			t.Error("Events field not found or not an array")
+		}
+		assertEqual(t, len(events), 1, "Expected 1 event for managed instance mode")
+		assertEqual(t, events[0], "SHUTDOWN", "Expected only SHUTDOWN event for managed instance mode")
+
+		w.Header().Add(extensionIdentiferHeader, "test-sumo-id")
+		w.WriteHeader(200)
+		respBytes, _ := json.Marshal(RegisterResponse{
+			FunctionName:    "test-function",
+			FunctionVersion: "$LATEST",
+			Handler:         "index.handler",
+		})
+		_, _ = w.Write(respBytes)
+	}))
+
+	defer srv.Close()
+	client := NewClient(srv.URL[7:], extensionName)
+
+	// Test with isManagedInstance = true
+	response, err := client.RegisterExtension(context.Background(), true)
+	commonAsserts(t, client, response, err)
+
+	// Verify the response is properly unmarshaled
+	if response.FunctionName != "test-function" {
+		t.Errorf("Expected function name 'test-function', got '%s'", response.FunctionName)
+	}
+}
+
+// TestRegisterExtension_ManagedInstanceModeWithoutContext tests managed instance mode without context
+func TestRegisterExtension_ManagedInstanceModeWithoutContext(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		reqBytes, err := ioutil.ReadAll(r.Body)
+		assertNoError(t, err, "Received error while reading request")
+		defer func() {
+			if err := r.Body.Close(); err != nil {
+				log.Printf("failed to close body: %v", err)
+			}
+		}()
+
+		var reqBody map[string]interface{}
+		err = json.Unmarshal(reqBytes, &reqBody)
+		assertNoError(t, err, "Failed to unmarshal request body")
+
+		events, ok := reqBody["events"].([]interface{})
+		if !ok {
+			t.Error("Events field not found or not an array")
+		}
+		assertEqual(t, len(events), 1, "Expected 1 event for managed instance mode")
+
+		w.Header().Add(extensionIdentiferHeader, "test-sumo-id")
+		w.WriteHeader(200)
+		respBytes, _ := json.Marshal(RegisterResponse{})
+		_, _ = w.Write(respBytes)
+	}))
+
+	defer srv.Close()
+	client := NewClient(srv.URL[7:], extensionName)
+
+	// Test with isManagedInstance = true and nil context
+	response, err := client.RegisterExtension(context.TODO(), true)
+	commonAsserts(t, client, response, err)
+}
+
+// TestRegisterExtension_ManagedInstanceModeEventValidation tests that managed instance mode registers correct events
+func TestRegisterExtension_ManagedInstanceModeEventValidation(t *testing.T) {
+	receivedEvents := make([]string, 0)
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		reqBytes, err := ioutil.ReadAll(r.Body)
+		assertNoError(t, err, "Received error while reading request")
+		defer func() {
+			if err := r.Body.Close(); err != nil {
+				log.Printf("failed to close body: %v", err)
+			}
+		}()
+
+		var reqBody map[string]interface{}
+		err = json.Unmarshal(reqBytes, &reqBody)
+		assertNoError(t, err, "Failed to unmarshal request body")
+
+		events, ok := reqBody["events"].([]interface{})
+		if !ok {
+			t.Error("Events field not found or not an array")
+		}
+
+		for _, e := range events {
+			receivedEvents = append(receivedEvents, e.(string))
+		}
+
+		w.Header().Add(extensionIdentiferHeader, "test-sumo-id")
+		w.WriteHeader(200)
+		respBytes, _ := json.Marshal(RegisterResponse{})
+		_, _ = w.Write(respBytes)
+	}))
+
+	defer srv.Close()
+	client := NewClient(srv.URL[7:], extensionName)
+
+	_, err := client.RegisterExtension(context.Background(), true)
+	assertNoError(t, err, "Failed to register extension in ManagedInstance mode")
+
+	// Validate that INVOKE event is NOT present in managed instance mode
+	for _, event := range receivedEvents {
+		if event == "INVOKE" {
+			t.Error("INVOKE event should not be registered in managed instance mode")
+		}
+	}
+
+	// Validate that SHUTDOWN event IS present
+	foundShutdown := false
+	for _, event := range receivedEvents {
+		if event == "SHUTDOWN" {
+			foundShutdown = true
+		}
+	}
+	if !foundShutdown {
+		t.Error("SHUTDOWN event should be registered in managed instance mode")
+	}
+}

lambda-extensions/lambdaapi/telemetryapiclient.go

@@ -14,15 +14,18 @@ const (
 	//telemetry_receiverPort = 4243
 )
 
-// SubscribeToLogsAPI is - Subscribe to Logs API to receive the Lambda Logs.
-func (client *Client) SubscribeToTelemetryAPI(ctx context.Context, logEvents []string, telemetryTimeoutMs int, telemetryMaxBytes int64, telemetryMaxItems int) ([]byte, error) {
+// SubscribeToTelemetryAPI is - Subscribe to Telemetry API to receive the Lambda Telemetry.
+func (client *Client) SubscribeToTelemetryAPI(ctx context.Context, logEvents []string, telemetryTimeoutMs int, telemetryMaxBytes int64, telemetryMaxItems int, isManagedInstance bool) ([]byte, error) {
 	URL := client.baseURL + telemetryURL
-
+	schemaVersion := "2022-07-01"
+	if isManagedInstance {
+		schemaVersion = "2025-01-29"
+	}
 	reqBody, err := json.Marshal(map[string]interface{}{
 		"destination":   map[string]interface{}{"protocol": "HTTP", "URI": fmt.Sprintf("http://sandbox:%v", receiverPort)},
 		"types":         logEvents,
 		"buffering":     map[string]interface{}{"timeoutMs": telemetryTimeoutMs, "maxBytes": telemetryMaxBytes, "maxItems": telemetryMaxItems},
-		"schemaVersion": "2022-07-01",
+		"schemaVersion": schemaVersion,
 	})
 	if err != nil {
 		return nil, err

lambda-extensions/lambdaapi/telemetryapiclient_test.go

@@ -2,6 +2,7 @@ package lambdaapi
 
 import (
 	"context"
+	"encoding/json"
 	ioutil "io"
 	"log"
 	"net/http"
@@ -31,10 +32,69 @@ func TestSubscribeToTelemetryAPI(t *testing.T) {
 	client := NewClient(srv.URL[7:], extensionName)
 
 	// Without Context
-	response, err := client.SubscribeToTelemetryAPI(context.TODO(), []string{"platform", "function", "extension"}, 1000, 262144, 10000)
+	response, err := client.SubscribeToTelemetryAPI(context.TODO(), []string{"platform", "function", "extension"}, 1000, 262144, 10000, false)
 	commonAsserts(t, client, response, err)
 
 	// With Context
-	response, err = client.SubscribeToTelemetryAPI(context.Background(), []string{"platform", "function", "extension"}, 1000, 262144, 10000)
+	response, err = client.SubscribeToTelemetryAPI(context.Background(), []string{"platform", "function", "extension"}, 1000, 262144, 10000, false)
+	commonAsserts(t, client, response, err)
+}
+
+// TestSubscribeToTelemetryAPI_ManagedInstanceMode tests telemetry API subscription in managed instance mode
+// In managed instance mode, schema version should be "2025-01-29" instead of "2022-07-01"
+func TestSubscribeToTelemetryAPI_ManagedInstanceMode(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assertEqual(t, r.Method, http.MethodPut, "Method is not PUT")
+		assertNotEmpty(t, r.Header.Get(extensionNameHeader), "Extension Name Header not present")
+
+		reqBytes, err := ioutil.ReadAll(r.Body)
+		assertNoError(t, err, "Received error")
+		defer func() {
+			if err := r.Body.Close(); err != nil {
+				log.Printf("failed to close body: %v", err)
+			}
+		}()
+		assertNotEmpty(t, reqBytes, "Received error in request")
+
+		// Verify the request body contains managed instance mode schema version
+		var reqBody map[string]interface{}
+		err = json.Unmarshal(reqBytes, &reqBody)
+		assertNoError(t, err, "Failed to unmarshal request body")
+
+		schemaVersion, ok := reqBody["schemaVersion"].(string)
+		if !ok {
+			t.Error("schemaVersion field not found or not a string")
+		}
+		assertEqual(t, schemaVersion, "2025-01-29", "Expected managed instance mode schema version '2025-01-29'")
+
+		// Verify other required fields are present
+		_, destinationExists := reqBody["destination"]
+		if !destinationExists {
+			t.Error("destination field not found")
+		}
+
+		_, typesExists := reqBody["types"]
+		if !typesExists {
+			t.Error("types field not found")
+		}
+
+		_, bufferingExists := reqBody["buffering"]
+		if !bufferingExists {
+			t.Error("buffering field not found")
+		}
+
+		w.Header().Add(extensionIdentiferHeader, "test-sumo-id")
+		w.WriteHeader(200)
+	}))
+
+	defer srv.Close()
+	client := NewClient(srv.URL[7:], extensionName)
+
+	// Test with isManagedInstance = true (context)
+	response, err := client.SubscribeToTelemetryAPI(context.Background(), []string{"platform", "function", "extension"}, 1000, 262144, 10000, true)
+	commonAsserts(t, client, response, err)
+
+	// Test with isManagedInstance = true (without context)
+	response, err = client.SubscribeToTelemetryAPI(context.TODO(), []string{"platform", "function", "extension"}, 1000, 262144, 10000, true)
 	commonAsserts(t, client, response, err)
 }