Small refactoring

Author: gregfurman

cmd/localstack/main.go

@@ -11,8 +11,7 @@ import (
 	"runtime/debug"
 	"strconv"
 	"strings"
-	"sync"
-	"time"
+	"syscall"
 
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/localstack/lambda-runtime-init/internal/aws/lambda"
@@ -22,16 +21,17 @@ import (
 	"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/extensions"
 	"go.amzn.com/lambda/interop"
-	"go.amzn.com/lambda/rapidcore"
-	supv "go.amzn.com/lambda/supervisor"
+	"go.amzn.com/lambda/rapid"
 )
 
 func InitLsOpts() *localstack.Config {
@@ -66,7 +66,7 @@ func InitFunctionConfig() lambda.FunctionConfig {
 		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", "3008"),
+		FunctionMemorySizeMb: utils.GetEnvWithDefault("AWS_LAMBDA_FUNCTION_MEMORY_SIZE", "128"),
 		FunctionHandler:      utils.GetEnvWithDefault("AWS_LAMBDA_FUNCTION_HANDLER", os.Getenv("_HANDLER")),
 	}
 }
@@ -101,6 +101,10 @@ func UnsetLsEnvs() {
 	}
 }
 
+type closer struct{ fn func() }
+
+func (c *closer) Close() error { c.fn(); return nil }
+
 func main() {
 	// we're setting this to the same value as in the official RIE
 	debug.SetGCPercent(33)
@@ -116,31 +120,19 @@ func main() {
 	// 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)
-		xRayLogLevel = "error"
-	case "panic":
-		log.SetLevel(log.PanicLevel)
+	case log.ErrorLevel, log.FatalLevel, log.PanicLevel:
 		xRayLogLevel = "error"
-	default:
-		log.Fatal("Invalid value for LOCALSTACK_INIT_LOG_LEVEL")
 	}
 
 	// patch MaxPayloadSize
@@ -183,52 +175,62 @@ func main() {
 		}
 	}
 
-	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
 	defer stop()
 
-	// file watcher for hot-reloading
-	fileWatcherContext, cancelFileWatcher := context.WithCancel(ctx)
-	defer cancelFileWatcher()
-
-	// Custom Interop Server
-	defaultServer := rapidcore.NewServer()
+	// LocalStack client used for sending callbacks
 	lsClient := localstack.NewLocalStackClient(lsOpts.RuntimeEndpoint, lsOpts.RuntimeId)
-	interopServer := server.NewInteropServer(defaultServer, lsClient)
 
 	// Services required for Sandbox environment
+	interopServer := server.NewInteropServer(lsClient)
+	defer interopServer.Close()
+
 	logCollector := logging.NewLogCollector()
 	localStackLogsEgressApi := logging.NewLocalStackLogsEgressAPI(logCollector)
 	tracer := tracing.NewLocalStackTracer()
-	eventsListener := events.NewLocalStackEventsAPI(lsClient)
-
-	defaultSupv := supv.NewLocalSupervisor()
-	localStackSupv := supervisor.NewLocalStackSupervisor(ctx, defaultSupv, eventsListener)
+	lsEventsAPI := events.NewLocalStackEventsAPI(lsClient)
+	localStackSupv := supervisor.NewLocalStackSupervisor(ctx, lsEventsAPI)
 
 	// build sandbox
-	exitChan := make(chan struct{})
-	sandbox := rapidcore.
-		NewSandboxBuilder().
-		AddShutdownFunc(func() {
-			log.Debugln("Stopping file watcher")
-			cancelFileWatcher()
-		}).
-		AddShutdownFunc(func() {
-			exitChan <- struct{}{}
-		}).
-		SetExtensionsFlag(true).
-		SetInitCachingFlag(true).
-		SetLogsEgressAPI(localStackLogsEgressApi).
-		SetTracer(tracer).
-		SetInteropServer(interopServer).
-		SetSupervisor(localStackSupv).
-		SetHandler(handler)
+	sandboxConfig := rapid.Sandbox{
+		EnableTelemetryAPI: false,
+		StandaloneMode:     true,
+		InitCachingEnabled: true,
+
+		Tracer:        tracer,
+		EventsAPI:     lsEventsAPI,
+		Supervisor:    localStackSupv,
+		InteropServer: interopServer,
+		LogsEgressAPI: localStackLogsEgressApi,
+
+		Handler: handler,
+
+		RuntimeFsRootPath: "/",
+		RuntimeAPIHost:    "127.0.0.1",
+		RuntimeAPIPort:    9001,
+	}
+
+	extensions.Enable()
+
+	rapidCtx, internalStateFn, addr := rapid.Start(ctx, &sandboxConfig)
+	sandboxCtx, err := sandbox.CreateSandboxContext(rapidCtx, handler, addr)
+	if err != nil {
+		log.Fatalf("fatal error encountered when creating SandboxContext: %w", err)
+	}
+
+	// Populate our interop server
+	interopServer.SetSandboxContext(sandboxCtx)
+	interopServer.SetInternalStateGetter(internalStateFn)
 
 	// Start daemons
 
