UnrollSetMany now respects gRPC size limits, too

[kans]

Summary by CodeRabbit

• Refactor

  • Session storage batching now respects both per-request key limits and cumulative payload size, with lazy item marshaling and improved error propagation during batched operations.

• Tests

  • Expanded coverage for size-aware batching, iterator errors, and batched set failures.

• Chores

  • Session limits extracted and exposed as shared configurable constants for key-count and total-payload size.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 23.08% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately reflects the main change: adding size limit awareness to UnrollSetMany function.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch kans/size-chenked-set-many

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 0bf380e and 132b3df.

📒 Files selected for processing (5)

• pkg/dotc1z/session_store.go (2 hunks)
• pkg/dotc1z/session_store_test.go (13 hunks)
• pkg/session/session.go (2 hunks)
• pkg/session/session_test.go (5 hunks)
• pkg/types/sessions/sessions.go (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (2)

• pkg/dotc1z/session_store.go
• pkg/dotc1z/session_store_test.go

🧰 Additional context used

🧠 Learnings (2)

📚 Learning: 2025-11-26T21:36:04.901Z

Learnt from: kans
Repo: ConductorOne/baton-sdk PR: 563
File: pkg/dotc1z/session_store.go:287-294
Timestamp: 2025-11-26T21:36:04.901Z
Learning: In pkg/dotc1z/session_store.go, the GetMany function intentionally calculates item size using only len(value) rather than len(key)+len(value)+overhead. This differs from getAllChunk's calculation, but the difference is acceptable because the 4163584 byte message size limit includes sufficient padding to account for key overhead without explicit calculation.

Applied to files:

• pkg/types/sessions/sessions.go
• pkg/session/session_test.go
• pkg/session/session.go

📚 Learning: 2025-10-24T20:32:19.985Z

Learnt from: kans
Repo: ConductorOne/baton-sdk PR: 531
File: pkg/session/json_session.go:38-42
Timestamp: 2025-10-24T20:32:19.985Z
Learning: In the baton-sdk repository, SetManyJSON and similar batch operations in pkg/session/json_session.go should not add early-return guards for empty inputs. The design preference is to let server-side validation errors (such as min_pairs:1) propagate rather than silently succeeding on the client side.

Applied to files:

• pkg/session/session_test.go

🧬 Code graph analysis (2)

pkg/session/session_test.go (2)

pkg/session/session.go (2)

• SizedItem (68-72)
• UnrollSetMany (77-111)

pkg/types/sessions/sessions.go (3)

• SessionStoreOption (26-26)
• MaxSessionStoreSizeLimit (12-12)
• MaxKeysPerRequest (7-7)

pkg/session/session.go (1)

pkg/types/sessions/sessions.go (2)

• MaxKeysPerRequest (7-7)
• MaxSessionStoreSizeLimit (12-12)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: go-test (1.25.2, windows-latest)

🔇 Additional comments (9)

pkg/types/sessions/sessions.go (1)

7-12: LGTM! Well-documented constants for gRPC size limits.

The constants are correctly calculated (4MB - 30KB = 4163584 bytes) and appropriately centralized for reuse across modules. The comment clearly explains the rationale.

pkg/session/session.go (3)

67-72: LGTM! Clean struct design for carrying size metadata.

The SizedItem[T] struct effectively encapsulates the key-value pair with pre-computed size, enabling efficient size-aware batching without recalculating sizes during iteration.

---

74-111: Well-implemented size-aware batching with proper edge case handling.

The flush condition correctly handles:

• Key count limit: flushes when chunk reaches MaxKeysPerRequest
• Size limit: flushes when adding an item would exceed MaxSessionStoreSizeLimit, but only if chunk is non-empty (correctly allowing oversized single items to pass through)

The iterator-based approach enables lazy evaluation and proper error propagation.

---

34-34: LGTM!

Good refactor to use the centralized sessions.MaxKeysPerRequest constant.

pkg/session/session_test.go (5)

649-662: LGTM! Clean helper function for test setup.

The mapToSizedIter function correctly converts a map to a sized item iterator, computing size as len(key) + len(value) which matches the production usage pattern.

---

723-739: Good test coverage for iterator error propagation.

This test validates that errors yielded by the iterator are correctly surfaced, which is essential for the new iterator-based API.

---

756-800: Good test for size-based chunking behavior.

The test correctly validates that:

1. Multiple chunks are created when total size exceeds the limit
2. Each chunk stays within MaxSessionStoreSizeLimit

The assertion at line 797 using Less (strict less-than) is correct given the implementation flushes before exceeding the limit.

---

802-841: Excellent edge case test validating size-based chunking priority.

This test confirms that chunking is triggered by size constraints before hitting the key limit, which is the core behavior this PR introduces. The test design (10 items × 500KB = 5MB > 4MB limit, with only 10 keys < 100 key limit) effectively isolates the size-based chunking path.

---

682-684: Appropriate signature adaptation in existing tests.

The existing test cases are correctly updated to use mapToSizedIter to match the new iterator-based UnrollSetMany signature while preserving the original test intent.

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}