Skip to content

Commit 00744ce

Browse files
sumoanemasumoanema
committed
Changes in name for the runtime from elevator to managed instance
1 parent 089c2f4 commit 00744ce

9 files changed

Lines changed: 538 additions & 336 deletions

‎lambda-extensions/lambdaapi/extensionapiclient.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -46,17 +46,17 @@ const (
4646
)
4747

4848
var (
49-
lambdaEvents = []EventType{"INVOKE", "SHUTDOWN"}
50-
elevatorLambdaEvents = []EventType{"SHUTDOWN"}
49+
lambdaEvents = []EventType{"INVOKE", "SHUTDOWN"}
50+
managedInstanceLambdaEvents = []EventType{"SHUTDOWN"}
5151
)
5252

5353
// RegisterExtension is to register extension to Run Time API client. Call the following method on initialization as early as possible,
5454
// otherwise you may get a timeout error. Runtime initialization will start after all extensions are registered.
55-
func (client *Client) RegisterExtension(ctx context.Context, isElevator bool) (*RegisterResponse, error) {
55+
func (client *Client) RegisterExtension(ctx context.Context, isManagedInstance bool) (*RegisterResponse, error) {
5656
URL := client.baseURL + extensionURL + "register"
5757
events := lambdaEvents
58-
if isElevator {
59-
events = elevatorLambdaEvents
58+
if isManagedInstance {
59+
events = managedInstanceLambdaEvents
6060
}
6161
reqBody, err := json.Marshal(map[string]interface{}{
6262
"events": events,

‎lambda-extensions/lambdaapi/extensionapiclient_test.go‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,3 +147,146 @@ func TestExitError(t *testing.T) {
147147
response, err = client.ExitError(context.Background(), "EXIT ERROR")
148148
commonAsserts(t, client, response, err)
149149
}
150+
151+
// TestRegisterExtension_ManagedInstanceMode tests extension registration in managed instance mode
152+
// In ManagedInstance mode, only SHUTDOWN events are registered (not INVOKE)
153+
func TestRegisterExtension_ManagedInstanceMode(t *testing.T) {
154+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
155+
assertEqual(t, r.Method, http.MethodPost, "Method is not POST")
156+
assertNotEmpty(t, r.Header.Get(extensionNameHeader), "Extension Name Header not present")
157+
158+
reqBytes, err := ioutil.ReadAll(r.Body)
159+
assertNoError(t, err, "Received error while reading request")
160+
defer func() {
161+
if err := r.Body.Close(); err != nil {
162+
log.Printf("failed to close body: %v", err)
163+
}
164+
}()
165+
assertNotEmpty(t, reqBytes, "Received error in request")
166+
167+
// Verify the request body contains only SHUTDOWN event for managed instance mode
168+
var reqBody map[string]interface{}
169+
err = json.Unmarshal(reqBytes, &reqBody)
170+
assertNoError(t, err, "Failed to unmarshal request body")
171+
172+
events, ok := reqBody["events"].([]interface{})
173+
if !ok {
174+
t.Error("Events field not found or not an array")
175+
}
176+
assertEqual(t, len(events), 1, "Expected 1 event for managed instance mode")
177+
assertEqual(t, events[0], "SHUTDOWN", "Expected only SHUTDOWN event for managed instance mode")
178+
179+
w.Header().Add(extensionIdentiferHeader, "test-sumo-id")
180+
w.WriteHeader(200)
181+
respBytes, _ := json.Marshal(RegisterResponse{
182+
FunctionName: "test-function",
183+
FunctionVersion: "$LATEST",
184+
Handler: "index.handler",
185+
})
186+
_, _ = w.Write(respBytes)
187+
}))
188+
189+
defer srv.Close()
190+
client := NewClient(srv.URL[7:], extensionName)
191+
192+
// Test with isManagedInstance = true
193+
response, err := client.RegisterExtension(context.Background(), true)
194+
commonAsserts(t, client, response, err)
195+
196+
// Verify the response is properly unmarshaled
197+
if response.FunctionName != "test-function" {
198+
t.Errorf("Expected function name 'test-function', got '%s'", response.FunctionName)
199+
}
200+
}
201+
202+
// TestRegisterExtension_ManagedInstanceModeWithoutContext tests managed instance mode without context
203+
func TestRegisterExtension_ManagedInstanceModeWithoutContext(t *testing.T) {
204+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
205+
reqBytes, err := ioutil.ReadAll(r.Body)
206+
assertNoError(t, err, "Received error while reading request")
207+
defer func() {
208+
if err := r.Body.Close(); err != nil {
209+
log.Printf("failed to close body: %v", err)
210+
}
211+
}()
212+
213+
var reqBody map[string]interface{}
214+
err = json.Unmarshal(reqBytes, &reqBody)
215+
assertNoError(t, err, "Failed to unmarshal request body")
216+
217+
events, ok := reqBody["events"].([]interface{})
218+
if !ok {
219+
t.Error("Events field not found or not an array")
220+
}
221+
assertEqual(t, len(events), 1, "Expected 1 event for managed instance mode")
222+
223+
w.Header().Add(extensionIdentiferHeader, "test-sumo-id")
224+
w.WriteHeader(200)
225+
respBytes, _ := json.Marshal(RegisterResponse{})
226+
_, _ = w.Write(respBytes)
227+
}))
228+
229+
defer srv.Close()
230+
client := NewClient(srv.URL[7:], extensionName)
231+
232+
// Test with isManagedInstance = true and nil context
233+
response, err := client.RegisterExtension(context.TODO(), true)
234+
commonAsserts(t, client, response, err)
235+
}
236+
237+
// TestRegisterExtension_ManagedInstanceModeEventValidation tests that managed instance mode registers correct events
238+
func TestRegisterExtension_ManagedInstanceModeEventValidation(t *testing.T) {
239+
receivedEvents := make([]string, 0)
240+
241+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
242+
reqBytes, err := ioutil.ReadAll(r.Body)
243+
assertNoError(t, err, "Received error while reading request")
244+
defer func() {
245+
if err := r.Body.Close(); err != nil {
246+
log.Printf("failed to close body: %v", err)
247+
}
248+
}()
249+
250+
var reqBody map[string]interface{}
251+
err = json.Unmarshal(reqBytes, &reqBody)
252+
assertNoError(t, err, "Failed to unmarshal request body")
253+
254+
events, ok := reqBody["events"].([]interface{})
255+
if !ok {
256+
t.Error("Events field not found or not an array")
257+
}
258+
259+
for _, e := range events {
260+
receivedEvents = append(receivedEvents, e.(string))
261+
}
262+
263+
w.Header().Add(extensionIdentiferHeader, "test-sumo-id")
264+
w.WriteHeader(200)
265+
respBytes, _ := json.Marshal(RegisterResponse{})
266+
_, _ = w.Write(respBytes)
267+
}))
268+
269+
defer srv.Close()
270+
client := NewClient(srv.URL[7:], extensionName)
271+
272+
_, err := client.RegisterExtension(context.Background(), true)
273+
assertNoError(t, err, "Failed to register extension in ManagedInstance mode")
274+
275+
// Validate that INVOKE event is NOT present in managed instance mode
276+
for _, event := range receivedEvents {
277+
if event == "INVOKE" {
278+
t.Error("INVOKE event should not be registered in managed instance mode")
279+
}
280+
}
281+
282+
// Validate that SHUTDOWN event IS present
283+
foundShutdown := false
284+
for _, event := range receivedEvents {
285+
if event == "SHUTDOWN" {
286+
foundShutdown = true
287+
}
288+
}
289+
if !foundShutdown {
290+
t.Error("SHUTDOWN event should be registered in managed instance mode")
291+
}
292+
}

