feat: add allocation store - #51
Open
alanshaw wants to merge 2 commits into
Open
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces a new allocation store to track blob allocations (billing at allocation time), and refactors existing blob-registry-backed billing so that space_diff and metrics updates are driven by allocation lifecycle instead of blob registration/acceptance. It updates storage backends, dependency injection wiring, and service handlers/tests to use the new store.
Changes:
- Added
allocation.Store(Postgres + in-memory) with a Postgres migration and comprehensive store tests validating billing side-effects (space diffs + metrics). - Simplified
blob_registrystores (AWS/Postgres/memory) into plain record stores by removing embedded billing writes and their dependencies. - Updated blob lifecycle handlers (
/blob/add,/blob/remove,/blob/abort) and fx providers to create/remove allocation records at the appropriate points.
Reviewed changes
Copilot reviewed 23 out of 24 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| pkg/store/consumer/consumer.go | Adds CollectForSpace helper to page and collect consumers for billing attribution. |
| pkg/store/blob_registry/postgres/store.go | Removes transactional billing writes; registry becomes a plain Postgres record store. |
| pkg/store/blob_registry/memory/store.go | Removes billing side-effects; memory registry becomes a plain record store. |
| pkg/store/blob_registry/aws/store.go | Removes TransactWrite-based billing writes; uses simple PutItem/DeleteItem with conditional checks. |
| pkg/store/blob_registry/blob_registry_test.go | Refactors tests to validate blob registry behavior without consumer/metrics/space-diff dependencies. |
| pkg/store/allocation/allocation.go | Introduces the allocation store API, errors, pagination options, and record model. |
| pkg/store/allocation/postgres/store.go | Adds Postgres allocation implementation coordinating allocation + space_diff + metrics in a transaction. |
| pkg/store/allocation/memory/store.go | Adds in-memory allocation implementation that performs billing side-effects via injected stores. |
| pkg/store/allocation/allocation_test.go | Adds tests covering allocation CRUD + pagination + billing side-effects (diffs/metrics) for memory & Postgres. |
| pkg/service/handlers/blob_add.go | Records allocation after successful provider allocation to drive billing even pre-acceptance. |
| pkg/service/handlers/blob_add_test.go | Updates tests to include allocation store and assert allocation record creation/persistence semantics. |
| pkg/service/handlers/blob_remove.go | Removes allocation record for accepted blobs during /blob/remove (while preserving “unregistered keeps allocation” behavior). |
| pkg/service/handlers/blob_remove_test.go | Updates tests to provision allocation records and assert correct retention/removal behavior. |
| pkg/service/handlers/blob_abort.go | Removes allocation record after successful reject on /blob/abort (keeps record when blob was accepted). |
| pkg/service/handlers/blob_abort_test.go | Adds allocation store dependencies and asserts allocation removal/retention on abort paths. |
| pkg/service/handlers/blob_list_test.go | Updates blob registry setup to the new no-deps constructor; keeps consumer store for unrelated provisioning needs. |
| pkg/service/handlers/ucan_conclude_http_put_test.go | Updates blob registry construction to new signature (no billing deps). |
| internal/migrations/sql/00002_allocation.sql | Adds allocation table schema for Postgres backend. |
| internal/fx/store/postgres/provider.go | Wires allocation.Store (Postgres) and updates blob registry provider signature. |
| internal/fx/store/memory/provider.go | Wires allocation.Store (memory) and updates blob registry provider signature. |
| internal/fx/store/aws/provider.go | Wires allocation using the in-memory impl for the AWS backend; updates blob registry provider signature. |
| go.mod | Bumps libforge and go-cid versions. |
| go.sum | Updates sums for the bumped dependencies. |
| CLAUDE.md | Updates architecture/store documentation to include the new allocation store and backend guidance. |
Suppressed comments (1)
pkg/store/allocation/memory/store.go:165
- spaceDiffStore.Put errors are ignored during Remove, so billing diffs can silently fail. Also, the allocation record is deleted before diffs/metrics are written; if a billing write fails, the record is already gone and the operation can’t be retried to repair billing state.
s.allocs[space] = append(s.allocs[space][:idx], s.allocs[space][idx+1:]...)
// There should only be one subscription per provider, but in theory you
// could have multiple providers for the same consumer (space).
for _, c := range consumers {
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+61
to
+83
| s.allocs[space] = append(s.allocs[space], allocation.Record{ | ||
| Space: space, | ||
| Blob: blob, | ||
| Cause: cause, | ||
| InsertedAt: time.Now(), | ||
| }) | ||
|
|
||
| // There should only be one subscription per provider, but in theory you | ||
| // could have multiple providers for the same consumer (space). | ||
| for _, c := range consumers { | ||
| s.spaceDiffStore.Put(ctx, c.Provider, space, c.Subscription, cause, int64(blob.Size), time.Now()) | ||
| } | ||
|
|
||
| inc := map[string]uint64{ | ||
| metrics.BlobAddTotalMetric: 1, | ||
| metrics.BlobAddSizeTotalMetric: blob.Size, | ||
| } | ||
| if err := s.spaceMetrics.IncrementTotals(ctx, space, inc); err != nil { | ||
| return fmt.Errorf("incrementing space metrics: %w", err) | ||
| } | ||
| if err := s.adminMetrics.IncrementTotals(ctx, inc); err != nil { | ||
| return fmt.Errorf("incrementing admin metrics: %w", err) | ||
| } |
Comment on lines
+239
to
+247
| return allocation.Record{ | ||
| Space: space, | ||
| Blob: allocation.Blob{ | ||
| Digest: digest, | ||
| Size: uint64(size), | ||
| }, | ||
| Cause: cause, | ||
| InsertedAt: insertedAt, | ||
| }, nil |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a store for tracking allocations, and switches space diff accounting and metrics to be driven by it instead of the blob registry.