| 1 | // 🤖 Package idalloc is the shared, global, monotonic id allocator behind |
| 2 | // region-neutral space addressing (multi-edge Phase 1, docs/multi-edge.md). |
| 3 | // |
| 4 | // The problem it solves: once a space's hostname is <sqids(globalID)>.dev.11ch.at |
| 5 | // — ONE namespace, no country label — the id feeding sqids must be unique across |
| 6 | // the WHOLE fleet, or two region clouds minting concurrently would produce the |
| 7 | // same hostname. Home is the primary caller, but demo and direct/operator |
| 8 | // creates bypass home (they go straight to a box's control daemon), so every |
| 9 | // place that mints an id must draw from the SAME sequence. |
| 10 | // |
| 11 | // The source of truth is a single small object in the shared S3 bucket (the |
| 12 | // mesh's substrate — LCHAT_HIBERNATE_BUCKET) holding the high-water mark: the |
| 13 | // count of ids reserved so far. Advancing it is a compare-and-swap on S3's ETag, |
| 14 | // exactly the conditional-write style the hibernation lease (internal/activator) |
| 15 | // and the home vault (home/s3store.go) already ride — no new dependency, no |
| 16 | // second S3 client. |
| 17 | // |
| 18 | // Block allocation is the scaling trick. Id allocation happens once per space |
| 19 | // create, and Paul wants this to reach billions of spaces; one S3 round-trip per |
| 20 | // create would put the counter object on the hot path of every create fleet-wide. |
| 21 | // Instead an allocator reserves a RANGE of N ids in one CAS, then hands them out |
| 22 | // from memory, refilling with another CAS only when the block is spent. That |
| 23 | // bounds S3 ops to ~(creates / N) — with the default N=100, a rounding error on |
| 24 | // the bill even at scale (#836's lesson: request count IS the S3 cost). |
| 25 | // |
| 26 | // Guarantees: ids are unique and strictly increasing across all allocators (the |
| 27 | // CAS makes overlap impossible — a conflict just retries against the new |
| 28 | // high-water). The high-water lives in S3, so a restarted process continues past |
| 29 | // it, skipping the unused tail of its last block; ids need only be unique and |
| 30 | // increasing, never gapless, so skipped tails are fine. In-process hand-outs are |
| 31 | // mutex-guarded. |
| 32 | // |
| 33 | // Where it plugs in (the NEXT slice — this package does not wire it): cloud's |
| 34 | // createSpaceHandler (cloud.go) currently takes the id from the cloud.db spaces |
| 35 | // PK via cloudDB.allocateSpace; that id source becomes Allocate(ctx). The demo |
| 36 | // path (demo.go) and controld's direct add-instance create the same way. Home, |
| 37 | // cloud, and controld each construct an Allocator from the shared bucket and use |
| 38 | // it identically. |
| 39 | package idalloc |
| 40 | |
| 41 | import ( |
| 42 | "bytes" |
| 43 | "context" |
| 44 | "errors" |
| 45 | "fmt" |
| 46 | "io" |
| 47 | "math/rand" |
| 48 | "net/http" |
| 49 | "path" |
| 50 | "strconv" |
| 51 | "strings" |
| 52 | "sync" |
| 53 | "time" |
| 54 | |
| 55 | "github.com/aws/aws-sdk-go-v2/aws" |
| 56 | awsconfig "github.com/aws/aws-sdk-go-v2/config" |
| 57 | awss3 "github.com/aws/aws-sdk-go-v2/service/s3" |
| 58 | "github.com/aws/smithy-go" |
| 59 | ) |
| 60 | |
| 61 | // DefaultBlockSize is the number of ids reserved per S3 CAS when a caller names |
| 62 | // none. 100 keeps S3 ops to ~1% of creates while wasting at most 99 ids on a |
| 63 | // restart — ids are cheap (uint64) and gaps are allowed, so a larger block only |
| 64 | // trades a slightly bigger restart gap for even fewer S3 ops. |
| 65 | const DefaultBlockSize uint64 = 100 |
| 66 | |
| 67 | // counterKey is the object under the bucket prefix that holds the high-water |
| 68 | // mark. It sits beside the mesh's membership heartbeats (<prefix>/mesh/members), |
| 69 | // under the same "mesh/" namespace — the bare label never collides with a |
| 70 | // per-instance prefix, which is always a dotted FQDN. |
| 71 | const counterKey = "mesh/id-sequence" |
| 72 | |
| 73 | // maxReserveAttempts caps the CAS-conflict retry loop. Block allocation makes a |
| 74 | // conflict a 1-in-N event and each conflict means a competitor made progress |
| 75 | // (the high-water strictly rose), so the loser wins on a later attempt; the |
| 76 | // number of real allocator instances is tiny (home + a handful of region clouds |
| 77 | // + controld), so contention never approaches this. It exists only so a |
| 78 | // pathological store (e.g. a swap that always reports conflict) fails loudly |
| 79 | // instead of spinning forever. The context deadline is the primary bound. |
| 80 | const maxReserveAttempts = 100 |
| 81 | |
| 82 | // s3CallTimeout bounds each individual S3 call so a stalled counter object can't |
| 83 | // hang a space create indefinitely (the caller's ctx is the outer bound; this is |
| 84 | // the per-request floor, matching s3blobstore.go's per-call timeout). |
| 85 | const s3CallTimeout = 30 * time.Second |
| 86 | |
| 87 | // counterStore is the minimal compare-and-swap surface the allocator needs over |
| 88 | // the one shared counter object. Splitting it out keeps the block/retry logic |
| 89 | // pure and unit-testable against an in-memory fake, while the S3-backed |
| 90 | // implementation stays a thin adapter. |
| 91 | // |
| 92 | // The version token is opaque (S3's ETag in production). By convention an empty |
| 93 | // token means "the object does not exist yet" — load returns "" when absent, and |
| 94 | // swap treats "" as a create-only precondition. S3 never returns an empty ETag |
| 95 | // for an existing object, so the convention is unambiguous there. |
| 96 | type counterStore interface { |
| 97 | // load reads the current high-water value and its version token. A missing |
| 98 | // object reads as value 0 with an empty token (the first writer's state). |
| 99 | load(ctx context.Context) (value uint64, token string, err error) |
| 100 | // swap writes newValue only if the object still matches token (or is still |
| 101 | // absent, when token is ""). ok=false means a concurrent writer won the race |
| 102 | // and the caller must reload and retry; a non-nil err is a real failure. |
| 103 | swap(ctx context.Context, token string, newValue uint64) (ok bool, err error) |
| 104 | } |
| 105 | |
| 106 | // Allocator hands out globally-unique, strictly-increasing ids drawn from blocks |
| 107 | // it reserves in the shared S3 counter. Safe for concurrent use. |
| 108 | type Allocator struct { |
| 109 | store counterStore |
| 110 | blockSize uint64 |
| 111 | |
| 112 | // mu guards the in-memory block AND serializes refills: holding it across the |
| 113 | // refill CAS means two goroutines can never both reserve a block (which would |
| 114 | // skip a whole block's worth of ids — still correct, but wasteful). Refills |
| 115 | // are 1-in-blockSize, so the serialization is negligible. |
| 116 | mu sync.Mutex |
| 117 | next uint64 // next id to hand out; ids are 1-based |
| 118 | end uint64 // one past the last id in the reserved block; next>=end ⇒ exhausted |
| 119 | } |
| 120 | |
| 121 | // Config constructs an S3-backed Allocator. Bucket/Prefix/Region mirror the |
| 122 | // hibernation config (LCHAT_HIBERNATE_*) so a caller that already reads that env |
| 123 | // can build one with no new plumbing; the counter object lands at |
| 124 | // <Prefix>/mesh/id-sequence in Bucket. |
| 125 | type Config struct { |
| 126 | Bucket string // shared S3 bucket (the mesh substrate); required |
| 127 | Prefix string // key prefix inside the bucket (may be empty) |
| 128 | Region string // optional; the AWS default chain resolves "" from env/profile |
| 129 | Endpoint string // optional S3-compatible endpoint (R2, a test fake); implies path-style |
| 130 | BlockSize uint64 // ids per CAS; 0 → DefaultBlockSize |
| 131 | } |
| 132 | |
| 133 | // New builds an S3-backed Allocator. Credentials ride the standard AWS chain |
| 134 | // (the same as the hibernation lease and blob store); Endpoint lets an |
| 135 | // S3-compatible store or a test fake stand in. |
| 136 | func New(ctx context.Context, cfg Config) (*Allocator, error) { |
| 137 | if strings.TrimSpace(cfg.Bucket) == "" { |
| 138 | return nil, errors.New("idalloc: bucket is required") |
| 139 | } |
| 140 | var opts []func(*awsconfig.LoadOptions) error |
| 141 | if cfg.Region != "" { |
| 142 | opts = append(opts, awsconfig.WithRegion(cfg.Region)) |
| 143 | } |
| 144 | awsCfg, err := awsconfig.LoadDefaultConfig(ctx, opts...) |
| 145 | if err != nil { |
| 146 | return nil, fmt.Errorf("idalloc: load aws config: %w", err) |
| 147 | } |
| 148 | client := awss3.NewFromConfig(awsCfg, func(o *awss3.Options) { |
| 149 | if cfg.Endpoint != "" { |
| 150 | o.BaseEndpoint = &cfg.Endpoint |
| 151 | o.UsePathStyle = true // custom endpoints (R2, fakes) address by path |
| 152 | } |
| 153 | // Match s3blobstore.go: the SDK's default-on flexible checksums wrap |
| 154 | // PutObject bodies in aws-chunked framing that non-AWS stores reject; the |
| 155 | // counter body is tiny with a known length, so plain bodies are correct. |
| 156 | o.RequestChecksumCalculation = aws.RequestChecksumCalculationWhenRequired |
| 157 | o.ResponseChecksumValidation = aws.ResponseChecksumValidationWhenRequired |
| 158 | }) |
| 159 | store := &s3Counter{ |
| 160 | client: client, |
| 161 | bucket: cfg.Bucket, |
| 162 | key: path.Join(cfg.Prefix, counterKey), |
| 163 | } |
| 164 | return newWithStore(store, cfg.BlockSize), nil |
| 165 | } |
| 166 | |
| 167 | // newWithStore builds an Allocator over any counterStore — the seam the tests |
| 168 | // drive with an in-memory fake, and the shared constructor path for New. |
| 169 | func newWithStore(store counterStore, blockSize uint64) *Allocator { |
| 170 | if blockSize == 0 { |
| 171 | blockSize = DefaultBlockSize |
| 172 | } |
| 173 | return &Allocator{store: store, blockSize: blockSize} |
| 174 | } |
| 175 | |
| 176 | // Allocate returns the next globally-unique id, reserving a fresh block from S3 |
| 177 | // when the current one is spent. |
| 178 | func (a *Allocator) Allocate(ctx context.Context) (uint64, error) { |
| 179 | a.mu.Lock() |
| 180 | defer a.mu.Unlock() |
| 181 | if a.next >= a.end { |
| 182 | if err := a.reserveLocked(ctx); err != nil { |
| 183 | return 0, err |
| 184 | } |
| 185 | } |
| 186 | id := a.next |
| 187 | a.next++ |
| 188 | return id, nil |
| 189 | } |
| 190 | |
| 191 | // reserveLocked reserves the next block via a CAS on the shared counter, |
| 192 | // retrying on conflict. Called with a.mu held. |
| 193 | // |
| 194 | // The CAS: read the high-water H (the count reserved so far) and its version |
| 195 | // token, then conditional-write H+blockSize keyed on that token. Success means |
| 196 | // no other allocator advanced the counter in between, so this process now owns |
| 197 | // ids [H+1, H+blockSize]. A conflict means a competitor reserved a block first; |
| 198 | // its write bumped the token, so ours is stale and the store reports ok=false — |
| 199 | // we reload the now-higher H and try again, never reusing the ids the competitor |
| 200 | // took. |
| 201 | func (a *Allocator) reserveLocked(ctx context.Context) error { |
| 202 | for attempt := 0; attempt < maxReserveAttempts; attempt++ { |
| 203 | if err := ctx.Err(); err != nil { |
| 204 | return err |
| 205 | } |
| 206 | high, token, err := a.store.load(ctx) |
| 207 | if err != nil { |
| 208 | return fmt.Errorf("idalloc: load counter: %w", err) |
| 209 | } |
| 210 | newHigh := high + a.blockSize |
| 211 | ok, err := a.store.swap(ctx, token, newHigh) |
| 212 | if err != nil { |
| 213 | return fmt.Errorf("idalloc: reserve block: %w", err) |
| 214 | } |
| 215 | if ok { |
| 216 | a.next = high + 1 // ids are 1-based, so the block starts at H+1 |
| 217 | a.end = newHigh + 1 // exclusive: hand out H+1 .. H+blockSize |
| 218 | return nil |
| 219 | } |
| 220 | // Lost the race — back off a touch before reloading, so a burst of |
| 221 | // allocators contending on a fresh counter doesn't livelock hammering it. |
| 222 | if err := sleepCtx(ctx, backoff(attempt)); err != nil { |
| 223 | return err |
| 224 | } |
| 225 | } |
| 226 | return fmt.Errorf("idalloc: reserve block: gave up after %d CAS conflicts", maxReserveAttempts) |
| 227 | } |
| 228 | |
| 229 | // backoff is capped exponential with full jitter — small, since a conflict |
| 230 | // resolves in one or two retries in any realistic fleet. |
| 231 | func backoff(attempt int) time.Duration { |
| 232 | const base = 5 * time.Millisecond |
| 233 | const cap = 250 * time.Millisecond |
| 234 | d := base << attempt |
| 235 | if d > cap || d <= 0 { // <=0 guards the shift overflowing at high attempts |
| 236 | d = cap |
| 237 | } |
| 238 | return time.Duration(rand.Int63n(int64(d) + 1)) |
| 239 | } |
| 240 | |
| 241 | // sleepCtx sleeps for d unless ctx is cancelled first. |
| 242 | func sleepCtx(ctx context.Context, d time.Duration) error { |
| 243 | if d <= 0 { |
| 244 | return ctx.Err() |
| 245 | } |
| 246 | t := time.NewTimer(d) |
| 247 | defer t.Stop() |
| 248 | select { |
| 249 | case <-ctx.Done(): |
| 250 | return ctx.Err() |
| 251 | case <-t.C: |
| 252 | return nil |
| 253 | } |
| 254 | } |
| 255 | |
| 256 | // --- S3-backed counterStore -------------------------------------------------- |
| 257 | |
| 258 | // s3Counter is the production counterStore: the high-water lives as a decimal |
| 259 | // string in one S3 object, and CAS rides S3's ETag conditional writes — the same |
| 260 | // mechanism home/s3store.go uses for its user docs, over the same aws-sdk-go-v2 |
| 261 | // client s3blobstore.go builds. |
| 262 | type s3Counter struct { |
| 263 | client *awss3.Client |
| 264 | bucket string |
| 265 | key string |
| 266 | } |
| 267 | |
| 268 | func (c *s3Counter) load(ctx context.Context) (uint64, string, error) { |
| 269 | ctx, cancel := context.WithTimeout(ctx, s3CallTimeout) |
| 270 | defer cancel() |
| 271 | out, err := c.client.GetObject(ctx, &awss3.GetObjectInput{Bucket: &c.bucket, Key: &c.key}) |
| 272 | if err != nil { |
| 273 | if isNotFound(err) { |
| 274 | return 0, "", nil // never allocated yet: value 0, empty (create-only) token |
| 275 | } |
| 276 | return 0, "", err |
| 277 | } |
| 278 | defer out.Body.Close() |
| 279 | // The body is a short decimal; cap the read so a corrupt/huge object can't |
| 280 | // balloon memory. |
| 281 | b, err := io.ReadAll(io.LimitReader(out.Body, 64)) |
| 282 | if err != nil { |
| 283 | return 0, "", err |
| 284 | } |
| 285 | v, err := strconv.ParseUint(strings.TrimSpace(string(b)), 10, 64) |
| 286 | if err != nil { |
| 287 | return 0, "", fmt.Errorf("idalloc: corrupt counter object %q: %w", strings.TrimSpace(string(b)), err) |
| 288 | } |
| 289 | token := "" |
| 290 | if out.ETag != nil { |
| 291 | token = *out.ETag // S3 ETags are quoted; kept verbatim for If-Match |
| 292 | } |
| 293 | return v, token, nil |
| 294 | } |
| 295 | |
| 296 | func (c *s3Counter) swap(ctx context.Context, token string, newValue uint64) (bool, error) { |
| 297 | ctx, cancel := context.WithTimeout(ctx, s3CallTimeout) |
| 298 | defer cancel() |
| 299 | body := []byte(strconv.FormatUint(newValue, 10)) |
| 300 | n := int64(len(body)) |
| 301 | in := &awss3.PutObjectInput{ |
| 302 | Bucket: &c.bucket, |
| 303 | Key: &c.key, |
| 304 | Body: bytes.NewReader(body), |
| 305 | ContentLength: &n, |
| 306 | } |
| 307 | if token == "" { |
| 308 | star := "*" |
| 309 | in.IfNoneMatch = &star // create-only: fails 412 if someone created it first |
| 310 | } else { |
| 311 | in.IfMatch = &token // replace-only: fails 412 if the object moved on |
| 312 | } |
| 313 | if _, err := c.client.PutObject(ctx, in); err != nil { |
| 314 | if isCASConflict(err) { |
| 315 | return false, nil |
| 316 | } |
| 317 | return false, err |
| 318 | } |
| 319 | return true, nil |
| 320 | } |
| 321 | |
| 322 | // isNotFound folds the shapes a missing object answers with (mirrors |
| 323 | // s3blobstore.go's s3NotFound — kept local so idalloc doesn't import package |
| 324 | // main). |
| 325 | func isNotFound(err error) bool { |
| 326 | var apiErr smithy.APIError |
| 327 | if errors.As(err, &apiErr) { |
| 328 | switch apiErr.ErrorCode() { |
| 329 | case "NoSuchKey", "NotFound", "404": |
| 330 | return true |
| 331 | } |
| 332 | } |
| 333 | return false |
| 334 | } |
| 335 | |
| 336 | // isCASConflict reports whether a PutObject error is a lost conditional write: |
| 337 | // 412 (If-Match / If-None-Match precondition failed) or 409 (S3's |
| 338 | // ConditionalRequestConflict when two conditional writes race). Matched by both |
| 339 | // the S3 error code and the raw HTTP status, so an S3-compatible store that only |
| 340 | // sets one of them still classifies correctly. |
| 341 | func isCASConflict(err error) bool { |
| 342 | var apiErr smithy.APIError |
| 343 | if errors.As(err, &apiErr) { |
| 344 | switch apiErr.ErrorCode() { |
| 345 | case "PreconditionFailed", "ConditionalRequestConflict": |
| 346 | return true |
| 347 | } |
| 348 | } |
| 349 | var httpErr interface{ HTTPStatusCode() int } |
| 350 | if errors.As(err, &httpErr) { |
| 351 | switch httpErr.HTTPStatusCode() { |
| 352 | case http.StatusPreconditionFailed, http.StatusConflict: |
| 353 | return true |
| 354 | } |
| 355 | } |
| 356 | return false |
| 357 | } |
| 358 | |
amadan