Unified Storage: Init at startup, fix traces, and speed up indexing

[marchugon]

Changes:

• Initializes unified storage when the ResourceServer is created, instead of doing it inside the first gRPC call it receives. This was causing context cancelled errors since the index takes too long to build within the context of the gRPC call.
• Fixes trace propagation by passing span contexts down.
• Improves index build speed using finer grain locking when writing the index cache. We were locking the whole BuildIndex function, which was slowing things down when building many namespaces with high concurrency.

Notes for reviewer:

• Does anything else depend on initializing US lazily? What could this break?

[friendlyreviewer-staging]

Hi there 👋

🌥️ Tech	🌥️ Feat
2 high, 8 medium, 4 low	2 issues, 2 mitigated

The MR correctly moves unified storage initialization to ResourceServer construction, which resolves context-cancelled errors in gRPC calls, and applies the intended trace-propagation fix in several spots. However, the finer-grain locking change introduces serious concurrency defects: TotalDocs reads the cache map without a lock, and concurrent BuildIndex calls can run for the same key, risking file-index corruption. Trace propagation remains incomplete in WriteEvent, ReadResource, ListIterator, and BuildIndex. Startup error handling can hard-fail the whole server on a single namespace build failure, and the test coverage has reliability gaps (fixed-sleep waits, unguarded goroutines, missing timeouts, and a postgres skip). These issues must be addressed before merge.

---

These might need a close look

• 🔴 pkg/storage/unified/search/bleve.go (L145)
  TotalDocs() iterates b.cache (a map) without holding cacheMu, while BuildIndex writes b.cache[key] under finer-grain locking. This is a data race that can occur during concurrent builds and Prometheus scrapes. Use cacheMu.RLock() for the iteration.
• 🔴 pkg/storage/unified/search/bleve.go (L99)
  Removing the function-wide cacheMu.Lock() allows two BuildIndex calls for the same key to run concurrently. This can corrupt file-based indexes (same directory) and cause duplicate work / double tenant counters. Use per-key single-flight (e.g., keyed mutex or singleflight) to serialize builds per key while allowing parallelism across different keys.

Worth checking

• 🟡 pkg/storage/unified/search/bleve.go (L97)
  BuildIndex still uses _, span := b.tracer.Start(...) and does not thread the derived context into the indexing work (builder closure takes no ctx). The BuildIndex span remains a dangling leaf, inconsistent with the trace-propagation fix applied elsewhere. Use ctx, span := and pass ctx into the builder.
• 🟡 pkg/storage/unified/resource/search.go (L187)
  totalBatchesIndexed is a plain int incremented inside concurrent goroutines and read afterwards, causing a data race. Use an atomic.Int64 or collect per-worker counts.
• 🟡 pkg/storage/unified/resource/server.go (L258)
  When s.Init(ctx) fails, NewResourceServer returns nil/err without calling s.cancel() or tearing down partially initialized state. Consider calling s.cancel() on the error path. Also note that a single namespace build failure now hard-fails server startup; consider tolerating partial failures.
• 🟡 pkg/storage/unified/sql/backend.go (L158)
  Trace propagation is still dropped in WriteEvent, ReadResource, and ListIterator: they use _, span := and pass the original ctx to child DB operations, so their spans are leaves. Convert to ctx, span := and pass the derived ctx down to nest child spans correctly.
• 🟡 pkg/storage/unified/sql/backend.go (L112)
  sql backend methods (IsHealthy, GetResourceStats, create/update/delete, ReadResource, ListIterator, poller) rely on external Init to populate b.db and have no nil check or self-guard. Since per-call init guards were removed from server.go, callers using NewBackend without Init will panic on nil b.db. Document/guard the Init contract or add defensive nil checks.
• 🟡 pkg/server/module_server_test.go (L56)
  The fixed 500ms sleep before polling /metrics is flaky and races the new synchronous startup init. If init takes longer, /metrics can be polled before readiness and Shutdown may cancel an in-flight init causing spurious failure. Replace with a bounded poll/retry (e.g., require.Eventually) that waits on server readiness.
• 🟡 pkg/server/module_server_test.go (L52)
  err.Error() in the goroutine panics if ms.Run() returns nil, and the exact-string comparison is brittle. Use errors.Is(err, context.Canceled) and guard for nil. Also the goroutine is never joined, racing t's use after the test returns; capture the error on a buffered channel and assert in the main goroutine.
• 🟡 pkg/server/module_server_test.go (L65)
  ms.Shutdown has no timeout and doesn't assert in-flight init cancellation; if shutdown hangs, the test hangs CI. Wrap shutdown in a timeout context so hangs surface as failures.

Small things (take or leave)

• 🔵 pkg/storage/unified/sql/backend.go (L580)
  The poller span is not ended on error paths (listLatestRVs/poll errors take continue and skip span.End()), leaking un-ended spans. Use defer span.End() immediately after starting the span.
• 🔵 pkg/server/module_server_test.go (L58)
  The http.Client has no Timeout and /metrics GET uses a background context, so it can hang indefinitely. Use context.WithTimeout and/or client.Timeout.
• 🔵 pkg/server/module_server_test.go (L63)
  The test only asserts /metrics HTTP 200, which is served by the independent InstrumentationServer and does not verify US initialization. Add an assertion that the storage server actually became ready (e.g., wait for its gRPC address or a readiness check).
• 🔵 pkg/server/module_server_test.go (L36)
  The added postgres skip reduces CI coverage for the startup-init lifecycle. The TODO to fix the postgres test should be tracked so the coverage gap is not lost.

---

✅ Feature-level checklist

Not there yet

• ❌ Trace propagation is fixed by passing span contexts down through all backend operations.
  Multiple spans still discard the derived context (WriteEvent, ReadResource, ListIterator in backend.go; BuildIndex in bleve.go), so the fix is incomplete.
• ❌ Index build speed improved using finer grain locking when writing the index cache.
  The finer-grain locking introduces data races (TotalDocs, totalBatchesIndexed) and permits concurrent BuildIndex for the same key, which can corrupt file-based indexes.

Partially covered

• ⚠️ Unified Storage is initialized at ResourceServer creation (startup) instead of lazily on the first gRPC call.
  Implemented, but startup can hard-fail on a single namespace build error, and sql backend methods now require explicit Init prior to use. Integration tests and startup path call Init explicitly, but exported Backend usage without Init can panic.
• ⚠️ No existing dependency on lazy US initialization is broken.
  Integration tests and startup path call Init explicitly, so normal flows work. However, the sql Backend interface now has an implicit Init contract; callers using NewBackend without Init will hit nil-pointer panics.

Review time: 9m 25s