refactor: Modularise packages and use custom rapidcore components

Author: gregfurman

.github/workflows/build.yml

@@ -29,7 +29,7 @@ jobs:
         name: aws-lambda-rie
         path: bin/*
     - name: Release binaries
-      uses: softprops/action-gh-release@v2
+      uses: softprops/action-gh-release@v2.2.2
       if: startsWith(github.ref, 'refs/tags/')
       with:
         files: bin/*

.gitignore

@@ -7,3 +7,4 @@ tags
 .idea
 .DS_Store
 .venv
+.vscode

cmd/localstack/main.go

@@ -4,28 +4,41 @@ package main
 
 import (
 	"context"
+	"fmt"
+	"net"
 	"os"
+	"os/signal"
 	"runtime/debug"
 	"strconv"
 	"strings"
+	"syscall"
 
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/localstack/lambda-runtime-init/internal/aws/lambda"
 	"github.com/localstack/lambda-runtime-init/internal/aws/xray"
 	"github.com/localstack/lambda-runtime-init/internal/bootstrap"
+	"github.com/localstack/lambda-runtime-init/internal/events"
 	"github.com/localstack/lambda-runtime-init/internal/hotreloading"
+	"github.com/localstack/lambda-runtime-init/internal/localstack"
 	"github.com/localstack/lambda-runtime-init/internal/logging"
+	"github.com/localstack/lambda-runtime-init/internal/sandbox"
+
 	"github.com/localstack/lambda-runtime-init/internal/server"
+
+	"github.com/localstack/lambda-runtime-init/internal/supervisor"
 	"github.com/localstack/lambda-runtime-init/internal/tracing"
 	"github.com/localstack/lambda-runtime-init/internal/utils"
 	log "github.com/sirupsen/logrus"
 	"go.amzn.com/lambda/core/directinvoke"
+	"go.amzn.com/lambda/interop"
 	"go.amzn.com/lambda/rapidcore"
 )
 
-func InitLsOpts() *server.LsOpts {
-	return &server.LsOpts{
+func InitLsOpts() *localstack.Config {
+	return &localstack.Config{
 		// required
-		RuntimeEndpoint: utils.GetEnvOrDie("LOCALSTACK_RUNTIME_ENDPOINT"),
-		RuntimeId:       utils.GetEnvOrDie("LOCALSTACK_RUNTIME_ID"),
+		RuntimeEndpoint: utils.MustGetEnv("LOCALSTACK_RUNTIME_ENDPOINT"),
+		RuntimeId:       utils.MustGetEnv("LOCALSTACK_RUNTIME_ID"),
 		AccountId:       utils.GetEnvWithDefault("LOCALSTACK_FUNCTION_ACCOUNT_ID", "000000000000"),
 		// optional with default
 		InteropPort:     utils.GetEnvWithDefault("LOCALSTACK_INTEROP_PORT", "9563"),
@@ -45,6 +58,19 @@ func InitLsOpts() *server.LsOpts {
 	}
 }
 
+func InitFunctionConfig() lambda.FunctionConfig {
+	return lambda.FunctionConfig{
+		FunctionName:         utils.GetEnvWithDefault("AWS_LAMBDA_FUNCTION_NAME", "test_function"),
+		FunctionVersion:      utils.GetEnvWithDefault("AWS_LAMBDA_FUNCTION_VERSION", "$LATEST"),
+		FunctionTimeoutSec:   utils.GetEnvWithDefault("AWS_LAMBDA_FUNCTION_TIMEOUT", "30"),
+		InitializationType:   utils.GetEnvWithDefault("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand"),
+		LogGroupName:         utils.GetEnvWithDefault("AWS_LAMBDA_LOG_GROUP_NAME", "/aws/lambda/Functions"),
+		LogStreamName:        utils.GetEnvWithDefault("AWS_LAMBDA_LOG_STREAM_NAME", "$LATEST"),
+		FunctionMemorySizeMb: utils.GetEnvWithDefault("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "128"),
+		FunctionHandler:      utils.GetEnvWithDefault("AWS_LAMBDA_FUNCTION_HANDLER", os.Getenv("_HANDLER")),
+	}
+}
+
 // UnsetLsEnvs unsets environment variables specific to LocalStack to achieve better runtime parity with AWS
 func UnsetLsEnvs() {
 	unsetList := [...]string{
@@ -75,42 +101,54 @@ func UnsetLsEnvs() {
 	}
 }
 
+func setupUserPermissions(user string) {
+	logger := utils.UserLogger()
+	// Switch to non-root user and drop root privileges
+	if utils.IsRootUser() && user != "" && user != "root" {
+		uid := 993
+		gid := 990
+		utils.AddUser(user, uid, gid)
+		if err := os.Chown("/tmp", uid, gid); err != nil {
+			log.WithError(err).Warn("Could not change owner of directory /tmp.")
+		}
+		logger.Debug("Process running as root user.")
+		err := utils.DropPrivileges(user)
+		if err != nil {
+			log.WithError(err).Warn("Could not drop root privileges.")
+		} else {
+			logger.Debug("Process running as non-root user.")
+		}
+	}
+}
+
 func main() {
 	// we're setting this to the same value as in the official RIE
 	debug.SetGCPercent(33)
 
 	// configuration parsing
 	lsOpts := InitLsOpts()
+	functionConf := InitFunctionConfig()
+	awsEnvConf, _ := config.NewEnvConfig()
+	awsEnvConf.Credentials.AccountID = lsOpts.AccountId
+
 	UnsetLsEnvs()
 
 	// set up logging following the Logrus logging levels: https://github.com/sirupsen/logrus#level-logging
 	log.SetReportCaller(true)
 	// https://docs.aws.amazon.com/xray/latest/devguide/xray-daemon-configuration.html
-	xRayLogLevel := "info"
-	switch lsOpts.InitLogLevel {
-	case "trace":
+
+	logLevel, err := log.ParseLevel(lsOpts.InitLogLevel)
+	if err != nil {
+		log.Fatal("Invalid value for LOCALSTACK_INIT_LOG_LEVEL")
+	}
+	log.SetLevel(logLevel)
+
+	xRayLogLevel := lsOpts.InitLogLevel
+	switch logLevel {
+	case log.TraceLevel:
 		log.SetFormatter(&log.JSONFormatter{})
-		log.SetLevel(log.TraceLevel)
-		xRayLogLevel = "debug"
-	case "debug":
-		log.SetLevel(log.DebugLevel)
-		xRayLogLevel = "debug"
-	case "info":
-		log.SetLevel(log.InfoLevel)
-	case "warn":
-		log.SetLevel(log.WarnLevel)
-		xRayLogLevel = "warn"
-	case "error":
-		log.SetLevel(log.ErrorLevel)
-		xRayLogLevel = "error"
-	case "fatal":
-		log.SetLevel(log.FatalLevel)
+	case log.ErrorLevel, log.FatalLevel, log.PanicLevel:
 		xRayLogLevel = "error"
-	case "panic":
-		log.SetLevel(log.PanicLevel)
-		xRayLogLevel = "error"
-	default:
-		log.Fatal("Invalid value for LOCALSTACK_INIT_LOG_LEVEL")
 	}
 
 	// patch MaxPayloadSize
@@ -119,6 +157,10 @@ func main() {
 		log.Panicln("Please specify a number for LOCALSTACK_MAX_PAYLOAD_SIZE")
 	}
 	directinvoke.MaxDirectResponseSize = int64(payloadSize)
+	if directinvoke.MaxDirectResponseSize > interop.MaxPayloadSize {
+		log.Infof("Large response size detected (%d bytes), forcing streaming mode", directinvoke.MaxDirectResponseSize)
+		directinvoke.InvokeResponseMode = interop.InvokeResponseModeStreaming
+	}
 
 	// download code archive if env variable is set
 	if err := utils.DownloadCodeArchives(lsOpts.CodeArchives); err != nil {
@@ -133,97 +175,119 @@ func main() {
 	bootstrap, handler := bootstrap.GetBootstrap(os.Args)
 
 	// Switch to non-root user and drop root privileges
-	if utils.IsRootUser() && lsOpts.User != "" && lsOpts.User != "root" {
-		uid := 993
-		gid := 990
-		utils.AddUser(lsOpts.User, uid, gid)
-		if err := os.Chown("/tmp", uid, gid); err != nil {
-			log.Warnln("Could not change owner of directory /tmp:", err)
-		}
-		utils.UserLogger().Debugln("Process running as root user.")
-		err := utils.DropPrivileges(lsOpts.User)
-		if err != nil {
-			log.Warnln("Could not drop root privileges.", err)
-		} else {
-			utils.UserLogger().Debugln("Process running as non-root user.")
-		}
-	}
+	setupUserPermissions(lsOpts.User)
 
-	// file watcher for hot-reloading
-	fileWatcherContext, cancelFileWatcher := context.WithCancel(context.Background())
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer stop()
+
+	// LocalStack client used for sending callbacks
+	lsClient := localstack.NewLocalStackClient(lsOpts.RuntimeEndpoint, lsOpts.RuntimeId)
+
+	// Services required for Sandbox environment
+	interopServer := server.NewInteropServer(lsClient)
 
 	logCollector := logging.NewLogCollector()
 	localStackLogsEgressApi := logging.NewLocalStackLogsEgressAPI(logCollector)
 	tracer := tracing.NewLocalStackTracer()
+	lsEventsAPI := events.NewLocalStackEventsAPI(lsClient)
+	localStackSupv := supervisor.NewLocalStackSupervisor(ctx, lsEventsAPI)
 
 	// build sandbox
-	sandbox := rapidcore.
+	builder := rapidcore.
 		NewSandboxBuilder().
-		//SetTracer(tracer).
-		AddShutdownFunc(func() {
-			log.Debugln("Stopping file watcher")
-			cancelFileWatcher()
-		}).
-		SetExtensionsFlag(true).
-		SetInitCachingFlag(true).
+		SetTracer(tracer).
+		SetEventsAPI(lsEventsAPI).
+		SetHandler(handler).
+		SetInteropServer(interopServer).
 		SetLogsEgressAPI(localStackLogsEgressApi).
-		SetTracer(tracer)
+		SetRuntimeAPIAddress("127.0.0.1:9000").
+		SetSupervisor(localStackSupv).
+		SetRuntimeFsRootPath("/")
+
+	builder.AddShutdownFunc(func() {
+		interopServer.Close()
+	})
+
+	// Start daemons
+
+	// file watcher for hot-reloading
+	fileWatcherContext, cancelFileWatcher := context.WithCancel(ctx)
+	builder.AddShutdownFunc(func() {
+		cancelFileWatcher()
+	})
+
+	go hotreloading.RunHotReloadingListener(interopServer, lsOpts.HotReloadingPaths, fileWatcherContext, lsOpts.FileWatcherStrategy)
 
-	// xray daemon
-	endpoint := "http://" + lsOpts.LocalstackIP + ":" + lsOpts.EdgePort
+	// Start xray daemon
+	endpoint := "http://" + net.JoinHostPort(lsOpts.LocalstackIP, lsOpts.EdgePort)
 	xrayConfig := xray.NewConfig(endpoint, xRayLogLevel)
 	d := xray.NewDaemon(xrayConfig, lsOpts.EnableXRayTelemetry == "1")
-	sandbox.AddShutdownFunc(func() {
+	builder.AddShutdownFunc(func() {
 		log.Debugln("Shutting down xray daemon")
 		d.Stop()
 		log.Debugln("Flushing segments in xray daemon")
 		d.Close()
 	})
-	d.Run() // async
+	d.Run() // served async
 
-	defaultInterop := sandbox.DefaultInteropServer()
-	interopServer := server.NewCustomInteropServer(lsOpts, defaultInterop, logCollector)
-	sandbox.SetInteropServer(interopServer)
-	if len(handler) > 0 {
-		sandbox.SetHandler(handler)
-	}
-	exitChan := make(chan struct{})
-	sandbox.AddShutdownFunc(func() {
-		exitChan <- struct{}{}
-	})
+	// Populate our interop server
+	sandboxCtx, internalStateFn := builder.Create()
 
-	// initialize all flows and start runtime API
-	sandboxContext, internalStateFn := sandbox.Create()
-	// Populate our custom interop server
-	interopServer.SetSandboxContext(sandboxContext)
+	interopServer.SetSandboxContext(sandboxCtx)
 	interopServer.SetInternalStateGetter(internalStateFn)
 
-	// get timeout
-	invokeTimeoutEnv := utils.GetEnvOrDie("AWS_LAMBDA_FUNCTION_TIMEOUT") // TODO: collect all AWS_* env parsing
-	invokeTimeoutSeconds, err := strconv.Atoi(invokeTimeoutEnv)
-	if err != nil {
-		log.Fatalln(err)
-	}
-	go hotreloading.RunHotReloadingListener(interopServer, lsOpts.HotReloadingPaths, fileWatcherContext, lsOpts.FileWatcherStrategy)
+	// TODO(gregfurman): The default interops server has a shutdown function added in the builder.
+	// This calls sandbox.Reset(), where if sandbox is not set, this will panic.
+	// So instead, just set the sandbox to a noop version.
+	builder.DefaultInteropServer().SetSandboxContext(sandbox.NewNoopSandboxContext())
+
+	// Create the LocalStack service
+	localStackService := server.NewLocalStackService(
+		interopServer, logCollector, lsClient, localStackSupv, xrayConfig.Endpoint, lsOpts, functionConf, awsEnvConf,
+	)
+	defer localStackService.Close()
 
 	// start runtime init. It is important to start `InitHandler` synchronously because we need to ensure the
 	// notification channels and status fields are properly initialized before `AwaitInitialized`
 	log.Debugln("Starting runtime init.")
-	server.InitHandler(sandbox.LambdaInvokeAPI(), utils.GetEnvOrDie("AWS_LAMBDA_FUNCTION_VERSION"), int64(invokeTimeoutSeconds), bootstrap, lsOpts.AccountId) // TODO: replace this with a custom init
+	if err := localStackService.Initialize(bootstrap); err != nil {
+		log.WithError(err).Warnf("Failed to initialize runtime. Initialization will be retried in the next invoke.")
+	}
+
+	invokeServer := server.NewServer(lsOpts.InteropPort, localStackService)
+	defer invokeServer.Close()
+
+	serverErr := make(chan error, 1)
+	go func() {
+		listener, err := net.Listen("tcp", fmt.Sprintf(":%s", lsOpts.InteropPort))
+		if err != nil {
+			log.Fatalf("failed to start LocalStack Lambda Runtime Interface server: %s", err)
+		}
+		go func() { serverErr <- invokeServer.Serve(listener); close(serverErr) }()
+		log.Debugf("LocalStack API gateway listening on %s", listener.Addr().String())
+	}()
 
 	log.Debugln("Awaiting initialization of runtime init.")
 	if err := interopServer.AwaitInitialized(); err != nil {
 		// Error cases: ErrInitDoneFailed or ErrInitResetReceived
 		log.Errorln("Runtime init failed to initialize: " + err.Error() + ". Exiting.")
 		// NOTE: Sending the error status to LocalStack is handled beforehand in the custom_interop.go through the
 		// callback SendInitErrorResponse because it contains the correct error response payload.
-		return
+		// return
+	} else {
+		log.Debugln("Completed initialization of runtime init. Sending status ready to LocalStack.")
+		if err := localStackService.SendStatus(localstack.Ready, []byte{}); err != nil {
+			log.Fatalln("Failed to send status ready to LocalStack " + err.Error() + ". Exiting.")
+		}
 	}
 
-	log.Debugln("Completed initialization of runtime init. Sending status ready to LocalStack.")
-	if err := interopServer.SendStatus(server.Ready, []byte{}); err != nil {
-		log.Fatalln("Failed to send status ready to LocalStack " + err.Error() + ". Exiting.")
+	// Block until context is cancelled OR the server errors out
+	select {
+	case <-ctx.Done():
+		log.Info("Shutdown signal received.")
+	case <-serverErr:
+		if err != nil {
+			log.Errorf("Server error: %v", err)
+		}
 	}
-
-	<-exitChan
 }

cmd/ls-api/main.go

@@ -5,11 +5,12 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"io"
+	"net/http"
+
 	"github.com/go-chi/chi"
 	"github.com/go-chi/chi/middleware"
 	log "github.com/sirupsen/logrus"
-	"io"
-	"net/http"
 )
 
 const apiPort = 9563
@@ -33,7 +34,7 @@ func main() {
 		invokeRequest, _ := json.Marshal(InvokeRequest{InvokeId: uid, Payload: "{\"counter\":0}"})
 		_, err := http.Post(invokeUrl, "application/json", bytes.NewReader(invokeRequest))
 		if err != nil {
-			log.Fatal(err)
+			log.Error(err)
 		}
 
 		w.WriteHeader(200)
@@ -47,7 +48,7 @@ func main() {
 		invokeRequest, _ := json.Marshal(InvokeRequest{InvokeId: uid, Payload: "{\"counter\":0, \"fail\": \"yes\"}"})
 		_, err := http.Post(invokeUrl, "application/json", bytes.NewReader(invokeRequest))
 		if err != nil {
-			log.Fatal(err)
+			log.Error(err)
 		}
 
 		w.WriteHeader(200)
@@ -57,6 +58,7 @@ func main() {
 		}
 	})
 
+	log.Infof("Listening on port :%d", listenPort)
 	err := http.ListenAndServe(fmt.Sprintf(":%d", listenPort), router)
 	if err != nil {
 		log.Fatal(err)
@@ -74,7 +76,7 @@ func invokeLogsHandler(w http.ResponseWriter, r *http.Request) {
 }
 
 type InvokeRequest struct {
-	InvokeId string `json:"invoke-id"`
+	InvokeId string `json:"request-id"`
 	Payload  string `json:"payload"`
 }