Skip to content

Refill tenant pool at a fixed rate with async dispatch#770

Draft
mornyx wants to merge 1 commit into
mainfrom
tenant-pool-async-refill
Draft

Refill tenant pool at a fixed rate with async dispatch#770
mornyx wants to merge 1 commit into
mainfrom
tenant-pool-async-refill

Conversation

@mornyx

@mornyx mornyx commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

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 — TenantPoolRefillBatchSize per TenantPoolRefillInterval, default 1/s — decoupled from both claim bursts and batchCreate latency (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

  • In-flight visibility: 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 CountTenantPoolFreeSlotsacross pods and the dispatcher never over-provisions. Placeholder tenants stayPending`: claims (Active-only) can't take them, the metadata-resume path skips them.
  • No resurrection on reclaim: the async finisher persists under the same pool locks used by shrink/delete and rechecks tenant status first; if the placeholder was reclaimed (shrink, pool delete, or the reaper) while the API call ran, the created cluster is deprovisioned instead of resurrecting the slot. Shrink/delete skips cloud deprovision for placeholder slots.
  • Orphan reaper + circuit breaker: placeholder tenants older than 15 min (e.g. abandoned by a crashed pod) are failed to free the slot; the worker stops after 5 consecutive creation failures and is re-kicked by the next claim.

createFreePoolTenants is split into batchProvisionFreePoolClusters + persistFreePoolClusters so the admin batch path and the async finisher share the persistence logic.

Config changes (breaking)

  • Removed: DRIVE9_TENANT_POOL_REFILL_FREE_RATIO (silently ignored now — call out in release notes)
  • Added: 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

  • New: placeholder slots visible in flight (counted as slots, not claimable), stale-placeholder reaper frees and refills slots, reclaimed placeholder deprovisions the created cluster, worker stops after consecutive failures, coalesced kicks, pool-full/deleted worker exit.
  • Updated: refill batching expectations (async one-cluster calls), env parsing tests, fake provisioner generates per-call unique cluster IDs (real uk_tidbcloud_org_cluster_branch constraint).
  • Verified: make test TEST_PKGS='./pkg/server' full suite, ./pkg/meta, ./cmd/drive9-server, make lint — all green.

Summary by CodeRabbit

  • New Features

    • Tenant pool refilling now occurs in configurable batches at a steady interval.
    • Added configuration options for refill batch size and interval.
    • Pool capacity is replenished automatically after tenants are claimed.
  • Bug Fixes

    • Improved handling of stalled or abandoned tenant provisioning.
    • Prevented incomplete provisioning from occupying pool capacity indefinitely.
    • Ensured reclaimed in-progress slots are cleaned up safely.
  • Documentation

    • Updated command-line help text with the new refill settings and rate-limiting behavior.

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).
@coderabbitai

coderabbitai Bot commented Jul 21, 2026

Copy link
Copy Markdown

Review Change Stack

📝 Walkthrough

Walkthrough

Tenant 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.

Changes

Tenant pool refill

Layer / File(s) Summary
Refill configuration and wiring
pkg/server/server.go, cmd/drive9-server/*, pkg/server/upload_test.go, cmd/drive9-server/main_test.go
Free-ratio configuration is replaced by batch-size and interval settings with environment parsing, defaults, normalization, help text, and updated tests.
Placeholder provisioning and persistence
pkg/server/admin_tenant_pool.go, pkg/meta/meta.go, pkg/server/quota_test.go
Batch provisioning and persistence are separated; placeholder tenants are excluded from resume, protected during deletion, and stale placeholders can be failed transactionally.
Paced refill worker lifecycle
pkg/server/admin_tenant_pool.go, pkg/server/admin_tenants.go, pkg/server/server.go, pkg/server/admin_tenant_pool_test.go, pkg/server/quota_test.go
Per-pool workers dispatch fixed batches at intervals, coalesce kicks, stop on failures or pool deletion, reap stale placeholders, and deprovision clusters created for reclaimed slots.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

  • mem9-ai/drive9#632: Introduces the tenant-pool APIs and earlier refill flows that this change refactors.
  • mem9-ai/drive9#645: Modifies the same tenant-pool provisioning and refill pipeline.
  • mem9-ai/drive9#763: Overlaps with batch provisioning, quota, and cloud-configuration persistence paths.

Suggested reviewers: srstack, qiffang

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
Loading
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly matches the main change: moving tenant pool refill to fixed-rate async dispatch.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch tenant-pool-async-refill

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot 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.

🧹 Nitpick comments (2)
pkg/server/admin_tenant_pool.go (2)

1753-1762: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low value

Reclaimed placeholders are counted as creation failures.

The deferred counter treats every non-success outcome as a failure, so the reclaimed path (pool shrink/delete or reaper) increments consecutiveFailures even though no cluster creation actually failed. Impact is bounded in practice (a shrink makes the next round return refillRoundFull/refillRoundExit), but semantically a reclaim shouldn't push the worker toward its failure breaker. Consider resetting/skipping the failure count when reclaimed is 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 tradeoff

Round-level errors have no circuit breaker; worker retries forever.

consecutiveFailures is only incremented by async creation outcomes in finishAsyncPoolClusterProvision. A refillRoundRound that returns refillRoundError (e.g. GetTenantPoolByID, CountTenantPoolFreeSlots, or insertPendingPoolTenant persistently failing) is neither refillRoundExit nor refillRoundFull, so the loop paces at interval and 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

📥 Commits

Reviewing files that changed from the base of the PR and between cd8e004 and e8537d4.

📒 Files selected for processing (9)
  • cmd/drive9-server/main.go
  • cmd/drive9-server/main_test.go
  • pkg/meta/meta.go
  • pkg/server/admin_tenant_pool.go
  • pkg/server/admin_tenant_pool_test.go
  • pkg/server/admin_tenants.go
  • pkg/server/quota_test.go
  • pkg/server/server.go
  • pkg/server/upload_test.go

@chatgpt-codex-connector chatgpt-codex-connector Bot 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.

💡 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".

Comment on lines +1602 to 1604
existing.setCred(cred)
existing.rerun.Store(true)
return

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge 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 👍 / 👎.

Comment on lines +1763 to +1767
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

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge 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 👍 / 👎.

Comment thread pkg/server/server.go
s.startProvisionedTenantSchemaInit(r.Context(), res)
}
s.replenishTenantPoolAsync(r.Context(), pool, *credentialReq)
s.kickTenantPoolRefill(r.Context(), pool, *credentialReq)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge 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 👍 / 👎.

Comment on lines +1726 to +1728
s.startServerWorker(ctx, func(workerCtx context.Context) {
s.finishAsyncPoolClusterProvision(workerCtx, poolID, orgID, tenantID, cred, job)
})

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge 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 👍 / 👎.

@mornyx
mornyx marked this pull request as draft July 21, 2026 13:07
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.

1 participant