| 1 | // Command sessioncapacity exercises the production session service at a |
| 2 | // configurable scale. It intentionally keeps acceptance sizes in flags rather |
| 3 | // than production constants: 100k events and GiB-sized data sets are release |
| 4 | // evidence, not product limits. |
| 5 | package main |
| 6 | |
| 7 | import ( |
| 8 | "context" |
| 9 | "encoding/json" |
| 10 | "errors" |
| 11 | "flag" |
| 12 | "fmt" |
| 13 | "io" |
| 14 | "os" |
| 15 | "os/signal" |
| 16 | "path/filepath" |
| 17 | "runtime" |
| 18 | "slices" |
| 19 | "strings" |
| 20 | "sync" |
| 21 | "time" |
| 22 | |
| 23 | "reasonix/internal/provider" |
| 24 | "reasonix/internal/session" |
| 25 | "reasonix/internal/sessioncontent" |
| 26 | ) |
| 27 | |
| 28 | const mib = int64(1 << 20) |
| 29 | |
| 30 | type config struct { |
| 31 | Root string |
| 32 | SessionID string |
| 33 | HistoryMessages int |
| 34 | HistoryBytes int64 |
| 35 | AttachmentBytes int64 |
| 36 | AttachmentChunkBytes int64 |
| 37 | WorksetBytes int64 |
| 38 | FlushEvery int |
| 39 | PageSamples int |
| 40 | } |
| 41 | |
| 42 | type report struct { |
| 43 | SessionID string `json:"sessionId"` |
| 44 | HistoryMessages int `json:"historyMessages"` |
| 45 | HistoryLogicalBytes int64 `json:"historyLogicalBytes"` |
| 46 | AttachmentObjects int `json:"attachmentObjects"` |
| 47 | AttachmentLogicalBytes int64 `json:"attachmentLogicalBytes"` |
| 48 | EventSequence uint64 `json:"eventSequence"` |
| 49 | DurableSequence uint64 `json:"durableSequence"` |
| 50 | ModelWorksetBytes int64 `json:"modelWorksetBytes"` |
| 51 | WriteDurationMS int64 `json:"writeDurationMs"` |
| 52 | ColdOpenDurationMS int64 `json:"coldOpenDurationMs"` |
| 53 | HistoryIndexBuildMS int64 `json:"historyIndexBuildDurationMs"` |
| 54 | IndexedPageP95MS int64 `json:"indexedPageP95Ms"` |
| 55 | SessionDiskBytes int64 `json:"sessionDiskBytes"` |
| 56 | ContentDiskBytes int64 `json:"contentDiskBytes"` |
| 57 | QueryCacheDiskBytes int64 `json:"queryCacheDiskBytes"` |
| 58 | BaselineHeapAllocBytes uint64 `json:"baselineHeapAllocBytes"` |
| 59 | PeakHeapAllocBytes uint64 `json:"peakHeapAllocBytes"` |
| 60 | BaselineRuntimeSysBytes uint64 `json:"baselineRuntimeSysBytes"` |
| 61 | PeakRuntimeSysBytes uint64 `json:"peakRuntimeSysBytes"` |
| 62 | ProcessPeakRSSBytes uint64 `json:"processPeakRssBytes,omitempty"` |
| 63 | WritePeakHeapBytes uint64 `json:"writePeakHeapBytes"` |
| 64 | ColdOpenPeakHeapBytes uint64 `json:"coldOpenPeakHeapBytes"` |
| 65 | HistoryPeakHeapBytes uint64 `json:"historyPeakHeapBytes"` |
| 66 | } |
| 67 | |
| 68 | func main() { |
| 69 | var cfg config |
| 70 | flag.StringVar(&cfg.Root, "root", "", "sessions-v4 root (a temporary root is used when empty)") |
| 71 | flag.StringVar(&cfg.SessionID, "session", "capacity-acceptance", "session identity") |
| 72 | flag.IntVar(&cfg.HistoryMessages, "history-messages", 2_000, "number of durable history messages") |
| 73 | flag.Int64Var(&cfg.HistoryBytes, "history-bytes", 32*mib, "logical bytes spread across history messages") |
| 74 | flag.Int64Var(&cfg.AttachmentBytes, "attachment-bytes", 32*mib, "bytes streamed through the shared content store") |
| 75 | flag.Int64Var(&cfg.AttachmentChunkBytes, "attachment-chunk-bytes", 4*mib, "maximum bytes per attachment object") |
| 76 | flag.Int64Var(&cfg.WorksetBytes, "workset-bytes", 4*mib, "provider model workset retained across cold open") |
| 77 | flag.IntVar(&cfg.FlushEvery, "flush-every", 128, "flush interval in history messages") |
| 78 | flag.IntVar(&cfg.PageSamples, "page-samples", 7, "indexed newest-page timing samples") |
| 79 | flag.Parse() |
| 80 | |
| 81 | rootWasTemporary := cfg.Root == "" |
| 82 | if rootWasTemporary { |
| 83 | temp, err := os.MkdirTemp("", "reasonix-session-capacity-*") |
| 84 | if err != nil { |
| 85 | fmt.Fprintln(os.Stderr, err) |
| 86 | os.Exit(1) |
| 87 | } |
| 88 | cfg.Root = filepath.Join(temp, "sessions-v4") |
| 89 | defer os.RemoveAll(temp) |
| 90 | } |
| 91 | ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt) |
| 92 | defer cancel() |
| 93 | result, err := run(ctx, cfg) |
| 94 | if err != nil { |
| 95 | fmt.Fprintln(os.Stderr, err) |
| 96 | os.Exit(1) |
| 97 | } |
| 98 | encoded, err := json.MarshalIndent(result, "", " ") |
| 99 | if err != nil { |
| 100 | fmt.Fprintln(os.Stderr, err) |
| 101 | os.Exit(1) |
| 102 | } |
| 103 | fmt.Println(string(encoded)) |
| 104 | if !rootWasTemporary { |
| 105 | fmt.Fprintf(os.Stderr, "capacity data retained at %s\n", cfg.Root) |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | // releaseCapacityService drops the query cache together with the writer lease |
| 110 | // and recovery store. Closing the query alone leaves recovery-v1.bolt open, and |
| 111 | // Windows refuses to remove a directory that still holds it. |
| 112 | func releaseCapacityService(service *session.Service) { |
| 113 | service.Query().Close() |
| 114 | _ = service.CloseAll(context.Background()) |
| 115 | } |
| 116 | |
| 117 | func run(ctx context.Context, cfg config) (result report, err error) { |
| 118 | if err := validateConfig(cfg); err != nil { |
| 119 | return report{}, err |
| 120 | } |
| 121 | if err := os.MkdirAll(cfg.Root, 0o700); err != nil { |
| 122 | return report{}, err |
| 123 | } |
| 124 | runtime.GC() |
| 125 | baseline := readMemory() |
| 126 | peaks := newMemoryPeaks(baseline) |
| 127 | stopMemory := make(chan struct{}) |
| 128 | memoryDone := make(chan struct{}) |
| 129 | go sampleMemory(stopMemory, memoryDone, peaks) |
| 130 | defer func() { |
| 131 | close(stopMemory) |
| 132 | <-memoryDone |
| 133 | peaks.observe(readMemory()) |
| 134 | peak := peaks.totalPeak() |
| 135 | result.BaselineHeapAllocBytes = baseline.heap |
| 136 | result.PeakHeapAllocBytes = peak.heap |
| 137 | result.BaselineRuntimeSysBytes = baseline.sys |
| 138 | result.PeakRuntimeSysBytes = peak.sys |
| 139 | result.ProcessPeakRSSBytes = peak.rss |
| 140 | }() |
| 141 | |
| 142 | service, err := session.NewService("capacity", session.NewFilesystemPersistence(cfg.Root)) |
| 143 | if err != nil { |
| 144 | return report{}, err |
| 145 | } |
| 146 | runtimeSession, err := service.Create(ctx, session.CreateOptions{SessionID: cfg.SessionID}) |
| 147 | if err != nil { |
| 148 | return report{}, err |
| 149 | } |
| 150 | ref := runtimeSession.Ref() |
| 151 | defer service.Query().Close() |
| 152 | |
| 153 | writeStarted := time.Now() |
| 154 | written, err := appendHistory(ctx, runtimeSession.Session(), cfg) |
| 155 | if err != nil { |
| 156 | _ = service.Close(context.Background(), ref) |
| 157 | return report{}, err |
| 158 | } |
| 159 | attachments, err := writeAttachments(ctx, filepath.Join(cfg.Root, ".content-v1"), cfg.AttachmentBytes, cfg.AttachmentChunkBytes) |
| 160 | if err != nil { |
| 161 | _ = service.Close(context.Background(), ref) |
| 162 | return report{}, err |
| 163 | } |
| 164 | workset := provider.Message{ID: "capacity-workset", Role: provider.RoleUser, Content: deterministicString(cfg.WorksetBytes, uint64(cfg.HistoryMessages)+1)} |
| 165 | worksetPayload, err := json.Marshal(map[string]any{"messages": []provider.Message{workset}, "reason": "capacity acceptance workset"}) |
| 166 | if err != nil { |
| 167 | return report{}, err |
| 168 | } |
| 169 | if _, err := runtimeSession.Session().AppendBatch(ctx, "capacity-workset", []session.Event{{Kind: "model/context-replace", Payload: worksetPayload}}); err != nil { |
| 170 | return report{}, err |
| 171 | } |
| 172 | receipt, err := runtimeSession.Session().Flush(ctx) |
| 173 | if err != nil { |
| 174 | return report{}, err |
| 175 | } |
| 176 | result.WriteDurationMS = time.Since(writeStarted).Milliseconds() |
| 177 | result.SessionID = cfg.SessionID |
| 178 | result.HistoryMessages = cfg.HistoryMessages |
| 179 | result.HistoryLogicalBytes = written |
| 180 | result.AttachmentObjects = attachments |
| 181 | result.AttachmentLogicalBytes = cfg.AttachmentBytes |
| 182 | result.EventSequence = runtimeSession.Session().EventSequence() |
| 183 | result.DurableSequence = receipt.DurableSequence |
| 184 | if err := service.Close(ctx, ref); err != nil { |
| 185 | return report{}, err |
| 186 | } |
| 187 | result.WritePeakHeapBytes = peaks.nextStage().heap |
| 188 | |
| 189 | second, err := session.NewService("capacity", session.NewFilesystemPersistence(cfg.Root)) |
| 190 | if err != nil { |
| 191 | return report{}, err |
| 192 | } |
| 193 | defer releaseCapacityService(second) |
| 194 | openStarted := time.Now() |
| 195 | binding, err := second.Open(ctx, ref) |
| 196 | if err != nil { |
| 197 | return report{}, err |
| 198 | } |
| 199 | result.ColdOpenDurationMS = time.Since(openStarted).Milliseconds() |
| 200 | model := binding.Runtime().Session().DeriveMessages() |
| 201 | for _, message := range model { |
| 202 | result.ModelWorksetBytes += int64(len(message.Content)) |
| 203 | } |
| 204 | if result.ModelWorksetBytes != cfg.WorksetBytes { |
| 205 | _ = binding.Release(context.Background()) |
| 206 | return report{}, fmt.Errorf("capacity: cold model workset is %d bytes, expected %d", result.ModelWorksetBytes, cfg.WorksetBytes) |
| 207 | } |
| 208 | result.ColdOpenPeakHeapBytes = peaks.nextStage().heap |
| 209 | |
| 210 | indexStarted := time.Now() |
| 211 | page, err := waitForHistoryPage(ctx, second.Query(), ref) |
| 212 | if err != nil { |
| 213 | _ = binding.Release(context.Background()) |
| 214 | return report{}, err |
| 215 | } |
| 216 | result.HistoryIndexBuildMS = time.Since(indexStarted).Milliseconds() |
| 217 | if len(page.Messages) == 0 && cfg.HistoryMessages > 0 { |
| 218 | _ = binding.Release(context.Background()) |
| 219 | return report{}, errors.New("capacity: rebuilt history index returned no messages") |
| 220 | } |
| 221 | durations := make([]time.Duration, 0, cfg.PageSamples) |
| 222 | for range cfg.PageSamples { |
| 223 | started := time.Now() |
| 224 | if _, err := second.Query().HistoryPage(ctx, ref, "", 100); err != nil { |
| 225 | _ = binding.Release(context.Background()) |
| 226 | return report{}, err |
| 227 | } |
| 228 | durations = append(durations, time.Since(started)) |
| 229 | } |
| 230 | result.IndexedPageP95MS = percentile95(durations).Milliseconds() |
| 231 | if err := binding.Release(ctx); err != nil { |
| 232 | return report{}, err |
| 233 | } |
| 234 | result.HistoryPeakHeapBytes = peaks.nextStage().heap |
| 235 | |
| 236 | result.SessionDiskBytes, err = treeBytes(filepath.Join(cfg.Root, cfg.SessionID)) |
| 237 | if err != nil { |
| 238 | return report{}, err |
| 239 | } |
| 240 | result.ContentDiskBytes, err = treeBytes(filepath.Join(cfg.Root, ".content-v1")) |
| 241 | if err != nil { |
| 242 | return report{}, err |
| 243 | } |
| 244 | result.QueryCacheDiskBytes, err = treeBytes(filepath.Join(cfg.Root, ".query-cache")) |
| 245 | return result, err |
| 246 | } |
| 247 | |
| 248 | func waitForHistoryPage(ctx context.Context, query *session.Query, ref session.SessionRef) (session.MessageHistoryPage, error) { |
| 249 | for { |
| 250 | page, err := query.HistoryPage(ctx, ref, "", 100) |
| 251 | if err != nil || page.Status == "ready" { |
| 252 | return page, err |
| 253 | } |
| 254 | if page.Status != "preparing" { |
| 255 | return session.MessageHistoryPage{}, fmt.Errorf("capacity: history locator status %q", page.Status) |
| 256 | } |
| 257 | select { |
| 258 | case <-ctx.Done(): |
| 259 | return session.MessageHistoryPage{}, ctx.Err() |
| 260 | case <-time.After(5 * time.Millisecond): |
| 261 | } |
| 262 | } |
| 263 | } |
| 264 | |
| 265 | func validateConfig(cfg config) error { |
| 266 | if strings.TrimSpace(cfg.Root) == "" || strings.TrimSpace(cfg.SessionID) == "" { |
| 267 | return errors.New("capacity: root and session are required") |
| 268 | } |
| 269 | if cfg.HistoryMessages < 1 || cfg.HistoryBytes < 0 || cfg.AttachmentBytes < 0 || cfg.AttachmentChunkBytes < 1 || cfg.WorksetBytes < 0 { |
| 270 | return errors.New("capacity: sizes must be non-negative and history/chunk counts must be positive") |
| 271 | } |
| 272 | if cfg.FlushEvery < 1 || cfg.PageSamples < 1 { |
| 273 | return errors.New("capacity: flush interval and page samples must be positive") |
| 274 | } |
| 275 | return nil |
| 276 | } |
| 277 | |
| 278 | func appendHistory(ctx context.Context, target *session.Session, cfg config) (int64, error) { |
| 279 | emptyContext := json.RawMessage(`{"messages":[],"reason":"capacity bounded workset"}`) |
| 280 | var written int64 |
| 281 | for i := range cfg.HistoryMessages { |
| 282 | remainingMessages := int64(cfg.HistoryMessages - i) |
| 283 | contentBytes := (cfg.HistoryBytes - written + remainingMessages - 1) / remainingMessages |
| 284 | content := deterministicString(contentBytes, uint64(i)+1) |
| 285 | message := provider.Message{ID: fmt.Sprintf("capacity-%08d", i), Role: provider.RoleUser, Content: content} |
| 286 | payload, err := json.Marshal(map[string]any{"message": message}) |
| 287 | if err != nil { |
| 288 | return written, err |
| 289 | } |
| 290 | _, err = target.AppendBatch(ctx, fmt.Sprintf("capacity-%08d", i), []session.Event{ |
| 291 | {Kind: "message/complete", Payload: payload}, |
| 292 | {Kind: "model/context-replace", Payload: emptyContext}, |
| 293 | }) |
| 294 | if err != nil { |
| 295 | return written, err |
| 296 | } |
| 297 | written += contentBytes |
| 298 | if (i+1)%cfg.FlushEvery == 0 { |
| 299 | if _, err := target.Flush(ctx); err != nil { |
| 300 | return written, err |
| 301 | } |
| 302 | } |
| 303 | } |
| 304 | return written, nil |
| 305 | } |
| 306 | |
| 307 | func writeAttachments(ctx context.Context, root string, total, chunk int64) (int, error) { |
| 308 | store := sessioncontent.New(root) |
| 309 | objects := 0 |
| 310 | for offset := int64(0); offset < total; { |
| 311 | size := min(chunk, total-offset) |
| 312 | reader := io.LimitReader(&deterministicReader{state: uint64(objects) + 0x9e3779b97f4a7c15}, size) |
| 313 | ref, err := store.Put(ctx, reader, sessioncontent.Metadata{MediaType: "application/octet-stream", Name: fmt.Sprintf("capacity-%06d.bin", objects)}) |
| 314 | if err != nil { |
| 315 | return objects, err |
| 316 | } |
| 317 | if ref.Bytes != size { |
| 318 | return objects, fmt.Errorf("capacity: attachment %d stored %d bytes, expected %d", objects, ref.Bytes, size) |
| 319 | } |
| 320 | offset += size |
| 321 | objects++ |
| 322 | } |
| 323 | return objects, nil |
| 324 | } |
| 325 | |
| 326 | func deterministicString(size int64, seed uint64) string { |
| 327 | if size <= 0 { |
| 328 | return "" |
| 329 | } |
| 330 | buf := make([]byte, int(size)) |
| 331 | r := deterministicReader{state: seed + 0x9e3779b97f4a7c15} |
| 332 | _, _ = io.ReadFull(&r, buf) |
| 333 | return string(buf) |
| 334 | } |
| 335 | |
| 336 | type deterministicReader struct{ state uint64 } |
| 337 | |
| 338 | func (r *deterministicReader) Read(p []byte) (int, error) { |
| 339 | for i := range p { |
| 340 | r.state ^= r.state << 13 |
| 341 | r.state ^= r.state >> 7 |
| 342 | r.state ^= r.state << 17 |
| 343 | p[i] = byte(33 + r.state%90) |
| 344 | } |
| 345 | return len(p), nil |
| 346 | } |
| 347 | |
| 348 | type memorySample struct{ heap, sys, rss uint64 } |
| 349 | |
| 350 | type memoryPeaks struct { |
| 351 | mu sync.Mutex |
| 352 | stage memorySample |
| 353 | total memorySample |
| 354 | } |
| 355 | |
| 356 | func newMemoryPeaks(initial memorySample) *memoryPeaks { |
| 357 | return &memoryPeaks{stage: initial, total: initial} |
| 358 | } |
| 359 | |
| 360 | func (p *memoryPeaks) observe(sample memorySample) { |
| 361 | p.mu.Lock() |
| 362 | p.stage.max(sample) |
| 363 | p.total.max(sample) |
| 364 | p.mu.Unlock() |
| 365 | } |
| 366 | |
| 367 | func (p *memoryPeaks) nextStage() memorySample { |
| 368 | latest := readMemory() |
| 369 | p.mu.Lock() |
| 370 | p.stage.max(latest) |
| 371 | p.total.max(latest) |
| 372 | finished := p.stage |
| 373 | p.stage = latest |
| 374 | p.mu.Unlock() |
| 375 | return finished |
| 376 | } |
| 377 | |
| 378 | func (p *memoryPeaks) totalPeak() memorySample { |
| 379 | p.mu.Lock() |
| 380 | defer p.mu.Unlock() |
| 381 | return p.total |
| 382 | } |
| 383 | |
| 384 | func readMemory() memorySample { |
| 385 | var stats runtime.MemStats |
| 386 | runtime.ReadMemStats(&stats) |
| 387 | return memorySample{heap: stats.HeapAlloc, sys: stats.Sys, rss: processPeakRSSBytes()} |
| 388 | } |
| 389 | |
| 390 | func (m *memorySample) max(other memorySample) { |
| 391 | m.heap = max(m.heap, other.heap) |
| 392 | m.sys = max(m.sys, other.sys) |
| 393 | m.rss = max(m.rss, other.rss) |
| 394 | } |
| 395 | |
| 396 | func sampleMemory(stop <-chan struct{}, done chan<- struct{}, peaks *memoryPeaks) { |
| 397 | defer close(done) |
| 398 | ticker := time.NewTicker(20 * time.Millisecond) |
| 399 | defer ticker.Stop() |
| 400 | for { |
| 401 | select { |
| 402 | case <-stop: |
| 403 | return |
| 404 | case <-ticker.C: |
| 405 | peaks.observe(readMemory()) |
| 406 | } |
| 407 | } |
| 408 | } |
| 409 | |
| 410 | func percentile95(values []time.Duration) time.Duration { |
| 411 | if len(values) == 0 { |
| 412 | return 0 |
| 413 | } |
| 414 | copyOfValues := append([]time.Duration(nil), values...) |
| 415 | slices.Sort(copyOfValues) |
| 416 | index := (95*len(copyOfValues) + 99) / 100 |
| 417 | return copyOfValues[index-1] |
| 418 | } |
| 419 | |
| 420 | func treeBytes(root string) (int64, error) { |
| 421 | var total int64 |
| 422 | err := filepath.WalkDir(root, func(path string, entry os.DirEntry, err error) error { |
| 423 | if err != nil { |
| 424 | if os.IsNotExist(err) { |
| 425 | return nil |
| 426 | } |
| 427 | return err |
| 428 | } |
| 429 | if entry.IsDir() { |
| 430 | return nil |
| 431 | } |
| 432 | info, err := entry.Info() |
| 433 | if err != nil { |
| 434 | return err |
| 435 | } |
| 436 | total += info.Size() |
| 437 | return nil |
| 438 | }) |
| 439 | return total, err |
| 440 | } |
| 441 |