‎lambda-extensions/lambdaapi/telemetryapiclient.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,11 @@ const (
1414
//telemetry_receiverPort = 4243
1515
)
1616

17-
// SubscribeToLogsAPI is - Subscribe to Logs API to receive the Lambda Logs.
18-
func (client *Client) SubscribeToTelemetryAPI(ctx context.Context, logEvents []string, telemetryTimeoutMs int, telemetryMaxBytes int64, telemetryMaxItems int, isElevator bool) ([]byte, error) {
17+
// SubscribeToTelemetryAPI is - Subscribe to Telemetry API to receive the Lambda Telemetry.
18+
func (client *Client) SubscribeToTelemetryAPI(ctx context.Context, logEvents []string, telemetryTimeoutMs int, telemetryMaxBytes int64, telemetryMaxItems int, isManagedInstance bool) ([]byte, error) {
1919
URL := client.baseURL + telemetryURL
2020
schemaVersion := "2022-07-01"
21-
if isElevator {
21+
if isManagedInstance {
2222
schemaVersion = "2025-01-29"
2323
}
2424
reqBody, err := json.Marshal(map[string]interface{}{

‎lambda-extensions/lambdaapi/telemetryapiclient_test.go‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package lambdaapi
22

33
import (
44
"context"
5+
"encoding/json"
56
ioutil "io"
67
"log"
78
"net/http"
@@ -38,3 +39,62 @@ func TestSubscribeToTelemetryAPI(t *testing.T) {
3839
response, err = client.SubscribeToTelemetryAPI(context.Background(), []string{"platform", "function", "extension"}, 1000, 262144, 10000, false)
3940
commonAsserts(t, client, response, err)
4041
}
42+
43+
// TestSubscribeToTelemetryAPI_ManagedInstanceMode tests telemetry API subscription in managed instance mode
44+
// In managed instance mode, schema version should be "2025-01-29" instead of "2022-07-01"
45+
func TestSubscribeToTelemetryAPI_ManagedInstanceMode(t *testing.T) {
46+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
47+
assertEqual(t, r.Method, http.MethodPut, "Method is not PUT")
48+
assertNotEmpty(t, r.Header.Get(extensionNameHeader), "Extension Name Header not present")
49+
50+
reqBytes, err := ioutil.ReadAll(r.Body)
51+
assertNoError(t, err, "Received error")
52+
defer func() {
53+
if err := r.Body.Close(); err != nil {
54+
log.Printf("failed to close body: %v", err)
55+
}
56+
}()
57+
assertNotEmpty(t, reqBytes, "Received error in request")
58+
59+
// Verify the request body contains managed instance mode schema version
60+
var reqBody map[string]interface{}
61+
err = json.Unmarshal(reqBytes, &reqBody)
62+
assertNoError(t, err, "Failed to unmarshal request body")
63+
64+
schemaVersion, ok := reqBody["schemaVersion"].(string)
65+
if !ok {
66+
t.Error("schemaVersion field not found or not a string")
67+
}
68+
assertEqual(t, schemaVersion, "2025-01-29", "Expected managed instance mode schema version '2025-01-29'")
69+
70+
// Verify other required fields are present
71+
_, destinationExists := reqBody["destination"]
72+
if !destinationExists {
73+
t.Error("destination field not found")
74+
}
75+
76+
_, typesExists := reqBody["types"]
77+
if !typesExists {
78+
t.Error("types field not found")
79+
}
80+
81+
_, bufferingExists := reqBody["buffering"]
82+
if !bufferingExists {
83+
t.Error("buffering field not found")
84+
}
85+
86+
w.Header().Add(extensionIdentiferHeader, "test-sumo-id")
87+
w.WriteHeader(200)
88+
}))
89+
90+
defer srv.Close()
91+
client := NewClient(srv.URL[7:], extensionName)
92+
93+
// Test with isManagedInstance = true (context)
94+
response, err := client.SubscribeToTelemetryAPI(context.Background(), []string{"platform", "function", "extension"}, 1000, 262144, 10000, true)
95+
commonAsserts(t, client, response, err)
96+
97+
// Test with isManagedInstance = true (without context)
98+
response, err = client.SubscribeToTelemetryAPI(context.TODO(), []string{"platform", "function", "extension"}, 1000, 262144, 10000, true)
99+
commonAsserts(t, client, response, err)
100+
}

‎lambda-extensions/sumologic-extension.go‎

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,12 @@ var (
2626

2727
var producer workers.TaskProducer
2828
var consumer workers.TaskConsumer
29-
var elevatorProducer workers.ElevatorTaskProducer
30-
var elevatorConsumer workers.ElevatorTaskConsumer
29+
var managedInstanceProducer workers.ManagedInstanceTaskProducer
30+
var managedInstanceConsumer workers.ManagedInstanceTaskConsumer
3131
var config *cfg.LambdaExtensionConfig
3232
var dataQueue chan []byte
3333
var flushSignal chan string
34-
var isElevator bool
34+
var isManagedInstance bool
3535

3636
func init() {
3737
Formatter := new(logrus.TextFormatter)
@@ -51,30 +51,30 @@ func init() {
5151
logger.Logger.SetLevel(config.LogLevel)
5252
dataQueue = make(chan []byte, config.MaxDataQueueLength)
5353

54-
// Check initialization type to determine if elevator mode should be used
54+
// Check initialization type to determine if managed instance mode should be used
5555
initializationType := os.Getenv("AWS_LAMBDA_INITIALIZATION_TYPE")
5656
if initializationType == "lambda-managed-instances" {
57-
isElevator = true
58-
logger.Debug("Initializing in Elevator mode")
57+
isManagedInstance = true
58+
logger.Debug("Initializing in Managed Instance mode")
5959

60-
// Initialize flushSignal channel for elevator mode communication
60+
// Initialize flushSignal channel for managed instance mode communication
6161
flushSignal = make(chan string, 10) // Buffered channel to prevent blocking
6262

63-
// Initialize Elevator Producer and start it in a goroutine
64-
elevatorProducer = workers.NewElevatorTaskProducer(dataQueue, flushSignal, logger)
63+
// Initialize Managed Instance Producer and start it in a goroutine
64+
managedInstanceProducer = workers.NewManagedInstanceTaskProducer(dataQueue, flushSignal, logger)
6565
go func() {
66-
if err := elevatorProducer.Start(); err != nil {
67-
logger.Errorf("elevatorProducer Start failed: %v", err)
66+
if err := managedInstanceProducer.Start(); err != nil {
67+
logger.Errorf("managedInstanceProducer Start failed: %v", err)
6868
}
6969
}()
7070

71-
// Initialize Elevator Consumer and start it
72-
elevatorConsumer = workers.NewElevatorTaskConsumer(dataQueue, flushSignal, config, logger)
71+
// Initialize Managed Instance Consumer and start it
72+
managedInstanceConsumer = workers.NewManagedInstanceTaskConsumer(dataQueue, flushSignal, config, logger)
7373
// Start the consumer's independent processing loop
7474
ctx := context.Background()
75-
elevatorConsumer.Start(ctx)
75+
managedInstanceConsumer.Start(ctx)
7676

77-
logger.Debug("Elevator mode initialization complete")
77+
logger.Debug("Managed Instance mode initialization complete")
7878
} else {
7979
logger.Debug("Initializing in standard mode")
8080
// Start HTTP Server before subscription in a goRoutine
@@ -90,29 +90,29 @@ func init() {
9090
logger.Debug("Standard mode initialization complete")
9191
}
9292

93-
logger.Debug("Is Elevator value: ", isElevator)
93+
logger.Debug("Is Managed Instance value: ", isManagedInstance)
9494
}
9595

9696
func runTimeAPIInit() (int64, error) {
9797
// Register early so Runtime could start in parallel
9898
logger.Debug("Registering Extension to Run Time API Client..........")
99-
registerResponse, err := extensionClient.RegisterExtension(context.TODO(), isElevator)
99+
registerResponse, err := extensionClient.RegisterExtension(context.TODO(), isManagedInstance)
100100
if err != nil {
101101
return 0, err
102102
}
103103
logger.Debug("Succcessfully Registered with Run Time API Client: ", utils.PrettyPrint(registerResponse))
104104

105105
// Subscribe to Telemetry API
106106
logger.Debug("Subscribing Extension to Telemetry API........")
107-
subscribeResponse, err := extensionClient.SubscribeToTelemetryAPI(context.TODO(), config.LogTypes, config.TelemetryTimeoutMs, config.TelemetryMaxBytes, config.TelemetryMaxItems, isElevator)
107+
subscribeResponse, err := extensionClient.SubscribeToTelemetryAPI(context.TODO(), config.LogTypes, config.TelemetryTimeoutMs, config.TelemetryMaxBytes, config.TelemetryMaxItems, isManagedInstance)
108108
if err != nil {
109109
return 0, err
110110
}
111111

112112
logger.Debug("Successfully subscribed to Telemetry API: ", utils.PrettyPrint(string(subscribeResponse)))
113113

114114
// Call next to say registration is successful and get the deadtimems
115-
if !isElevator {
115+
if !isManagedInstance {
116116
nextResponse, err := nextEvent(context.TODO())
117117
if err != nil {
118118
return 0, err
@@ -146,7 +146,7 @@ func processEvents(ctx context.Context) {
146146
consumer.FlushDataQueue(ctx)
147147
return
148148
default:
149-
if !isElevator {
149+
if !isManagedInstance {
150150
logger.Debugf("switching to other go routine")
151151
runtime.Gosched()
152152
logger.Infof("Calling DrainQueue from processEvents")

0 commit comments

Comments
 (0)