-	// Start hot-reloading watcher
+	// file watcher for hot-reloading
+	fileWatcherContext, cancelFileWatcher := context.WithCancel(ctx)
+	defer cancelFileWatcher()
+
 	go hotreloading.RunHotReloadingListener(interopServer, lsOpts.HotReloadingPaths, fileWatcherContext, lsOpts.FileWatcherStrategy)
 
-	// xray daemon
+	// Start xray daemon
 	endpoint := "http://" + net.JoinHostPort(lsOpts.LocalstackIP, lsOpts.EdgePort)
 	xrayConfig := xray.NewConfig(endpoint, xRayLogLevel)
 	d := xray.NewDaemon(xrayConfig, lsOpts.EnableXRayTelemetry == "1")
@@ -240,15 +242,11 @@ func main() {
 	}()
 	d.Run() // served async
 
-	// initialize all flows and start runtime API
-	sandboxContext, internalStateFn := sandbox.Create()
-	// Populate our interop server
-	interopServer.SetSandboxContext(sandboxContext)
-	interopServer.SetInternalStateGetter(internalStateFn)
-
+	// 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`
@@ -258,26 +256,18 @@ func main() {
 	}
 
 	invokeServer := server.NewServer(lsOpts.InteropPort, localStackService)
-	invokeServer.RegisterOnShutdown(localStackService.Close)
+	defer invokeServer.Close()
 
-	defer invokeServer.Shutdown(context.Background())
-
-	var wg sync.WaitGroup
-
-	wg.Add(1)
+	serverErr := make(chan error, 1)
 	go func() {
-		defer wg.Done()
 		listener, err := net.Listen("tcp", fmt.Sprintf(":%s", lsOpts.InteropPort))
-
 		if err != nil {
-			log.Fatalf("failed to start listener for custom interops server: %s", err)
+			log.Fatalf("failed to start LocalStack Lambda Runtime Interface server: %s", err)
 		}
-		go invokeServer.Serve(listener)
+		go func() { serverErr <- invokeServer.Serve(listener); close(serverErr) }()
 		log.Debugf("LocalStack API gateway listening on %s", listener.Addr().String())
 	}()
 
-	wg.Wait()
-
 	log.Debugln("Awaiting initialization of runtime init.")
 	if err := interopServer.AwaitInitialized(); err != nil {
 		// Error cases: ErrInitDoneFailed or ErrInitResetReceived
@@ -292,16 +282,13 @@ func main() {
 		}
 	}
 
+	// Block until context is cancelled OR the server errors out
 	select {
 	case <-ctx.Done():
-	case <-exitChan:
-	}
-
-	gracefulCtx, cancel := context.WithTimeout(ctx, time.Millisecond*500)
-	defer cancel()
-
-	if err := localStackService.AwaitCompleted(gracefulCtx); err != nil {
-		log.Warnf("Did not gracefully complete: %w", err)
+		log.Info("Shutdown signal received.")
+	case <-serverErr:
+		if err != nil {
+			log.Errorf("Server error: %v", err)
+		}
 	}
-
 }

internal/hotreloading/reloader.go

@@ -12,6 +12,7 @@ type Resetter interface {
 	Reset(reason string, timeoutMs int64) (*statejson.ResetDescription, error)
 }
 
