WIP: snapshotting etc.

Author: gregfurman

cmd/localstack/main.go

@@ -29,7 +29,6 @@ import (
 	"go.amzn.com/lambda/interop"
 	"go.amzn.com/lambda/rapidcore"
 	supv "go.amzn.com/lambda/supervisor"
-	"go.amzn.com/lambda/telemetry"
 )
 
 func InitLsOpts() *server.LsOpts {
@@ -197,32 +196,37 @@ func main() {
 	logCollector := logging.NewLogCollector()
 	localStackLogsEgressApi := logging.NewLocalStackLogsEgressAPI(logCollector)
 	tracer := tracing.NewLocalStackTracer()
-	eventsListener := events.NewEventsListener(lsAdapter, &telemetry.NoOpEventsAPI{}, logCollector, functionConf)
+	eventsListener := events.NewEventsListener(lsAdapter)
 
 	defaultSupv := supv.NewLocalSupervisor()
 	wrappedSupv := supervisor.NewLocalStackSupervisor(ctx, defaultSupv, eventsListener, interopServer.InternalState)
 
 	// build sandbox
+	exitChan := make(chan struct{})
 	sandbox := rapidcore.
 		NewSandboxBuilder().
-		//SetTracer(tracer).
 		AddShutdownFunc(func() {
 			log.Debugln("Stopping file watcher")
 			cancelFileWatcher()
 		}).
+		AddShutdownFunc(func() {
+			exitChan <- struct{}{}
+		}).
 		SetExtensionsFlag(true).
 		SetInitCachingFlag(true).
 		SetLogsEgressAPI(localStackLogsEgressApi).
 		SetTracer(tracer).
 		SetInteropServer(interopServer).
 		SetSupervisor(wrappedSupv).
-		SetEventsAPI(eventsListener)
+		SetHandler(handler)
+	sandbox.AddShutdownFunc(func() {
+		exitChan <- struct{}{}
+	})
 
-	// SetEventsAPI(eventsListener)
+	// Start daemons
 
-	// Externally set supervisor for metrics tracking
-	// sandbox.SetSupervisor(localSupervisor)
-	// sandbox.SetRuntimeFsRootPath(localSupervisor.RootPath)
+	// Start hot-reloading watcher
+	go hotreloading.RunHotReloadingListener(interopServer, lsOpts.HotReloadingPaths, fileWatcherContext, lsOpts.FileWatcherStrategy)
 
 	// xray daemon
 	endpoint := "http://" + net.JoinHostPort(lsOpts.LocalstackIP, lsOpts.EdgePort)
@@ -234,15 +238,7 @@ func main() {
 		log.Debugln("Flushing segments in xray daemon")
 		d.Close()
 	}()
-	d.Run() // async
-
-	if len(handler) > 0 {
-		sandbox.SetHandler(handler)
-	}
-	exitChan := make(chan struct{})
-	sandbox.AddShutdownFunc(func() {
-		exitChan <- struct{}{}
-	})
+	d.Run() // served async
 
 	// initialize all flows and start runtime API
 	sandboxContext, internalStateFn := sandbox.Create()
@@ -278,12 +274,6 @@ func main() {
 		log.Debugf("LocalStack API gateway listening on %s", listener.Addr().String())
 	}()
 
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		hotreloading.RunHotReloadingListener(interopServer, lsOpts.HotReloadingPaths, fileWatcherContext, lsOpts.FileWatcherStrategy)
-	}()
-
 	wg.Wait()
 
 	log.Debugln("Awaiting initialization of runtime init.")

internal/events/events.go

@@ -3,33 +3,23 @@ package events
 import (
 	"encoding/json"
 	"fmt"
-	"time"
 
-	"github.com/localstack/lambda-runtime-init/internal/logging"
 	"github.com/localstack/lambda-runtime-init/internal/server"
-	log "github.com/sirupsen/logrus"
 
 	"go.amzn.com/lambda/interop"
 	"go.amzn.com/lambda/telemetry"
 )
 
 type EventsListener struct {
 	interop.EventsAPI
-
-	adapter      *server.LocalStackAdapter
-	logCollector *logging.LogCollector
-
-	function server.FunctionConfig
-
-	initStart time.Time
-	initEnd   time.Time
+	adapter *server.LocalStackAdapter
 }
 
-func NewEventsListener(adapter *server.LocalStackAdapter, eventsAPI interop.EventsAPI, logger *logging.LogCollector, cfg server.FunctionConfig) *EventsListener {
+func NewEventsListener(adapter *server.LocalStackAdapter) *EventsListener {
 	return &EventsListener{
-		adapter:      adapter,
-		EventsAPI:    eventsAPI,
-		logCollector: logger,
+		adapter: adapter,
+		// For now just use the no-ops API to satisfy the EventsAPI interface
+		EventsAPI: &telemetry.NoOpEventsAPI{},
 	}
 
 }
@@ -47,60 +37,3 @@ func (ev *EventsListener) SendFault(data interop.FaultData) error {
 
 	return ev.adapter.SendStatus(server.Error, payload)
 }
-
-func (ev *EventsListener) SendInitStart(_ interop.InitStartData) error {
-	ev.initStart = time.Now() // need to use this event to retrieve when the INIT start happens
-	return nil
-}
-
-func (ev *EventsListener) SendInitRuntimeDone(done interop.InitRuntimeDoneData) error {
-	ev.initEnd = time.Now()
-	initDurationMs := float64(ev.initEnd.Sub(ev.initStart)) / float64(time.Millisecond) // convert µs -> ms
-	switch done.Status {
-	case telemetry.RuntimeDoneSuccess:
-		return ev.adapter.SendStatus(server.Ready, []byte{})
-	case telemetry.RuntimeDoneError:
-		if done.ErrorType != nil {
-			if done.Phase != telemetry.InitInsideInitPhase {
-				log.WithField("ErrorType", *done.ErrorType).Debug("Failed to initialise runtime during INIT phase. This will be retried during the next INVOKE.")
-			}
-			_, _ = fmt.Fprintf(ev.logCollector, "INIT_REPORT Init Duration: %.2f ms Phase: init Status: %s", initDurationMs, done.Status)
-			// return ev.adapter.SendStatus(server.Error, []byte(*done.ErrorType))
-		}
-	}
-	return nil
-}
-
-func (ev *EventsListener) SendInvokeStart(invoke interop.InvokeStartData) error {
-	_, _ = fmt.Fprintf(ev.logCollector, "START RequestId: %s Version: %s\n", invoke.RequestID, invoke.Version)
-	return nil
-}
-
-func (ev *EventsListener) SendInvokeRuntimeDone(invoke interop.InvokeRuntimeDoneData) error {
-	initDurationMs := float64(ev.initEnd.Sub(ev.initStart)) / float64(time.Millisecond) // convert µs -> ms
-	memUsed := invoke.Metrics.ProducedBytes / (1024 * 1024)                             // bytes -> mb
-
-	report := fmt.Sprintf("REPORT RequestId: %s\t"+
-		"Duration: %.2f ms\t"+
-		"Billed Duration: %.0f ms\t"+
-		"Memory Size: %s MB\t"+
-		"Max Memory Used: %d MB\t"+
-		"Init Duration: %.2f ms",
-		invoke.RequestID, invoke.Metrics.DurationMs, ev.function.FunctionMemorySizeMb, memUsed, initDurationMs)
-
-	_, _ = fmt.Fprintln(ev.logCollector, report)
-
-	// we dont send the logs from here since we may want to delay sending logs before we know for certain that
-	// the /response has been sent
-
-	return nil
-}
-
-func (ev *EventsListener) Reset() {
-	ev.initStart = time.Time{}
-	ev.initEnd = time.Time{}
-}
-
-// func (ev *EventsListener) SendInitReport(data interop.InitReportData) error {
-
-// }

internal/logging/collector.go

@@ -10,40 +10,27 @@ type LogResponse struct {
 }
 
 type LogCollector struct {
-	mutex       *sync.Mutex
-	RuntimeLogs []string
+	logs strings.Builder
+	mu   *sync.Mutex
 }
 
 func NewLogCollector() *LogCollector {
 	return &LogCollector{
-		RuntimeLogs: []string{},
-		mutex:       &sync.Mutex{},
+		mu: &sync.Mutex{},
 	}
 }
 
 func (lc *LogCollector) Write(p []byte) (n int, err error) {
-	lc.Put(string(p))
-	return len(p), nil
-}
-
-func (lc *LogCollector) Put(line string) {
-	lc.mutex.Lock()
-	defer lc.mutex.Unlock()
-	lc.RuntimeLogs = append(lc.RuntimeLogs, line)
-}
-
-func (lc *LogCollector) reset() {
-	lc.mutex.Lock()
-	defer lc.mutex.Unlock()
-	lc.RuntimeLogs = []string{}
+	lc.mu.Lock()
+	defer lc.mu.Unlock()
+	n, err = lc.logs.Write(p)
+	return
 }
 
 func (lc *LogCollector) GetLogs() LogResponse {
-	lc.mutex.Lock()
-	defer lc.mutex.Unlock()
-	response := LogResponse{
-		Logs: strings.Join(lc.RuntimeLogs, ""),
-	}
-	lc.RuntimeLogs = []string{}
-	return response
+	lc.mu.Lock()
+	defer lc.mu.Unlock()
+	resp := LogResponse{Logs: lc.logs.String()}
+	lc.logs.Reset()
+	return resp
 }

internal/server/handler.go

@@ -60,12 +60,10 @@ func InvokeHandler(api *LocalStackService) http.HandlerFunc {
 				log.Fatalln("unable to marshal json timeout response")
 			}
 			response = errorResponseJson
-			// default:
-			// 	log.Fatalln(err)
 		}
 
 		api.AfterInvoke(r.Context())
-		api.CollectAndSendLogs(req.InvokeId)
+		api.ForwardLogs(req.InvokeId)
 
 		invokeRespDecoder := json.NewDecoder(bytes.NewReader(response))
 		invokeRespDecoder.DisallowUnknownFields()

internal/server/interop.go

@@ -14,7 +14,6 @@ import (
 	"github.com/sirupsen/logrus"
 	log "github.com/sirupsen/logrus"
 	"go.amzn.com/lambda/core/directinvoke"
-	"go.amzn.com/lambda/fatalerror"
 	"go.amzn.com/lambda/interop"
 	"go.amzn.com/lambda/metering"
 	"go.amzn.com/lambda/rapi/model"
@@ -60,13 +59,14 @@ func NewInteropServer(server *rapidcore.Server, ls *LocalStackAdapter) *CustomIn
 	return &CustomInteropServer{
 		Server:            server,
 		localStackAdapter: ls,
+		mutex:             &sync.Mutex{},
 	}
 }
 
 type CustomInteropServer struct {
 	*rapidcore.Server
 	localStackAdapter *LocalStackAdapter
-	initOnce          sync.Once
+	mutex             *sync.Mutex
 }
 
 func (c *CustomInteropServer) Invoke(responseWriter http.ResponseWriter, invoke *interop.Invoke) error {
@@ -77,69 +77,78 @@ func (c *CustomInteropServer) Invoke(responseWriter http.ResponseWriter, invoke
 	defer close(releaseRespChan)
 
 	go func() {
-		_, err := c.Server.Reserve(invoke.ID, invoke.TraceID, invoke.LambdaSegmentID)
-		if err != nil && !errors.Is(err, rapidcore.ErrAlreadyReserved) {
+		reserveResp, err := c.Server.Reserve(invoke.ID, invoke.TraceID, invoke.LambdaSegmentID)
+		if err != nil {
 			releaseRespChan <- err
 			return
 		}
 
-		invoke.DeadlineNs = fmt.Sprintf("%d", metering.Monotime()+c.Server.GetInvokeTimeout().Nanoseconds())
+		invoke.DeadlineNs = fmt.Sprintf("%d", metering.Monotime()+reserveResp.Token.FunctionTimeout.Nanoseconds())
+
+		// Wait for initialization to complete
+		if err := c.Server.AwaitInitialized(); err != nil {
+			switch err {
+			case rapidcore.ErrInitDoneFailed:
+				// Init failed, reset and continue with suppressed init
+				if _, resetErr := c.Server.Reset("InitFailed", 2000); resetErr != nil {
+					log.Errorf("Reset failed: %v", resetErr)
+				}
+				// Reserve again after reset for suppressed init
+				if _, reserveErr := c.Server.Reserve(invoke.ID, invoke.TraceID, invoke.LambdaSegmentID); reserveErr != nil {
+					releaseRespChan <- reserveErr
+					return
+				}
+			default:
+				releaseRespChan <- err
+				return
+			}
+		}
+
 		go func() {
 			isDirect := directinvoke.MaxDirectResponseSize > interop.MaxPayloadSize
 			if err := c.Server.FastInvoke(responseWriter, invoke, isDirect); err != nil {
 				log.Debugf("FastInvoke() error: %s", err)
 			}
 		}()
 
-		// Waits for an invocation to finish before calling Release()
 		_, err = c.Server.AwaitRelease()
 		if err != nil {
-			log.Debugf("AwaitRelease() error: %s", err)
 			switch err {
 			case rapidcore.ErrReleaseReservationDone:
-				// not an error, expected return value when Reset is called
+				// Expected when Reset is called
+				releaseRespChan <- nil
 				return
 			case rapidcore.ErrInitDoneFailed, rapidcore.ErrInvokeDoneFailed:
-				c.Server.Reset("ReleaseFail", 2000)
+				if _, resetErr := c.Server.Reset("ReleaseFail", 2000); resetErr != nil {
+					log.Errorf("Reset failed: %v", resetErr)
+				}
+				releaseRespChan <- err
+				return
+			default:
+				if _, resetErr := c.Server.Reset("UnexpectedError", 2000); resetErr != nil {
+					log.Errorf("Reset failed: %v", resetErr)
+				}
+				releaseRespChan <- err
+				return
 			}
 		}
-		releaseRespChan <- err
+
+		releaseRespChan <- nil
 	}()
 
 	select {
 	case err := <-releaseRespChan:
-		if err != nil {
-			return err
-		}
+		return err
 	case <-ctx.Done():
 		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
-			c.Server.Reset("Timeout", 2000)
+			if _, resetErr := c.Server.Reset("Timeout", 2000); resetErr != nil {
+				log.Errorf("Reset failed: %v", resetErr)
+			}
 			return rapidcore.ErrInvokeTimeout
 		}
 	}
 
-	return c.Server.Release()
-}
-
-func (c *CustomInteropServer) AwaitInitialized() error {
-	err := c.Server.AwaitInitialized()
-	go func() {
-		state, stateErr := c.InternalState()
-		if stateErr != nil {
-			return
-		}
-
-		doneResponse := &interop.Done{}
-		if state.FirstFatalError != "" {
-			doneResponse.ErrorType = fatalerror.GetValidRuntimeOrFunctionErrorType(state.FirstFatalError)
-		}
-
-		c.InitDoneChan <- rapidcore.DoneWithState{
-			Done:  doneResponse,
-			State: c.InternalStateGetter(),
-		}
-	}()
-	return err
+	return nil
 }
 
 func (c *CustomInteropServer) SendInitErrorResponse(resp *interop.ErrorInvokeResponse) error {

internal/server/model.go

@@ -61,6 +61,7 @@ type ErrorResponse struct {
 	ErrorType    string   `json:"errorType,omitempty"`
 	RequestId    *string  `json:"requestId,omitempty"`
 	StackTrace   []string `json:"stackTrace,omitempty"`
+	Trace        []string `json:"trace,omitempty"`
 }
 
 type ResponseWriterProxy struct {