Refill tenant pool at a fixed rate with async dispatch#770
Conversation
Claim-triggered pool refill used to wait until free tenants dropped below 80% of pool size and then batch-create the entire deficit in one call, producing a request storm against the TiDB Cloud API (e.g. a 600-slot pool refilling 120+ clusters at once). Replace the watermark catch-up with a queue model: a per-pool background worker (kicked by claims) dispatches cluster creations at a fixed, configurable rate — TenantPoolRefillBatchSize per TenantPoolRefillInterval, default 1/s — decoupled from both claim bursts and batchCreate latency. Correctness under async dispatch: - Each dispatched creation first records a pending tenant plus a placeholder binding (cluster_id "pending:<tenantID>") inside the pool locks, so in-flight creations occupy slots in CountTenantPoolFreeSlots and concurrent dispatchers/pods never over-provision. Placeholder tenants stay Pending, so claims (Active-only) never take them, and the resume path skips them. - The async finisher persists under the pool locks and rechecks tenant status: if the placeholder was reclaimed by shrink/delete or the reaper while the API call ran, the created cluster is deprovisioned instead of resurrecting the slot. Pool shrink/delete skips cloud deprovision for placeholder slots. - A reaper fails placeholder tenants older than 15min (e.g. abandoned by a crashed pod), freeing the slot for redispatch; the worker stops after 5 consecutive creation failures and is re-kicked by the next claim. Config: DRIVE9_TENANT_POOL_REFILL_FREE_RATIO is removed; new env vars DRIVE9_TENANT_POOL_REFILL_BATCH_SIZE (default 1) and DRIVE9_TENANT_POOL_REFILL_INTERVAL (default 1s).
📝 WalkthroughWalkthroughTenant pool refilling now uses configurable batch sizes and intervals. Per-pool workers create placeholder tenants, provision clusters asynchronously, persist results, reap stale placeholders, handle reclaimed slots, and replace the previous free-ratio replenishment triggers. ChangesTenant pool refill
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant TenantClaim
participant RefillWorker
participant MetaStore
participant TiDBCloud
TenantClaim->>RefillWorker: kickTenantPoolRefill
RefillWorker->>MetaStore: create placeholder bindings
RefillWorker->>TiDBCloud: provision clusters in batches
TiDBCloud-->>RefillWorker: return cluster results
RefillWorker->>MetaStore: persist completed tenants
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
pkg/server/admin_tenant_pool.go (2)
1753-1762: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueReclaimed placeholders are counted as creation failures.
The deferred counter treats every non-
successoutcome as a failure, so the reclaimed path (pool shrink/delete or reaper) incrementsconsecutiveFailureseven though no cluster creation actually failed. Impact is bounded in practice (a shrink makes the next round returnrefillRoundFull/refillRoundExit), but semantically a reclaim shouldn't push the worker toward its failure breaker. Consider resetting/skipping the failure count whenreclaimedis true.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/server/admin_tenant_pool.go` around lines 1753 - 1762, Update the deferred counter logic in finishAsyncPoolClusterProvision to distinguish reclaimed outcomes from creation failures: when reclaimed is true, reset or skip consecutiveFailures instead of incrementing it; retain the existing success reset and failure increment behavior for actual creation outcomes.
1613-1650: 🩺 Stability & Availability | 🔵 Trivial | ⚖️ Poor tradeoffRound-level errors have no circuit breaker; worker retries forever.
consecutiveFailuresis only incremented by async creation outcomes infinishAsyncPoolClusterProvision. ArefillRoundRoundthat returnsrefillRoundError(e.g.GetTenantPoolByID,CountTenantPoolFreeSlots, orinsertPendingPoolTenantpersistently failing) is neitherrefillRoundExitnorrefillRoundFull, so the loop paces atintervaland retries indefinitely with no back-off escalation or stop condition other than context cancellation. Consider tracking consecutive round errors and stopping the worker (like the creation-failure breaker) so a durable meta-DB fault doesn't leave a permanently spinning per-pool worker.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/server/admin_tenant_pool.go` around lines 1613 - 1650, The runTenantPoolRefillWorker loop must apply the existing consecutive-failure circuit breaker to refillRoundError results, not only asynchronous creation failures. Track consecutive round errors, increment and log them when tenantPoolRefillRound returns refillRoundError, stop once the configured maximum is reached, and reset the round-error counter after a successful non-error round while preserving existing exit, full, and context behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@pkg/server/admin_tenant_pool.go`:
- Around line 1753-1762: Update the deferred counter logic in
finishAsyncPoolClusterProvision to distinguish reclaimed outcomes from creation
failures: when reclaimed is true, reset or skip consecutiveFailures instead of
incrementing it; retain the existing success reset and failure increment
behavior for actual creation outcomes.
- Around line 1613-1650: The runTenantPoolRefillWorker loop must apply the
existing consecutive-failure circuit breaker to refillRoundError results, not
only asynchronous creation failures. Track consecutive round errors, increment
and log them when tenantPoolRefillRound returns refillRoundError, stop once the
configured maximum is reached, and reset the round-error counter after a
successful non-error round while preserving existing exit, full, and context
behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 13d2585b-1d17-422a-a002-9a1a016b58b0
📒 Files selected for processing (9)
cmd/drive9-server/main.gocmd/drive9-server/main_test.gopkg/meta/meta.gopkg/server/admin_tenant_pool.gopkg/server/admin_tenant_pool_test.gopkg/server/admin_tenants.gopkg/server/quota_test.gopkg/server/server.gopkg/server/upload_test.go
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e8537d4d9f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| existing.setCred(cred) | ||
| existing.rerun.Store(true) | ||
| return |
There was a problem hiding this comment.
Do not drop kicks while the refill worker is exiting
When a claim arrives after the worker has already passed the rerun.Swap(false) full-pool check but before the goroutine's deferred Delete runs, this branch sets rerun and returns without starting another worker; the exiting worker never observes that flag and then removes the job. In that timing, the slot consumed by the claim is not refilled until some later successful claim kicks the pool again, so low-volume pools can remain below their configured size.
Useful? React with 👍 / 👎.
| clusters, batchCloudCfg, err := s.batchProvisionFreePoolClusters(ctx, poolID, []string{tenantID}, cred, nil) | ||
| if err != nil && len(clusters) == 0 { | ||
| // The batch helper already marked the tenant failed, which frees the | ||
| // placeholder slot for a later dispatch round. | ||
| return |
There was a problem hiding this comment.
Re-check reclaimed placeholders before marking failures
If shrink/delete reclaims this placeholder while the TiDB Cloud call is in flight and the provider later returns an error with zero clusters, batchProvisionFreePoolClusters marks the requested tenant failed before the status re-check below can run. Because UpdateTenantStatus is unconditional, this can change a tenant that deleteNewestFreePoolTenants just marked Deleted back to Failed, leaving orphan failed tenants after pool deletion/shrink instead of preserving the reclaimed state.
Useful? React with 👍 / 👎.
| s.startProvisionedTenantSchemaInit(r.Context(), res) | ||
| } | ||
| s.replenishTenantPoolAsync(r.Context(), pool, *credentialReq) | ||
| s.kickTenantPoolRefill(r.Context(), pool, *credentialReq) |
There was a problem hiding this comment.
Kick the refill worker after pool misses
Because this new kick is only executed inside the claimed branch, requests that find the pool but miss an Active free tenant never start the worker. When the new async path leaves only placeholders (for example, the pod that created them crashed, or the failure breaker drained the active slots), those misses fall through to direct provisioning and the only reaper/refill loop is never run, so the pool can stay empty or stuck until an admin operation intervenes.
Useful? React with 👍 / 👎.
| s.startServerWorker(ctx, func(workerCtx context.Context) { | ||
| s.finishAsyncPoolClusterProvision(workerCtx, poolID, orgID, tenantID, cred, job) | ||
| }) |
There was a problem hiding this comment.
Count in-flight failures before dispatching more creates
When each TiDB Cloud failure is slower than TenantPoolRefillInterval, this dispatch path keeps starting new async creates while earlier failed attempts are still in flight because the breaker at the top of the loop only sees completed failures. With a large deficit and bad/expired credentials, the default 1/s worker can launch many calls before the fifth failure is observed, defeating the intended max-consecutive-failure guard.
Useful? React with 👍 / 👎.
Motivation
Claim-triggered pool refill waited until free tenants dropped below 80% of pool size, then batch-created the entire deficit in one call — a 600-slot pool dipping below 480 free triggered a single catch-up batch of 120+ clusters, a request storm against the TiDB Cloud API.
New model: queue with a fixed-rate producer
A per-pool background worker (kicked by claims, credentials kept in memory only) treats the pool as a queue: whenever it is not full, it dispatches cluster creations at a fixed, configurable rate —
TenantPoolRefillBatchSizeperTenantPoolRefillInterval, default 1/s — decoupled from both claim bursts andbatchCreatelatency (tens of seconds in practice). Pool capacity now only absorbs peaks; the producer keeps pace with average consumption.Initial pool create (POST) and admin grow (PATCH) keep their explicit one-shot batch behavior.
Correctness under async dispatch
cluster_id = "pending:<tenantID>") inside the pool locks, so in-flight creations occupy slots inCountTenantPoolFreeSlotsacross pods and the dispatcher never over-provisions. Placeholder tenants stayPending`: claims (Active-only) can't take them, the metadata-resume path skips them.createFreePoolTenantsis split intobatchProvisionFreePoolClusters+persistFreePoolClustersso the admin batch path and the async finisher share the persistence logic.Config changes (breaking)
DRIVE9_TENANT_POOL_REFILL_FREE_RATIO(silently ignored now — call out in release notes)DRIVE9_TENANT_POOL_REFILL_BATCH_SIZE(default 1),DRIVE9_TENANT_POOL_REFILL_INTERVAL(default 1s)Note: with N pods each running a worker, aggregate creation rate is N × configured rate (same multi-pod behavior as the previous claim-triggered refill; DB locks still prevent over-provisioning).
Tests
uk_tidbcloud_org_cluster_branchconstraint).make test TEST_PKGS='./pkg/server'full suite,./pkg/meta,./cmd/drive9-server,make lint— all green.Summary by CodeRabbit
New Features
Bug Fixes
Documentation