+// TODO: Rework this into a struct/service where the ChangeListener is passed via dependency injection
 func RunHotReloadingListener(server Resetter, targetPaths []string, ctx context.Context, fileWatcherStrategy string) {
 	if len(targetPaths) == 1 && targetPaths[0] == "" {
 		log.Debugln("Hot reloading disabled.")
@@ -31,5 +32,4 @@ func RunHotReloadingListener(server Resetter, targetPaths []string, ctx context.
 
 	<-ctx.Done()
 	log.Infoln("Closing down filewatcher.")
-
 }

internal/sandbox/sandbox.go

@@ -0,0 +1,44 @@
+package sandbox
+
+import (
+	"fmt"
+	"reflect"
+	"unsafe"
+
+	"go.amzn.com/lambda/interop"
+	"go.amzn.com/lambda/rapidcore"
+)
+
+// CreateSandboxContext creates a SandboxContext using reflection
+// This is a dirty hack since the SandboxBuilder sets quite a few elements by default that cause issues
+func CreateSandboxContext(rapidCtx interop.RapidContext, handler string, runtimeAPIAddress string) (interop.SandboxContext, error) {
+	sandboxCtx := &rapidcore.SandboxContext{}
+	v := reflect.ValueOf(sandboxCtx).Elem()
+
+	setField := func(fieldName string, value interface{}) {
+		field := v.FieldByName(fieldName)
+		if !field.IsValid() {
+			return
+		}
+
+		var settableField reflect.Value
+		if field.CanSet() {
+			settableField = field
+		} else {
+
+			settableField = reflect.NewAt(field.Type(), unsafe.Pointer(field.UnsafeAddr())).Elem()
+		}
+
+		settableField.Set(reflect.ValueOf(value))
+	}
+
+	setField("rapidCtx", rapidCtx)
+	setField("handler", handler)
+	setField("runtimeAPIAddress", runtimeAPIAddress)
+
+	if sandboxCtx == nil || reflect.ValueOf(*sandboxCtx).IsZero() {
+		return nil, fmt.Errorf("failed to dynamically create SandboxContext.")
+	}
+
+	return sandboxCtx, nil
+}

internal/server/interop.go

@@ -24,9 +24,9 @@ type LocalStackInteropsServer struct {
 	localStackAdapter *localstack.LocalStackClient
 }
 
-func NewInteropServer(server *rapidcore.Server, ls *localstack.LocalStackClient) *LocalStackInteropsServer {
+func NewInteropServer(ls *localstack.LocalStackClient) *LocalStackInteropsServer {
 	return &LocalStackInteropsServer{
-		Server:            server,
+		Server:            rapidcore.NewServer(),
 		localStackAdapter: ls,
 	}
 }
@@ -166,3 +166,9 @@ func (c *LocalStackInteropsServer) SendInitErrorResponse(resp *interop.ErrorInvo
 
 	return c.Server.SendInitErrorResponse(resp)
 }
+
+func (c *LocalStackInteropsServer) Close() error {
+	log.Info("Shutting down...")
+	_, err := c.Reset("SandboxTerminated", 2000)
+	return err
+}

internal/server/service.go

@@ -259,9 +259,10 @@ func (ls *LocalStackService) AfterInvoke(ctx context.Context) error {
 	return nil
 }
 
-func (ls *LocalStackService) Close() {
+func (ls *LocalStackService) Close() error {
 	ls.isShuttingDown.Store(true)
 	log.Debug("Shutdown of LocalStackService triggered.")
+	return nil
 }
 
 func (ls *LocalStackService) shutdownTriggered() bool {
@@ -277,7 +278,10 @@ func (ls *LocalStackService) AwaitCompleted(ctx context.Context) error {
 		select {
 		case <-ctx.Done():
 			pending := atomic.LoadInt64(&ls.pendingCallbacks)
-			return fmt.Errorf("failed to gracefully complete all callbacks to LocalStack: %d remaining", pending)
+			if pending > 0 {
+				return fmt.Errorf("failed to gracefully complete all callbacks to LocalStack: %d remaining", pending)
+			}
+			return nil
 		case <-ls.allDone:
 			if atomic.LoadInt64(&ls.pendingCallbacks) == 0 {
 				return nil

internal/supervisor/supervisor.go

@@ -10,6 +10,7 @@ import (
 	"github.com/sirupsen/logrus"
 	"go.amzn.com/lambda/fatalerror"
 	"go.amzn.com/lambda/interop"
+	"go.amzn.com/lambda/supervisor"
 	"go.amzn.com/lambda/supervisor/model"
 )
 
@@ -24,10 +25,10 @@ type LocalStackSupervisor struct {
 	isShuttingDown *atomic.Bool
 }
 
-func NewLocalStackSupervisor(ctx context.Context, supv model.ProcessSupervisor, evs interop.EventsAPI) *LocalStackSupervisor {
+func NewLocalStackSupervisor(ctx context.Context, evs interop.EventsAPI) *LocalStackSupervisor {
 	var isShuttingDown atomic.Bool
 	ls := &LocalStackSupervisor{
-		ProcessSupervisor: supv,
+		ProcessSupervisor: supervisor.NewLocalSupervisor(),
 		eventsAPI:         evs,
 		eventsChan:        make(chan model.Event),
 		isShuttingDown:    &isShuttingDown,