test(runtime): add persistence observer streaming benchmarks - #4077
test(runtime): add persistence observer streaming benchmarks#4077salignatmoandal wants to merge 3 commits into
Conversation
36cfbd2 to
084471c
Compare
|
👋 Some commits in this PR are not signed and verified by GitHub. Please sign your commits with a GPG or SSH key registered in your GitHub account, then force-push. Commits that are not verified: See GitHub's guide on signing commits for setup instructions. I've added |
084471c to
5cbf63a
Compare
Document the per-chunk AddMessage/UpdateMessage contract and establish in-memory vs SQLite baselines for streaming assistant persistence.
5cbf63a to
e40f6b6
Compare
|
Hi @aheritier, thanks for the heads-up. I've signed the commit with my SSH signing key and force-pushed (e40f6b6). GitHub now shows it as verified on my side. Happy to adjust anything else if needed. |
Sayt-0
left a comment
There was a problem hiding this comment.
Useful addition: the streaming persistence contract (1 AddMessage + N-1 UpdateMessage per turn) was not pinned by any focused test, and both characterization tests verify correctly against persistStreamingContent. Tests pass locally and follow the surrounding conventions (t.Parallel, t.Context, testify, :memory: + SetMaxOpenConns(1) pattern from pkg/session/store_memory_test.go).
Main issue: the in-memory benchmark result is a function of b.N (details inline), which defeats the stated goal of establishing a baseline. The SQLite benchmark is fine (reproduced ~2.9 ms/op, stable across benchtime values).
| Area | Status |
|---|---|
| Characterization tests | correct, verified against persistence_observer.go |
| SQLite benchmark | stable, no change needed |
| In-memory benchmark | numbers depend on b.N, needs fix (inline comment) |
| Conventions, vet | clean |
Non-blocking: the file name persistence_observer_bench_test.go also hosts two unit tests; moving them to persistence_observer_streaming_test.go would keep the bench file benchmark-only, matching pkg/tui/components/message/bench_test.go.
| for range b.N { | ||
| emitStreamingChunks(ctx, obs, sess, streamingBenchChunks) | ||
| finalizeStreamingMessage(ctx, obs, sess) | ||
| } |
There was a problem hiding this comment.
The shared session makes ns/op a function of b.N: each iteration leaves one extra message row in bench-session, and InMemorySessionStore.UpdateMessage scans every message of every session on each call, so later iterations pay O(iterations) per chunk.
Measured on this branch (darwin/arm64):
| benchtime | ns/op |
|---|---|
| 200x | 195,198 |
| 2000x | 624,997 |
| 8000x | 2,398,435 |
The PR's headline number (~2.0 ms/op) is therefore an artifact of the iteration count picked by the harness, not a baseline that future optimizations can be compared against. The "in-memory drift" caveat in the description understates this: the number is not noisy, it is unbounded.
Keeping the store bounded flattens the result (~175 µs/op at 200x, 2000x and 8000x, verified):
for i := range b.N {
sess := session.New(session.WithID(strconv.Itoa(i)), session.WithUserMessage("hi"))
if err := store.AddSession(ctx, sess); err != nil {
b.Fatal(err)
}
emitStreamingChunks(ctx, obs, sess, streamingBenchChunks)
finalizeStreamingMessage(ctx, obs, sess)
if err := store.DeleteSession(ctx, sess.ID); err != nil {
b.Fatal(err)
}
}Requires the strconv import and the store returned by setupPersistenceObserverBench. Note that a fresh session per iteration alone is not enough: UpdateMessage ranges over all sessions, so the store must not accumulate them, hence the DeleteSession. If per-iteration setup cost is a concern, it is ~2 allocations against ~2,500 per iteration, so it does not move the numbers.
There was a problem hiding this comment.
Good catch — confirmed in InMemorySessionStore.UpdateMessage: it Ranges every session and scans Messages, so cost is O(sessions × messages). With a shared session the store grew unbounded with b.N, so ns/op was not a baseline.
Applied your suggested harness in a0ddcd4: each iteration does AddSession → work → DeleteSession (fresh session alone is not enough because of the session Range). Also simplified the helper to (obs, store).
Re-measured on darwin/arm64 (Apple M3 Pro):
| benchtime | before (unbounded) | after (bounded) |
|---|---|---|
| 200x | ~203 µs | 165 µs |
| 2000x | ~673 µs | 160 µs |
| 8000x | ~2575 µs | 160 µs |
ns/op is now flat (~160 µs), matching your ~175 µs result. Allocs stay ~2522/op.
Happy to follow up on the slog io.Discard nit and splitting unit tests into persistence_observer_streaming_test.go if you want those in this PR too.
| return s.InMemorySessionStore.UpdateMessage(ctx, messageID, msg) | ||
| } | ||
|
|
||
| func setupPersistenceObserverBench(tb testing.TB) (*PersistenceObserver, *session.Session, *session.InMemorySessionStore) { |
There was a problem hiding this comment.
Nit: the third return value is discarded by the only caller. With the bounded-store fix suggested in the other comment, the session return becomes the unused one instead. Returning (obs, store) and letting benchmarks create their own sessions would keep the helper minimal.
There was a problem hiding this comment.
Done in a0ddcd4 — setupPersistenceObserverBench now returns (obs, store); the benchmark creates/deletes its own sessions per iteration.
| b.Cleanup(func() { _ = db.Close() }) | ||
| db.SetMaxOpenConns(1) | ||
|
|
||
| store, err := session.NewSQLiteSessionStoreFromDB(b.Context(), db) |
There was a problem hiding this comment.
Nit: migration INFO logs are emitted on every harness re-invocation of the benchmark function, interleaving with the -bench output. Not counted in ns/op since setup runs before b.ResetTimer, but redirecting the default slog handler to io.Discard for the benchmark would keep the output clean for tools that parse it.
Reset the session each iteration so UpdateMessage's O(sessions×messages) scan cannot make ns/op a function of b.N.
|
Good catch — confirmed in Fixed in a0ddcd4 with your suggested harness: each iteration does Re-measured on darwin/arm64 (Apple M3 Pro):
Happy to follow up on the slog |
Discard the default slog handler during SQLite store setup so harness re-invocations do not interleave migration INFO with -bench output.
Summary
Adds characterization tests and benchmarks for
PersistenceObserverstreaming persistence — the path that mirrors assistant token deltas (AgentChoice/AgentChoiceReasoning) into a single growing message row in the session store.This establishes a baseline before any future optimization (e.g. debounced flushes) and documents the per-chunk store write contract that the observer currently implements.
Context
During a streaming assistant turn, the runtime emits one
AgentChoiceEventper delta.PersistenceObserver.persistStreamingContentkeeps a single in-flight row and:AddMessage) on the first chunkUpdateMessage) on every subsequent chunkMessageAddedEventto finalise the row with the canonical payloadThis behaviour is easy to regress when refactoring persistence or the store layer, but was not previously covered by a focused unit test or benchmark.
Changes
New file:
pkg/runtime/persistence_observer_bench_test.goTests
TestPersistenceObserver_UpdateCountPerChunkAddMessage+(N-1)UpdateMessagecalls for N streaming chunks, plus one finalUpdateMessageonMessageAddedTestPersistenceObserver_StreamingContentAccumulates"hel"+"lo"→"hello")Both tests use a
countingStorewrapper aroundInMemorySessionStoreto assert store call counts without mocking the observer.Benchmarks
Each iteration simulates a long assistant turn: 500
AgentChoicedeltas + 1MessageAddedfinalisation (streamingBenchChunks = 500).BenchmarkPersistenceObserver_StreamingChunksInMemorySessionStoreBenchmarkPersistenceObserver_StreamingChunks_SQLite:memory:json.Marshal+UPDATE session_itemsper chunkSample results (darwin/arm64, Apple M3 Pro)
Per chunk (500 chunks/iter): ~4 µs in-memory, ~6 µs SQLite, ~5 vs ~17 allocs.
SQLite is ~1.5× slower and ~3× more alloc-heavy — expected given JSON marshal + SQL per update.
Known benchmark caveats
b.Niterations, soUpdateMessageeventually scans an ever-growing message list (O(messages) per update). Later iterations are slower than the first. SQLite stays flat because updates are keyed bymessage_id.NewSQLiteSessionStoreFromDB, which logs migration info beforeb.ResetTimer()— noisy stdout but not included inns/op.These caveats are acceptable for a baseline but worth keeping in mind when comparing future numbers.
Why now?
Streaming persistence is on the hot path for every assistant response when a session store is configured. Having explicit tests + benchmarks makes it safer to:
persistStreamingContent(avoidstrings.Builder.String()copies on every delta)Test plan
go test ./pkg/runtime -run TestPersistenceObserver_UpdateCountPerChunk -vgo test ./pkg/runtime -run TestPersistenceObserver_StreamingContentAccumulates -vgo test ./pkg/runtime -run=^$ -bench=BenchmarkPersistenceObserver -benchmem -count=1task test(full suite)task lint