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/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 669333ed..30285c72 100644 --- a/server/cmd/api/main.go +++ b/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