Fix concurrent map crash in dev server's Lambda bridge

cmd/sst/mosaic/aws/function.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -80,13 +81,34 @@ func function(ctx context.Context, input input) {
 	workerShutdownChan := make(chan *WorkerInfo, 1000)
 	nextChan := map[string]chan io.Reader{}
 	workers := map[string]*WorkerInfo{}
+	// nextChan and workers are written by the event loop below while the
+	// lambda runtime-API handlers read them from request goroutines; a bare
+	// map access on either side crashes the process with a concurrent
+	// map read/write fatal under cold-start bursts.
+	var stateMu sync.Mutex
+	loadNextChan := func(workerID string) chan io.Reader {
+		stateMu.Lock()
+		defer stateMu.Unlock()
+		ch, ok := nextChan[workerID]
+		if !ok {
+			ch = make(chan io.Reader, 100)
+			nextChan[workerID] = ch
+		}
+		return ch
+	}
+	loadWorker := func(workerID string) (*WorkerInfo, bool) {
+		stateMu.Lock()
+		defer stateMu.Unlock()
+		info, ok := workers[workerID]
+		return info, ok
+	}
 	evts := bus.Subscribe(&watcher.FileChangedEvent{}, &project.CompleteEvent{}, &runtime.BuildInput{}, &FunctionInvokedEvent{})
 	go fileLogger(input.project)
 
 	input.server.Mux.HandleFunc(`/lambda/{workerID}/2018-06-01/runtime/invocation/next`, func(w http.ResponseWriter, r *http.Request) {
 		log.Info("got next request", "workerID", r.PathValue("workerID"))
 		workerID := r.PathValue("workerID")
-		ch := nextChan[workerID]
+		ch := loadNextChan(workerID)
 		select {
 		case <-r.Context().Done():
 			log.Info("worker disconnected", "workerID", workerID)
@@ -105,7 +127,7 @@ func function(ctx context.Context, input input) {
 			var buf bytes.Buffer
 			tee := io.TeeReader(resp.Body, &buf)
 			io.Copy(w, tee)
-			workerInfo, ok := workers[workerID]
+			workerInfo, ok := loadWorker(workerID)
 			if ok {
 				bus.Publish(&FunctionInvokedEvent{
 					FunctionID: workerInfo.FunctionID,
@@ -126,7 +148,7 @@ func function(ctx context.Context, input input) {
 		io.Copy(writer, tee)
 		writer.Close()
 		w.WriteHeader(200)
-		info, ok := workers[workerID]
+		info, ok := loadWorker(workerID)
 		if ok {
 			fee := &FunctionErrorEvent{
 				FunctionID: info.FunctionID,
@@ -148,7 +170,7 @@ func function(ctx context.Context, input input) {
 		io.Copy(writer, tee)
 		writer.Close()
 		w.WriteHeader(202)
-		info, ok := workers[workerID]
+		info, ok := loadWorker(workerID)
 		if ok {
 			bus.Publish(&FunctionResponseEvent{
 				FunctionID: info.FunctionID,
@@ -170,7 +192,7 @@ func function(ctx context.Context, input input) {
 		io.Copy(writer, tee)
 		writer.Close()
 		w.WriteHeader(202)
-		info, ok := workers[workerID]
+		info, ok := loadWorker(workerID)
 		if ok {
 			fee := &FunctionErrorEvent{
 				FunctionID: info.FunctionID,
@@ -244,16 +266,21 @@ func function(ctx context.Context, input input) {
 			scanner := bufio.NewScanner(logs)
 			for scanner.Scan() {
 				line := scanner.Text()
+				stateMu.Lock()
+				requestID := info.CurrentRequestID
+				stateMu.Unlock()
 				bus.Publish(&FunctionLogEvent{
 					FunctionID: functionID,
 					WorkerID:   workerID,
-					RequestID:  info.CurrentRequestID,
+					RequestID:  requestID,
 					Line:       line,
 				})
 			}
 			workerShutdownChan <- info
 		}()
+		stateMu.Lock()
 		workers[workerID] = info
+		stateMu.Unlock()
 
 		return true
 	}
@@ -265,19 +292,15 @@ func function(ctx context.Context, input input) {
 		case msg := <-input.msg:
 			switch msg.Type {
 			case bridge.MessageInit:
-				ch, ok := nextChan[msg.Source]
-				if !ok {
-					ch = make(chan io.Reader, 100)
-					nextChan[msg.Source] = ch
-				}
+				loadNextChan(msg.Source)
 				init := bridge.InitBody{}
 				json.NewDecoder(msg.Body).Decode(&init)
 				if _, ok := targets[init.FunctionID]; !ok {
 					log.Error("function not found", "functionID", init.FunctionID)
 					continue
 				}
 				workerID := msg.Source
-				if _, ok := workers[workerID]; ok {
+				if _, ok := loadWorker(workerID); ok {
 					log.Error("got reboot but worker already exists", "workerID", workerID, "functionID", init.FunctionID)
 					continue
 				}
@@ -317,12 +340,8 @@ func function(ctx context.Context, input input) {
 				writer := input.client.NewWriter(bridge.MessagePing, input.prefix+"/"+msg.Source+"/in")
 				json.NewEncoder(writer).Encode(bridge.PingBody{})
 				writer.Close()
-				ch, ok := nextChan[msg.Source]
-				if !ok {
-					ch = make(chan io.Reader, 100)
-					nextChan[msg.Source] = ch
-				}
-				_, ok = workers[msg.Source]
+				ch := loadNextChan(msg.Source)
+				_, ok := loadWorker(msg.Source)
 				if !ok {
 					log.Info("asking for reboot", "workerID", msg.Source)
 					writer := input.client.NewWriter(bridge.MessageReboot, input.prefix+"/"+msg.Source+"/in")
@@ -335,25 +354,26 @@ func function(ctx context.Context, input input) {
 
 		case info := <-workerShutdownChan:
 			log.Info("worker died", "workerID", info.WorkerID)
+			stateMu.Lock()
 			existing, ok := workers[info.WorkerID]
-			if !ok {
-				continue
-			}
 			// only delete if a new worker hasn't already been started
-			if existing == info {
+			if ok && existing == info {
 				log.Info("deleting worker", "workerID", info.WorkerID)
 				delete(workers, info.WorkerID)
 				delete(nextChan, info.WorkerID)
 			}
+			stateMu.Unlock()
 			break
 		case unknown := <-evts:
 			switch evt := unknown.(type) {
 			case *FunctionInvokedEvent:
-				info, ok := workers[evt.WorkerID]
+				info, ok := loadWorker(evt.WorkerID)
 				if !ok {
 					continue
 				}
+				stateMu.Lock()
 				info.CurrentRequestID = evt.RequestID
+				stateMu.Unlock()
 			case *project.CompleteEvent:
 				if evt.Old {
 					continue

cmd/sst/mosaic/aws/function_test.go

@@ -0,0 +1,72 @@
+package aws
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/sst/sst/v3/cmd/sst/mosaic/aws/bridge"
+	"github.com/sst/sst/v3/pkg/server"
+)
+
+// The lambda runtime-API handlers read worker state from request goroutines
+// while the event loop writes it. Driving both sides concurrently surfaces
+// the conflict under `go test -race`; the same interleaving without the race
+// detector crashes the dev server with a concurrent map fatal (#6567).
+//
+// MessageInit for an unknown function id exercises the event loop's map
+// write and returns before any bridge or project access, so the test needs
+// no infrastructure.
+func TestFunctionConcurrentInitAndNext(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	srv := &server.Server{Mux: http.NewServeMux()}
+	msg := make(chan bridge.Message)
+	go function(ctx, input{server: srv, msg: msg})
+
+	send := func(worker int) {
+		m := bridge.Message{
+			Type:   bridge.MessageInit,
+			Source: fmt.Sprintf("worker-%d", worker),
+			Body:   strings.NewReader(`{"functionID":"does-not-exist"}`),
+		}
+		select {
+		case msg <- m:
+		case <-time.After(5 * time.Second):
+			t.Error("event loop stopped accepting messages")
+		}
+	}
+
+	// The first accepted message proves the handlers are registered: the
+	// unbuffered channel is only read once function's event loop is running,
+	// which happens after registration.
+	send(0)
+
+	requests, stopRequests := context.WithCancel(context.Background())
+	var handlers sync.WaitGroup
+	for i := 0; i < 64; i++ {
+		handlers.Add(1)
+		go func(worker int) {
+			defer handlers.Done()
+			r := httptest.NewRequest(
+				http.MethodGet,
+				fmt.Sprintf("/lambda/worker-%d/2018-06-01/runtime/invocation/next", worker%8),
+				nil,
+			).WithContext(requests)
+			srv.Mux.ServeHTTP(httptest.NewRecorder(), r)
+		}(i)
+	}
+
+	for i := 1; i < 512; i++ {
+		send(i % 8)
+	}
+
+	stopRequests()
+	handlers.Wait()
+}

cmd/sst/mosaic/errors/errors.go

@@ -8,7 +8,6 @@ import (
 	"github.com/sst/sst/v3/cmd/sst/mosaic/aws"
 	"github.com/sst/sst/v3/cmd/sst/mosaic/aws/appsync"
 	"github.com/sst/sst/v3/internal/util"
-	"github.com/sst/sst/v3/pkg/js"
 	"github.com/sst/sst/v3/pkg/project"
 	"github.com/sst/sst/v3/pkg/project/provider"
 	"github.com/sst/sst/v3/pkg/server"
@@ -36,7 +35,6 @@ var transformers = []ErrorTransformer{
 	exact(project.ErrProtectedStage, "Cannot remove protected stage. To remove a protected stage edit your sst.config.ts and remove the `protect` property."),
 	exact(provider.ErrLockNotFound, "This app / stage is not locked"),
 	exact(aws.ErrAppsyncNotReady, "SST creates an appsync event api to power live lambda. After 10 seconds of waiting this cli could not connect to it."),
-	exact(js.ErrTopLevelImport, "Your sst.config.ts has top level imports - this is not allowed. Move imports inside the function they are used and do a dynamic import: `const mod = await import(\"./mod\")`"),
 	match(func(err *project.ErrBuildFailed) string {
 		result := "Failed to build sst.config.ts"
 		for _, msg := range err.Errors {

pkg/js/js.go

@@ -11,8 +11,6 @@ import (
 	esbuild "github.com/evanw/esbuild/pkg/api"
 )
 
-var ErrTopLevelImport = fmt.Errorf("ErrTopLevelImport")
-
 type EvalOptions struct {
 	Dir     string
 	Outfile string
@@ -59,7 +57,6 @@ func Build(input EvalOptions) (esbuild.BuildResult, error) {
 		outfile = filepath.Join(input.Dir, ".sst", "platform", fmt.Sprintf("sst.config.%v.mjs", time.Now().UnixMilli()))
 	}
 	slog.Info("esbuild building", "out", outfile)
-	var err error
 	result := esbuild.Build(esbuild.BuildOptions{
 		Banner: map[string]string{
 			"js": `
@@ -85,15 +82,24 @@ const __dirname = topLevelFileUrlToPath(new topLevelURL(".", import.meta.url))
 		},
 		Plugins: []esbuild.Plugin{
 			{
-				Name: "DisallowImports",
+				Name: "StubImports",
 				Setup: func(build esbuild.PluginBuild) {
 					build.OnResolve(esbuild.OnResolveOptions{Filter: ".*"}, func(args esbuild.OnResolveArgs) (esbuild.OnResolveResult, error) {
 						if input.Globals == "" && filepath.Base(args.Importer) == "sst.config.ts" && args.Kind == esbuild.ResolveJSImportStatement {
-							err = ErrTopLevelImport
-							return esbuild.OnResolveResult{}, ErrTopLevelImport
+							return esbuild.OnResolveResult{
+								Path:      args.Path,
+								Namespace: "stub",
+							}, nil
 						}
 						return esbuild.OnResolveResult{}, nil
 					})
+					build.OnLoad(esbuild.OnLoadOptions{Filter: ".*", Namespace: "stub"}, func(args esbuild.OnLoadArgs) (esbuild.OnLoadResult, error) {
+						contents := "module.exports = {}"
+						return esbuild.OnLoadResult{
+							Contents: &contents,
+							Loader:   esbuild.LoaderJS,
+						}, nil
+					})
 				},
 			},
 			{
@@ -138,9 +144,6 @@ const __dirname = topLevelFileUrlToPath(new topLevelURL(".", import.meta.url))
 		Bundle:   true,
 		Metafile: true,
 	})
-	if err != nil {
-		return esbuild.BuildResult{}, err
-	}
 	if len(result.Errors) > 0 {
 		for _, err := range result.Errors {
 			slog.Error("esbuild error", "text", err.Text)