Skip to content

feat: add allocation store - #51

Open
alanshaw wants to merge 2 commits into
ash/fix/blob-release-causefrom
ash/feat/allocation-store
Open

feat: add allocation store#51
alanshaw wants to merge 2 commits into
ash/fix/blob-release-causefrom
ash/feat/allocation-store

Conversation

@alanshaw

@alanshaw alanshaw commented Aug 3, 2026

Copy link
Copy Markdown
Member

Adds a store for tracking allocations, and switches space diff accounting and metrics to be driven by it instead of the blob registry.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_registry stores (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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants