Point event sinks at the identity a fork actually takes

server/cmd/api/fork_identity.go

@@ -11,7 +11,10 @@ import (
 	"github.com/kernel/kernel-images/server/lib/forkidentity"
 )
 
-func forkIdentityHandler(log *slog.Logger) http.HandlerFunc {
+// onApplied runs once the guest has taken the identity, so services this
+// process owns can retarget themselves to it. It is called with the applied
+// payload and only on success; pass nil when there is nothing to retarget.
+func forkIdentityHandler(log *slog.Logger, onApplied func(forkidentity.Payload)) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		enabled, err := forkidentity.WaitEnabled()
 		if err != nil {
@@ -58,6 +61,9 @@ func forkIdentityHandler(log *slog.Logger) http.HandlerFunc {
 			http.Error(w, "fork identity not applied", http.StatusInternalServerError)
 			return
 		}
+		if onApplied != nil {
+			onApplied(payload)
+		}
 		w.WriteHeader(http.StatusNoContent)
 	}
 }

server/cmd/api/fork_identity_test.go

@@ -96,7 +96,7 @@ func TestForkIdentityHandlerReturnsConflictWhenDisabled(t *testing.T) {
 		"session_intel_url": "https://intel.example.test"
 	}`))
 	rec := httptest.NewRecorder()
-	forkIdentityHandler(slog.Default()).ServeHTTP(rec, req)
+	forkIdentityHandler(slog.Default(), nil).ServeHTTP(rec, req)
 
 	assert.Equal(t, http.StatusConflict, rec.Code)
 }
@@ -107,7 +107,7 @@ func TestForkIdentityHandlerRejectsBadPayload(t *testing.T) {
 
 	req := httptest.NewRequest(http.MethodPost, "/internal/fork-identity", strings.NewReader(`{}`))
 	rec := httptest.NewRecorder()
-	forkIdentityHandler(slog.Default()).ServeHTTP(rec, req)
+	forkIdentityHandler(slog.Default(), nil).ServeHTTP(rec, req)
 
 	assert.Equal(t, http.StatusBadRequest, rec.Code)
 }
@@ -122,7 +122,7 @@ func TestForkIdentityHandlerRejectsAfterApplied(t *testing.T) {
 		"session_intel_url": "https://intel.example.test"
 	}`))
 	rec := httptest.NewRecorder()
-	forkIdentityHandler(slog.Default()).ServeHTTP(rec, req)
+	forkIdentityHandler(slog.Default(), nil).ServeHTTP(rec, req)
 
 	assert.Equal(t, http.StatusConflict, rec.Code)
 }
@@ -136,9 +136,10 @@ func TestForkIdentityHandlerWritesPayloadAndWaitsForAppliedMarker(t *testing.T)
 		"session_intel_url": "https://intel.example.test"
 	}`))
 	rec := httptest.NewRecorder()
+	applied := make(chan forkidentity.Payload, 1)
 	done := make(chan struct{})
 	go func() {
-		forkIdentityHandler(slog.Default()).ServeHTTP(rec, req)
+		forkIdentityHandler(slog.Default(), func(p forkidentity.Payload) { applied <- p }).ServeHTTP(rec, req)
 		close(done)
 	}()
 
@@ -157,6 +158,28 @@ func TestForkIdentityHandlerWritesPayloadAndWaitsForAppliedMarker(t *testing.T)
 		t.Fatal("handler did not return after applied marker")
 	}
 	assert.Equal(t, http.StatusNoContent, rec.Code)
+
+	// The callback is what lets this process retarget the sinks that label
+	// events with the instance identity.
+	select {
+	case p := <-applied:
+		assert.Equal(t, "browser-1", p.InstanceName())
+	default:
+		t.Fatal("applied callback did not run")
+	}
+}
+
+func TestForkIdentityHandlerSkipsCallbackWhenNotApplied(t *testing.T) {
+	useTempForkIdentityFiles(t)
+	t.Setenv(forkidentity.WaitEnv, "true")
+
+	req := httptest.NewRequest(http.MethodPost, "/internal/fork-identity", strings.NewReader(`{}`))
+	rec := httptest.NewRecorder()
+	called := false
+	forkIdentityHandler(slog.Default(), func(forkidentity.Payload) { called = true }).ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusBadRequest, rec.Code)
+	assert.False(t, called)
 }
 
 func useTempForkIdentityFiles(t *testing.T) {

server/cmd/api/main.go

@@ -11,6 +11,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"strings"
+	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -28,6 +29,7 @@ import (
 	"github.com/kernel/kernel-images/server/lib/chromedriverproxy"
 	"github.com/kernel/kernel-images/server/lib/devtoolsproxy"
 	"github.com/kernel/kernel-images/server/lib/events"
+	"github.com/kernel/kernel-images/server/lib/forkidentity"
 	"github.com/kernel/kernel-images/server/lib/logger"
 	"github.com/kernel/kernel-images/server/lib/metrics"
 	"github.com/kernel/kernel-images/server/lib/nekoclient"
@@ -122,12 +124,30 @@ func main() {
 		slogger.Error("sysmon: kmsg OOM monitor disabled", "err", err)
 	}
 
-	// Optional S2 storage sink.
-	var s2Writer *events.S2StorageWriter
-	if config.S2Basin != "" && config.S2AccessToken != "" && config.S2Stream != "" {
-		slogger.Info("S2 storage enabled", "basin", config.S2Basin, "stream", config.S2Stream)
-		s2Writer = events.NewS2StorageWriter(eventStream, config.S2Basin, config.S2AccessToken, config.S2Stream, events.S2Config{}, slogger)
-		if err := s2Writer.Start(ctx); err != nil {
+	forkIdentityWait, err := forkidentity.WaitEnabled()
+	if err != nil {
+		slogger.Error("invalid fork identity configuration", "err", err)
+		os.Exit(1)
+	}
+
+	// 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.
+	var s2Writer atomic.Pointer[events.S2StorageWriter]
+	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)
+		return w.Start(ctx)
+	}
+	if !forkIdentityWait {
+		if err := startS2Writer(config.S2Stream); err != nil {
 			slogger.Error("failed to start S2 storage writer", "err", err)
 			os.Exit(1)
 		}
@@ -138,6 +158,7 @@ func main() {
 	// off until turned on per-session via the telemetry API, so a VM that always
 	// has the relay injected does not export (and get dropped) by default.
 	var otlpExporter api.OTLPExporter
+	var otlpController *events.OTLPExportController
 	var otlpMetrics *events.OTLPMetrics
 	if config.OTLPEndpoint != "" {
 		headers := map[string]string{}
@@ -152,7 +173,7 @@ func main() {
 		otlpMetrics = &events.OTLPMetrics{}
 		otelDiag := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo - 1})
 		otel.SetLogger(logr.FromSlogHandler(events.NewDropCountingHandler(otelDiag, otlpMetrics)))
-		otlpExporter = events.NewOTLPExportController(eventStream, events.OTLPConfig{
+		otlpController = events.NewOTLPExportController(eventStream, events.OTLPConfig{
 			Endpoint:       config.OTLPEndpoint,
 			URLPath:        config.OTLPPath,
 			Insecure:       config.OTLPInsecure,
@@ -165,6 +186,20 @@ func main() {
 			ExportTimeout:  config.OTLPExportTimeout,
 			Metrics:        otlpMetrics,
 		}, slogger)
+		otlpExporter = otlpController
+	}
+
+	// A forked instance boots carrying the identity of the instance it came
+	// from, so the sinks that label events with it are retargeted here, once
+	// the guest has taken an identity of its own.
+	onForkIdentityApplied := func(payload forkidentity.Payload) {
+		env := forkidentity.Env(payload)
+		if otlpController != nil {
+			otlpController.SetIdentity(env["INSTANCE_NAME"], env["METRO_NAME"])
+		}
+		if err := startS2Writer(forkidentity.FirstNonEmpty(env["S2_STREAM"], config.S2Stream)); err != nil {
+			slogger.Error("failed to start S2 storage writer for fork identity", "err", err)
+		}
 	}
 
 	apiService, err := api.New(
@@ -191,7 +226,7 @@ func main() {
 	oapi.HandlerFromMux(strictHandler, r)
 
 	// Fork identity endpoints - not part of OpenAPI spec.
-	r.Post("/internal/fork-identity", forkIdentityHandler(slogger))
+	r.Post("/internal/fork-identity", forkIdentityHandler(slogger, onForkIdentityApplied))
 	r.Get("/internal/fork-identity/config", forkIdentityConfigHandler(slogger))
 
 	// endpoints to expose the spec
@@ -384,10 +419,10 @@ func main() {
 
 	// 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.
-	if s2Writer != nil {
+	if w := s2Writer.Load(); w != nil {
 		stopCtx, stopCancel := context.WithTimeout(context.Background(), 10*time.Second)
 		defer stopCancel()
-		if err := s2Writer.Stop(stopCtx); err != nil {
+		if err := w.Stop(stopCtx); err != nil {
 			slogger.Error("s2 storage writer stop failed", "err", err)
 		}
 	}

server/lib/events/otlpstorage.go

@@ -361,3 +361,19 @@ func (c *OTLPExportController) Running() bool {
 	defer c.mu.Unlock()
 	return c.writer != nil
 }
+
+// SetIdentity replaces the instance identity attached to exported events, for
+// instances that learn theirs after boot through the fork identity handoff.
+// Empty values are ignored. The OTLP Resource carrying the identity is built
+// when export starts, so this affects subsequent Starts — which covers the
+// fork case, where export stays off until a session turns it on.
+func (c *OTLPExportController) SetIdentity(instanceName, metro string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if instanceName != "" {
+		c.cfg.InstanceName = instanceName
+	}
+	if metro != "" {
+		c.cfg.Metro = metro
+	}
+}

server/lib/events/otlpstorage_test.go

@@ -239,3 +239,16 @@ func TestChunkBySize(t *testing.T) {
 		assert.Len(t, chunks[0], 1)
 	})
 }
+
+func TestOTLPExportControllerSetIdentity(t *testing.T) {
+	c := NewOTLPExportController(nil, OTLPConfig{InstanceName: "seed", Metro: "seed-metro"}, slog.Default())
+
+	c.SetIdentity("fork", "fork-metro")
+	assert.Equal(t, "fork", c.cfg.InstanceName)
+	assert.Equal(t, "fork-metro", c.cfg.Metro)
+
+	// A payload that omits a value must not blank out the identity in place.
+	c.SetIdentity("", "")
+	assert.Equal(t, "fork", c.cfg.InstanceName)
+	assert.Equal(t, "fork-metro", c.cfg.Metro)
+}