From e72db8878a859495b1a1cdc959e0a2d03608e54f Mon Sep 17 00:00:00 2001 From: archandatta <35818003+archandatta@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:39:39 +0000 Subject: [PATCH 1/3] Drive the S2 storage sink from the controller The writer still opens at boot for image-provisioned instances and at fork-identity apply for forks; the controller now owns single-start, in-flight tracking, and the shutdown drain. Co-Authored-By: Claude Opus 5 --- server/cmd/api/main.go | 69 +++++++++++------------------------------- 1 file changed, 17 insertions(+), 52 deletions(-) diff --git a/server/cmd/api/main.go b/server/cmd/api/main.go index 669333ed..19be3d3d 100644 --- a/server/cmd/api/main.go +++ b/server/cmd/api/main.go @@ -12,8 +12,6 @@ import ( "os/exec" "os/signal" "strings" - "sync" - "sync/atomic" "syscall" "time" @@ -139,37 +137,22 @@ 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. + streamFn := func() string { + stream, _ := appliedS2Stream(config) + return stream } + s2Storage := events.NewS2StorageController(eventStream, config.S2Basin, config.S2AccessToken, streamFn, 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) } } @@ -220,14 +203,10 @@ func main() { // export is turned on per session, which the platform does after the // 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) + onForkIdentityApplied := func(forkidentity.Payload) { + // The handler contract forbids blocking the handoff on an optional sink. go func() { - defer s2Starting.Done() - if err := startS2Writer(stream); err != nil { + if err := s2Storage.Start(ctx); err != nil { slogger.Error("failed to start S2 storage writer for fork identity", "err", err) } }() @@ -463,27 +442,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 From 93a5726bbd4a2b7ec30484a72f8c3aec402c1d33 Mon Sep 17 00:00:00 2001 From: archandatta <35818003+archandatta@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:39:50 +0000 Subject: [PATCH 2/3] Hand the S2 storage controller to the api service Nothing reads it yet; it is the handle a later telemetry-side change needs to decide whether the sink may open. Co-Authored-By: Claude Opus 5 --- server/cmd/api/api/api.go | 14 ++++++++++++++ server/cmd/api/api/api_test.go | 2 +- server/cmd/api/api/display_test.go | 2 +- server/cmd/api/api/telemetry_test.go | 2 +- server/cmd/api/main.go | 1 + 5 files changed, 18 insertions(+), 3 deletions(-) diff --git a/server/cmd/api/api/api.go b/server/cmd/api/api/api.go index 6faa1b6b..d7a528d2 100644 --- a/server/cmd/api/api/api.go +++ b/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, diff --git a/server/cmd/api/api/api_test.go b/server/cmd/api/api/api_test.go index 43291073..8bd599b6 100644 --- a/server/cmd/api/api/api_test.go +++ b/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) { diff --git a/server/cmd/api/api/display_test.go b/server/cmd/api/api/display_test.go index d4838276..e012ba27 100644 --- a/server/cmd/api/api/display_test.go +++ b/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 } diff --git a/server/cmd/api/api/telemetry_test.go b/server/cmd/api/api/telemetry_test.go index cb14eebc..3e11e4fd 100644 --- a/server/cmd/api/api/telemetry_test.go +++ b/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 diff --git a/server/cmd/api/main.go b/server/cmd/api/main.go index 19be3d3d..431d1aaf 100644 --- a/server/cmd/api/main.go +++ b/server/cmd/api/main.go @@ -222,6 +222,7 @@ func main() { eventStream, config.DisplayNum, otlpExporter, + s2Storage, ) if err != nil { slogger.Error("failed to create api service", "err", err) From 68dabd8cc945c41948f547c23aff0a8577181467 Mon Sep 17 00:00:00 2001 From: archandatta <35818003+archandatta@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:09:55 +0000 Subject: [PATCH 3/3] Preserve applied S2 stream in fork hook --- server/cmd/api/fork_identity_test.go | 24 +++++++++++++ server/cmd/api/main.go | 52 +++++++++++++++++++++------- 2 files changed, 64 insertions(+), 12 deletions(-) diff --git a/server/cmd/api/fork_identity_test.go b/server/cmd/api/fork_identity_test.go index 854a1654..a8345380 100644 --- a/server/cmd/api/fork_identity_test.go +++ b/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) { diff --git a/server/cmd/api/main.go b/server/cmd/api/main.go index 431d1aaf..30285c72 100644 --- a/server/cmd/api/main.go +++ b/server/cmd/api/main.go @@ -12,6 +12,7 @@ import ( "os/exec" "os/signal" "strings" + "sync/atomic" "syscall" "time" @@ -143,11 +144,8 @@ func main() { // 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. - streamFn := func() string { - stream, _ := appliedS2Stream(config) - return stream - } - s2Storage := events.NewS2StorageController(eventStream, config.S2Basin, config.S2AccessToken, streamFn, events.S2Config{}, slogger) + 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 @@ -203,13 +201,8 @@ func main() { // export is turned on per session, which the platform does after the // handoff. An export started before then keeps the source's resource // attributes until it is restarted. - onForkIdentityApplied := func(forkidentity.Payload) { - // The handler contract forbids blocking the handoff on an optional sink. - go func() { - if err := s2Storage.Start(ctx); err != nil { - slogger.Error("failed to start S2 storage writer for fork identity", "err", err) - } - }() + onForkIdentityApplied := func(payload forkidentity.Payload) { + s2Streams.StartForAppliedPayload(ctx, payload, s2Storage.Start, slogger) } apiService, err := api.New( @@ -471,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