Rate limiting with a token bucket gate

// The SMS provider accepts 5 messages a second for each customer.
// A refused message comes back when its lease ends: here after 1 second.
await queen.queue('sms-out').config({ leaseTime: 1 }).create()

const limiter = await Stream
  .from(queen.queue('sms-out'))                          // one partition per customer
  .gate(tokenBucketGate({ capacity: 5, refillPerSec: 5 }))  // 5 at once, then 5 a second
  .foreach(async (sms) => sendSms(sms))                  // only what the gate lets through
  .run({ queryId: 'sms-gate', url, subscriptionMode: 'all' })

// Twelve messages for one customer, all at once.
await queen.queue('sms-out').partition('customer-42')
  .push(Array.from({ length: 12 }, (_, i) => ({ data: { n: i + 1, text: `#${i + 1}` } })))

# The SMS provider accepts 5 messages a second for each customer.
# A refused message comes back when its lease ends: here after 1 second.
await queen.queue("sms-out").config({"leaseTime": 1}).create()

limiter = await (
    Stream.from_(queen.queue("sms-out"))                    # one partition per customer
    .gate(token_bucket_gate(capacity=5, refill_per_sec=5))  # 5 at once, then 5 a second
    .foreach(send_sms)                                      # only what the gate allows
    .run(query_id="sms-gate", url=url, subscription_mode="all")
)

# Twelve messages for one customer, all at once.
await queen.queue("sms-out").partition("customer-42").push(
    [{"data": {"n": i + 1, "text": f"#{i + 1}"}} for i in range(12)]
)

// The SMS provider accepts 5 messages a second for each customer.
// A refused message comes back when its lease ends: here after 1 second.
_, err = client.Queue("sms-out").Config(queen.QueueConfig{LeaseTime: 1}).
	Create().Execute(ctx)
if err != nil {
	log.Fatal(err)
}

limiter, err := streams.
	From(client.Queue("sms-out").AsStreamSource()). // one partition per customer
	Gate(helpers.TokenBucketGate(helpers.TokenBucketGateOptions{
		Capacity: 5, RefillPerSec: 5, // 5 at once, then 5 a second
	})).
	Foreach(func(sms interface{}, _ streams.EmitCtx) error {
		return sendSms(sms) // only what the gate lets through
	}).
	Run(ctx, streams.RunOptions{QueryID: "sms-gate", URL: url,
		SubscriptionMode: queen.SubscriptionModeAll})
if err != nil {
	log.Fatal(err)
}

// Twelve messages for one customer, all at once.
burst := make([]interface{}, 12)
for i := range burst {
	burst[i] = map[string]interface{}{"n": i + 1, "text": fmt.Sprintf("#%d", i+1)}
}
_, err = client.Queue("sms-out").Partition("customer-42").Push(burst).Execute(ctx)
if err != nil {
	log.Fatal(err)
}

use queen_mq::streams::{GateCtx, Record};

// The Rust client has no tokenBucketGate helper; this is the same bucket. Its
// state is kept per partition in the broker, and saved only when a message passes.
fn token_bucket_gate(
    capacity: f64,
    refill_per_sec: f64,
) -> impl Fn(&Record, &mut GateCtx<'_>) -> bool + Send + Sync + 'static {
    move |_msg, ctx| {
        let now = ctx.stream_time_ms as f64;
        let elapsed = ((now - ctx.num("lastRefillAt", now)) / 1000.0).max(0.0);
        let tokens = (ctx.num("tokens", capacity) + elapsed * refill_per_sec).min(capacity);
        ctx.set_num("lastRefillAt", now);
        if tokens < 1.0 {
            return false; // refused: it comes back, in order, when its lease ends
        }
        ctx.set_num("tokens", tokens - 1.0);
        true
    }
}

// The SMS provider accepts 5 messages a second for each customer.
// A refused message comes back when its lease ends: here after 1 second.
queen
    .queue("sms-out")
    .configure(QueueOptions { lease_time: Some(1), ..Default::default() })
    .await?;

let limiter = Stream::from(queen.queue("sms-out")) // one partition per customer
    .gate(token_bucket_gate(5.0, 5.0))      // 5 at once, then 5 a second
    .foreach(|sms, _| async move { send_sms(sms).await }) // only what the gate lets through
    .run(&queen, RunOptions::new("sms-gate").subscription_mode(SubscriptionMode::All))
    .await?;

// Twelve messages for one customer, all at once.
queen
    .queue("sms-out")
    .partition("customer-42")
    .push_many((1..=12).map(|n| json!({ "n": n, "text": format!("#{n}") })))
    .await?;