Drive the S2 storage sink from its controller

server/cmd/api/api/api.go

@@ -43,6 +43,17 @@ type OTLPExporter interface {
 
 var _ OTLPExporter = (*events.OTLPExportController)(nil)
 
+// S2Storage controls the optional S2 storage sink, which no-ops when the VM
+// has no S2 credentials. Implemented by *events.S2StorageController.
+type S2Storage interface {
+	Start(ctx context.Context) error
+	Stop(ctx context.Context) error
+	Running() bool
+	EverStarted() bool
+}
+
+var _ S2Storage = (*events.S2StorageController)(nil)
+
 type webMCPClient interface {
 	Tools(ctx context.Context) ([]webmcpclient.Tool, error)
 	Invoke(ctx context.Context, toolRef string, input map[string]any) (webmcpclient.InvocationResult, error)
@@ -117,6 +128,7 @@ type ApiService struct {
 	telemetrySession *telemetry.TelemetrySession
 	cdpMonitor       cdpMonitorController
 	otlpExport       OTLPExporter
+	s2Storage        S2Storage
 	monitorMu        sync.Mutex
 	// exportMu serializes OTLP export reconciliation independently of monitorMu,
 	// so a toggle-off drain (bounded by otlpStopTimeout) never blocks concurrent
@@ -138,6 +150,7 @@ func New(
 	eventStream *events.EventStream,
 	displayNum int,
 	otlpExport OTLPExporter,
+	s2Storage S2Storage,
 ) (*ApiService, error) {
 	switch {
 	case recordManager == nil:
@@ -173,6 +186,7 @@ func New(
 		telemetrySession:  telemetrySession,
 		cdpMonitor:        mon,
 		otlpExport:        otlpExport,
+		s2Storage:         s2Storage,
 		webmcp:            webmcpclient.NewManager(upstreamMgr),
 		browserRepl:       newBrowserReplManager(),
 		lifecycleCtx:      ctx,

server/cmd/api/api/api_test.go

@@ -399,7 +399,7 @@ func newTelemetrySession(t *testing.T) (*telemetry.TelemetrySession, *events.Eve
 func newSvc(t *testing.T, mgr recorder.RecordManager) (*ApiService, error) {
 	t.Helper()
 	ts, es := newTelemetrySession(t)
-	return New(mgr, newMockFactory(), newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil)
+	return New(mgr, newMockFactory(), newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil, nil)
 }
 
 func TestApiService_PatchChromiumFlags(t *testing.T) {

server/cmd/api/api/display_test.go

@@ -36,7 +36,7 @@ func testFFmpegFactory(t *testing.T, tempDir string) recorder.FFmpegRecorderFact
 func newTestServiceWithFactory(t *testing.T, mgr recorder.RecordManager, factory recorder.FFmpegRecorderFactory) *ApiService {
 	t.Helper()
 	ts, es := newTelemetrySession(t)
-	svc, err := New(mgr, factory, newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil)
+	svc, err := New(mgr, factory, newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil, nil)
 	require.NoError(t, err)
 	return svc
 }

server/cmd/api/api/telemetry_test.go

@@ -408,7 +408,7 @@ func (m *mockRecordManager) StopAll(_ context.Context) error
 func newTestService(t *testing.T, mgr recorder.RecordManager) *ApiService {
 	t.Helper()
 	ts, es := newTelemetrySession(t)
-	svc, err := New(mgr, newMockFactory(), newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil)
+	svc, err := New(mgr, newMockFactory(), newTestUpstreamManager(), scaletozero.NewNoopController(), newMockNekoClient(t), ts, es, 0, nil, nil)
 	require.NoError(t, err)
 	svc.cdpMonitor = &stubCdpMonitor{}
 	return svc

server/cmd/api/fork_identity_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"log/slog"
 	"net/http"
@@ -268,6 +269,29 @@ func TestAppliedS2StreamIgnoresMarkerFromBeforeThisBoot(t *testing.T) {
 	assert.Equal(t, "seed-stream", stream)
 }
 
+func TestS2StreamResolverUsesHookPayloadWithoutReadyFile(t *testing.T) {
+	useTempForkIdentityFiles(t)
+	resolver := newS2StreamResolver(&config.Config{S2Stream: "seed-stream"})
+	started := make(chan string, 1)
+
+	resolver.StartForAppliedPayload(
+		context.Background(),
+		forkidentity.Payload{"s2_stream": "fork-stream"},
+		func(context.Context) error {
+			started <- resolver.Resolve()
+			return nil
+		},
+		slog.Default(),
+	)
+
+	select {
+	case stream := <-started:
+		assert.Equal(t, "fork-stream", stream)
+	case <-time.After(time.Second):
+		t.Fatal("S2 storage did not start")
+	}
+}
+
 // markForkIdentityWaitArmed stands in for the wrapper having entered the wait,
 // which is what makes an applied identity current for this boot.
 func markForkIdentityWaitArmed(t *testing.T) {

server/cmd/api/main.go

@@ -12,7 +12,6 @@ import (
 	"os/exec"
 	"os/signal"
 	"strings"
-	"sync"
 	"sync/atomic"
 	"syscall"
 	"time"
@@ -139,37 +138,19 @@ func main() {
 	}
 	// An instance that already took a fork identity keeps it across a restart of
 	// this process, which the env captured at boot does not reflect.
-	s2Stream, s2StreamApplied := appliedS2Stream(config)
+	_, s2StreamApplied := appliedS2Stream(config)
 
 	// Optional S2 storage sink. The append session is bound to one stream when
 	// it starts, and an instance still holding for a fork identity is carrying
 	// the stream name of the instance it was forked from, so defer the writer
-	// until the identity that owns the events arrives. Opening that session
-	// dials with no deadline while holding the writer's lock, so an in-flight
-	// start is tracked for shutdown rather than blocking Stop behind the dial.
-	var s2Writer atomic.Pointer[events.S2StorageWriter]
-	var s2Starting sync.WaitGroup
-	startS2Writer := func(streamName string) error {
-		if config.S2Basin == "" || config.S2AccessToken == "" || streamName == "" || s2Writer.Load() != nil {
-			return nil
-		}
-		w := events.NewS2StorageWriter(eventStream, config.S2Basin, config.S2AccessToken, streamName, events.S2Config{}, slogger)
-		if !s2Writer.CompareAndSwap(nil, w) {
-			return nil
-		}
-		slogger.Info("S2 storage enabled", "basin", config.S2Basin, "stream", streamName)
-		if err := w.Start(ctx); err != nil {
-			// Leave the slot empty so a later identity can still open a writer.
-			s2Writer.CompareAndSwap(w, nil)
-			return err
-		}
-		return nil
-	}
+	// until the identity that owns the events arrives.
+	s2Streams := newS2StreamResolver(config)
+	s2Storage := events.NewS2StorageController(eventStream, config.S2Basin, config.S2AccessToken, s2Streams.Resolve, events.S2Config{}, slogger)
 	if !forkIdentityWait || s2StreamApplied {
 		// An optional sink that cannot open must not take the browser down: the
 		// api runs under supervisord with autorestart, so exiting here would
 		// crashloop the VM over a misconfigured basin or token.
-		if err := startS2Writer(s2Stream); err != nil {
+		if err := s2Storage.Start(ctx); err != nil {
 			slogger.Error("failed to start S2 storage writer, continuing without it", "err", err)
 		}
 	}
@@ -221,16 +202,7 @@ func main() {
 	// handoff. An export started before then keeps the source's resource
 	// attributes until it is restarted.
 	onForkIdentityApplied := func(payload forkidentity.Payload) {
-		// Opening the S2 append session dials the network with no deadline, so
-		// it runs off the handoff's critical path.
-		stream := forkidentity.FirstNonEmpty(forkidentity.Env(payload)["S2_STREAM"], config.S2Stream)
-		s2Starting.Add(1)
-		go func() {
-			defer s2Starting.Done()
-			if err := startS2Writer(stream); err != nil {
-				slogger.Error("failed to start S2 storage writer for fork identity", "err", err)
-			}
-		}()
+		s2Streams.StartForAppliedPayload(ctx, payload, s2Storage.Start, slogger)
 	}
 
 	apiService, err := api.New(
@@ -243,6 +215,7 @@ func main() {
 		eventStream,
 		config.DisplayNum,
 		otlpExporter,
+		s2Storage,
 	)
 	if err != nil {
 		slogger.Error("failed to create api service", "err", err)
@@ -463,27 +436,13 @@ func main() {
 		slogger.Error("server failed to shutdown", "err", err)
 	}
 
-	// s2Writer shuts down after the servers above, since they might produce events we
-	// want to capture into the stream; we must let them finish before closing the writer.
-	// Stop takes the same lock the unbounded append-session dial holds, so only
-	// drain once a start that is still opening has finished. Skipping the drain
-	// loses at most the events of a writer that was never serving.
-	s2StartSettled := make(chan struct{})
-	go func() {
-		s2Starting.Wait()
-		close(s2StartSettled)
-	}()
-	select {
-	case <-s2StartSettled:
-		if w := s2Writer.Load(); w != nil {
-			stopCtx, stopCancel := context.WithTimeout(context.Background(), 10*time.Second)
-			defer stopCancel()
-			if err := w.Stop(stopCtx); err != nil {
-				slogger.Error("s2 storage writer stop failed", "err", err)
-			}
-		}
-	case <-time.After(2 * time.Second):
-		slogger.Warn("s2 storage writer still opening at shutdown, skipping drain")
+	// S2 storage shuts down after the servers above, since they might produce
+	// events we want to capture into the stream; we must let them finish before
+	// closing the writer.
+	s2StopCtx, s2StopCancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer s2StopCancel()
+	if err := s2Storage.Stop(s2StopCtx); err != nil {
+		slogger.Error("s2 storage writer stop failed", "err", err)
 	}
 
 	// Likewise stop OTLP export after the servers drain (a no-op if the toggle
@@ -505,6 +464,41 @@ func mustFFmpeg() {
 	}
 }
 
+// s2StreamResolver uses the hook payload once applied and the persisted payload
+// after an API process restart.
+type s2StreamResolver struct {
+	cfg        *config.Config
+	hookStream atomic.Pointer[string]
+}
+
+func newS2StreamResolver(cfg *config.Config) *s2StreamResolver {
+	return &s2StreamResolver{cfg: cfg}
+}
+
+func (r *s2StreamResolver) Resolve() string {
+	if stream := r.hookStream.Load(); stream != nil {
+		return *stream
+	}
+	stream, _ := appliedS2Stream(r.cfg)
+	return stream
+}
+
+func (r *s2StreamResolver) StartForAppliedPayload(
+	ctx context.Context,
+	payload forkidentity.Payload,
+	start func(context.Context) error,
+	log *slog.Logger,
+) {
+	stream := forkidentity.FirstNonEmpty(forkidentity.Env(payload)["S2_STREAM"], r.cfg.S2Stream)
+	r.hookStream.Store(&stream)
+	// The handler contract forbids blocking the handoff on an optional sink.
+	go func() {
+		if err := start(ctx); err != nil {
+			log.Error("failed to start S2 storage writer for fork identity", "err", err)
+		}
+	}()
+}
+
 // appliedS2Stream resolves the stream the S2 writer should bind, preferring a
 // fork identity the guest has already taken over the env this process started
 // with. The env belongs to the instance this one was forked from, and it is what