Serverless DeepDive #11：Extensionあり・なしの課金時間を比較してみる

package main

import (
"bytes"
"crypto/rand"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"net"
"net/http"
"os"
"strconv"
"sync"
"sync/atomic"
"time"
)

var outputMu sync.Mutex
var sequence atomic.Uint64
var initID string

func emit(record map[string]any) {
record["extension_init_id"] = initID
record["function_id"] = os.Getenv("LAB_FUNCTION_ID")
outputMu.Lock()
defer outputMu.Unlock()
_ = json.NewEncoder(os.Stdout).Encode(record)
}

func setting(name string, fallback, minimum, maximum int) int {
value := os.Getenv(name)
if value == "" {
return fallback
}
parsed, err := strconv.Atoi(value)
if err != nil || parsed < minimum || parsed > maximum {
panic(fmt.Sprintf("invalid non-secret setting %s", name))
}
return parsed
}

func api(client *http.Client, method, endpoint, identifier string, body any) (*http.Response, error) {
var reader io.Reader
if body != nil {
encoded, err := json.Marshal(body)
if err != nil {
return nil, err
}
reader = bytes.NewReader(encoded)
}
request, err := http.NewRequest(method, endpoint, reader)
if err != nil {
return nil, err
}
request.Header.Set("Content-Type", "application/json")
if identifier == "" {
request.Header.Set("Lambda-Extension-Name", "lab-extension")
} else {
request.Header.Set("Lambda-Extension-Identifier", identifier)
}
return client.Do(request)
}

func subscribe(client *http.Client, base, identifier string) error {
listener, err := net.Listen("tcp", ":4243")
if err != nil {
return err
}
mux := http.NewServeMux()
mux.HandleFunc("/", func(writer http.ResponseWriter, request *http.Request) {
if request.Method != http.MethodPost {
writer.WriteHeader(http.StatusMethodNotAllowed)
return
}
var events []struct {
Type string `json:"type"`
Time string `json:"time"`
Record json.RawMessage `json:"record"`
}
// Subscription uses at most 256 KiB batches. Keep a bounded margin for JSON framing.
if err := json.NewDecoder(http.MaxBytesReader(writer, request.Body, 1024*1024)).Decode(&events); err != nil {
writer.WriteHeader(http.StatusBadRequest)
return
}
for _, event := range events {
emit(map[string]any{
"phase": "telemetry-event", "type": event.Type, "time": event.Time,
"record": event.Record, "received_ms": time.Now().UnixMilli(),
"seq": sequence.Add(1),
})
}
writer.WriteHeader(http.StatusOK)
})
server := &http.Server{Handler: mux, ReadHeaderTimeout: time.Second}
go func() { _ = server.Serve(listener) }()
response, err := api(client, http.MethodPut, base+"/2022-07-01/telemetry", identifier, map[string]any{
"schemaVersion": "2025-01-29", "types": []string{"platform", "function"},
"buffering": map[string]int{"maxItems": 1000, "maxBytes": 262144, "timeoutMs": setting("EXT_BUFFER_MS", 1000, 25, 30000)},
"destination": map[string]string{"protocol": "HTTP", "URI": "http://sandbox.localdomain:4243"},
})
if err != nil {
return err
}
defer response.Body.Close()
if response.StatusCode != http.StatusOK {
return fmt.Errorf("telemetry subscribe HTTP %d", response.StatusCode)
}
emit(map[string]any{"phase": "extension-subscribed", "schema": "2025-01-29"})
return nil
}

func main() {
random := make([]byte, 16)
if _, err := rand.Read(random); err != nil {
panic(err)
}
initID = hex.EncodeToString(random)
started := time.Now().UnixMilli()
base := "http://" + os.Getenv("AWS_LAMBDA_RUNTIME_API")
client := &http.Client{} // Extension Next intentionally has no request timeout.
response, err := api(client, http.MethodPost, base+"/2020-01-01/extension/register", "", map[string]any{"events": []string{"INVOKE", "SHUTDOWN"}})
if err != nil {
panic(err)
}
identifier := response.Header.Get("Lambda-Extension-Identifier")
response.Body.Close()
if response.StatusCode != http.StatusOK || identifier == "" {
panic("extension registration failed")
}
if os.Getenv("EXT_TELEMETRY") == "1" {
if err := subscribe(client, base, identifier); err != nil {
panic(err)
}
}
initMS := setting("EXT_INIT_MS", 100, 0, 500)
time.Sleep(time.Duration(initMS) * time.Millisecond)
emit(map[string]any{"phase": "extension-init", "started_ms": started, "finished_ms": time.Now().UnixMilli(), "requested_ms": initMS})
for {
response, err := api(client, http.MethodGet, base+"/2020-01-01/extension/event/next", identifier, nil)
if err != nil {
panic(err)
}
var event struct {
EventType string `json:"eventType"`
RequestID string `json:"requestId"`
DeadlineMS int64 `json:"deadlineMs"`
ShutdownReason string `json:"shutdownReason"`
}
err = json.NewDecoder(io.LimitReader(response.Body, 16384)).Decode(&event)
response.Body.Close()
if err != nil || response.StatusCode != http.StatusOK {
panic("extension Next failed")
}
if event.EventType == "SHUTDOWN" {
emit(map[string]any{"phase": "extension-shutdown", "reason": event.ShutdownReason, "at_ms": time.Now().UnixMilli()})
return
}
invokeStarted := time.Now().UnixMilli()
postMS := setting("EXT_POST_MS", 150, 0, 500)
markerMS := int64(0)
if os.Getenv("EXT_WAIT_FOR_MARKER") != "0" {
markerPath := "/tmp/lab-extension-" + event.RequestID + ".json"
for time.Now().UnixMilli() < event.DeadlineMS {
data, readErr := os.ReadFile(markerPath)
if readErr == nil {
var marker struct {
PostMS int `json:"post_ms"`
FinishedMS int64 `json:"finished_ms"`
}
if json.Unmarshal(data, &marker) == nil && marker.PostMS >= 0 && marker.PostMS <= 500 {
postMS = marker.PostMS
markerMS = marker.FinishedMS
_ = os.Remove(markerPath)
break
}
}
time.Sleep(5 * time.Millisecond)
}
}
emit(map[string]any{"phase": "extension-post-start", "request_id": event.RequestID, "invoke_started_ms": invokeStarted, "marker_ms": markerMS, "post_started_ms": time.Now().UnixMilli(), "requested_ms": postMS})
time.Sleep(time.Duration(postMS) * time.Millisecond)
emit(map[string]any{"phase": "extension-done", "request_id": event.RequestID, "invoke_started_ms": invokeStarted, "marker_ms": markerMS, "finished_ms": time.Now().UnixMilli(), "requested_ms": postMS})
}
}