"fern/pages/components/frontend/frontend-guide.md" did not exist on "d59b9d72ed98eed50e8863de8259d0871efad336"
kv_router.rs 20.8 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
// SPDX-License-Identifier: Apache-2.0

Yan Ru Pei's avatar
Yan Ru Pei committed
4
use std::sync::Arc;
5
use std::time::{Duration, Instant};
6

7
use anyhow::Result;
Yan Ru Pei's avatar
Yan Ru Pei committed
8
use dynamo_kv_router::{ConcurrentRadixTree, ThreadPoolIndexer};
9
use dynamo_runtime::{
10
    component::{Client, Endpoint},
11
    discovery::DiscoveryQuery,
12
    pipeline::{
13
14
        AsyncEngine, AsyncEngineContextProvider, Error, ManyOut, ResponseStream, SingleIn,
        async_trait,
15
    },
16
    protocols::EndpointId,
17
    protocols::annotated::Annotated,
18
    traits::DistributedRuntimeProvider,
19
};
20
use futures::stream;
Yan Ru Pei's avatar
Yan Ru Pei committed
21
use tokio::sync::oneshot;
22
use tracing::Instrument;
23
use validator::Validate;
24

25
26
27
28
// Re-export from dynamo-kv-router crate
pub use dynamo_kv_router::approx;
pub use dynamo_kv_router::indexer;
pub use dynamo_kv_router::protocols;
29
30
pub use dynamo_kv_router::scheduling;
pub use dynamo_kv_router::selector;
31

32
pub mod cache_control;
33
pub mod config;
34
mod jetstream;
35
pub mod metrics;
36
pub mod prefill_router;
37
pub mod publisher;
38
pub mod push_router;
39
pub mod queue;
40
pub mod recorder;
41
pub mod scheduler;
42
pub mod sequence;
43
pub mod subscriber;
44
pub mod worker_query;
45

46
pub use cache_control::{CacheControlClient, spawn_pin_prefix};
47
pub use config::{KvRouterConfig, RouterConfigOverride};
48
pub use prefill_router::PrefillRouter;
49
pub use push_router::{DirectRoutingRouter, KvPushRouter};
50

51
use crate::{
52
    discovery::RuntimeConfigWatch,
53
    kv_router::{
54
        approx::PruneConfig,
Yan Ru Pei's avatar
Yan Ru Pei committed
55
        indexer::{GetWorkersRequest, KvIndexer, KvIndexerInterface, KvRouterError},
Yan Ru Pei's avatar
Yan Ru Pei committed
56
        protocols::{
57
            BlockExtraInfo, DpRank, LocalBlockHash, OverlapScores, RouterEvent, RouterRequest,
58
            RouterResponse, TokensWithHashes, WorkerId, WorkerWithDpRank,
Yan Ru Pei's avatar
Yan Ru Pei committed
59
            compute_block_hash_for_seq,
Yan Ru Pei's avatar
Yan Ru Pei committed
60
        },
61
        scheduler::{KvScheduler, PotentialLoad},
62
        sequence::{SequenceError, SequenceRequest},
63
    },
64
    local_model::runtime_config::ModelRuntimeConfig,
65
66
};

67
68
use std::collections::HashSet;

69
70
// [gluo TODO] shouldn't need to be public
// this should be discovered from the component
71
72
73
74
75

// for metric scraping (pull-based)
pub const KV_METRICS_ENDPOINT: &str = "load_metrics";

// for metric publishing (push-based)
76
pub const KV_EVENT_SUBJECT: &str = "kv-events";
77
78
79
80
81
pub const KV_METRICS_SUBJECT: &str = "kv_metrics";

// for inter-router comms
pub const PREFILL_SUBJECT: &str = "prefill_events";
pub const ACTIVE_SEQUENCES_SUBJECT: &str = "active_sequences_events";
82

83
84
85
86
// for radix tree snapshot storage
pub const RADIX_STATE_BUCKET: &str = "radix-bucket";
pub const RADIX_STATE_FILE: &str = "radix-state";

87
88
89
// for standalone indexer query
pub const KV_INDEXER_QUERY_ENDPOINT: &str = "kv_indexer_query";

90
91
92
// for worker-local kvindexer query
pub const WORKER_KV_INDEXER_BUFFER_SIZE: usize = 1024; // store 1024 most recent events in worker buffer

93
94
95
96
97
98
/// Generates a dp_rank-specific endpoint name for the worker KV indexer query service.
/// Each dp_rank has its own LocalKvIndexer and query endpoint to ensure per-dp_rank monotonicity.
pub fn worker_kv_indexer_query_endpoint(dp_rank: DpRank) -> String {
    format!("worker_kv_indexer_query_dp{dp_rank}")
}

99
// for router discovery registration
100
pub const KV_ROUTER_ENDPOINT: &str = "router-discovery";
101
102

/// Creates an EndpointId for the KV router in the given namespace.
103
pub fn router_endpoint_id(namespace: String, component: String) -> EndpointId {
104
105
    EndpointId {
        namespace,
106
        component,
107
108
109
110
111
        name: KV_ROUTER_ENDPOINT.to_string(),
    }
}

/// Creates a DiscoveryQuery for the KV router in the given namespace.
112
pub fn router_discovery_query(namespace: String, component: String) -> DiscoveryQuery {
113
114
    DiscoveryQuery::Endpoint {
        namespace,
115
        component,
116
117
118
119
        endpoint: KV_ROUTER_ENDPOINT.to_string(),
    }
}

120
121
122
/// Concrete `WorkerSelector` bound to the runtime config type.
pub type WorkerSelector =
    dyn dynamo_kv_router::selector::WorkerSelector<ModelRuntimeConfig> + Send + Sync;
123

Yan Ru Pei's avatar
Yan Ru Pei committed
124
#[derive(Clone)]
125
pub enum Indexer {
Yan Ru Pei's avatar
Yan Ru Pei committed
126
    /// Single-threaded radix tree with channel-based event processing.
127
    /// Supports TTL-based expiration and size-based pruning.
128
    /// Has the ability to persist and snapshot states.
129
    KvIndexer(KvIndexer),
130

Yan Ru Pei's avatar
Yan Ru Pei committed
131
132
133
134
135
    /// Concurrent radix tree with a thread pool for event processing.
    /// Uses sticky worker routing for per-worker event serialization.
    /// Does not support TTL/pruning.
    Concurrent(Arc<ThreadPoolIndexer<ConcurrentRadixTree>>),

136
137
138
    /// Used when we do not wish to use the indexer at all (e.g., when overlap_score_weight is 0).
    /// Note: This will cause KV events to accumulate in JetStream as we do not regularly purge them.
    None,
139
140
141
}

impl Indexer {
142
143
144
145
146
147
    pub fn new(
        component: &dynamo_runtime::component::Component,
        kv_router_config: &KvRouterConfig,
        block_size: u32,
    ) -> Self {
        if kv_router_config.overlap_score_weight == 0.0 {
Yan Ru Pei's avatar
Yan Ru Pei committed
148
149
150
            return Indexer::None;
        }

151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
        // Approximate mode (--no-kv-events): always use single-threaded KvIndexer
        // with TTL/pruning regardless of event_threads, since updates come from
        // routing decisions only, not live KV events from workers.
        if !kv_router_config.use_kv_events {
            let kv_indexer_metrics = indexer::KvIndexerMetrics::from_component(component);
            let cancellation_token = component.drt().primary_token();
            let prune_config = Some(PruneConfig {
                ttl: Duration::from_secs_f64(kv_router_config.router_ttl_secs),
                max_tree_size: kv_router_config.router_max_tree_size,
                prune_target_ratio: kv_router_config.router_prune_target_ratio,
            });
            return Indexer::KvIndexer(KvIndexer::new_with_frequency(
                cancellation_token,
                None,
                block_size,
                kv_indexer_metrics,
                prune_config,
            ));
        }

Yan Ru Pei's avatar
Yan Ru Pei committed
171
172
173
174
        if kv_router_config.router_event_threads > 1 {
            return Indexer::Concurrent(Arc::new(ThreadPoolIndexer::new(
                ConcurrentRadixTree::new(),
                kv_router_config.router_event_threads as usize,
175
                block_size,
Yan Ru Pei's avatar
Yan Ru Pei committed
176
            )));
177
        }
Yan Ru Pei's avatar
Yan Ru Pei committed
178
179

        let kv_indexer_metrics = indexer::KvIndexerMetrics::from_component(component);
180
        let cancellation_token = component.drt().primary_token();
Yan Ru Pei's avatar
Yan Ru Pei committed
181
182
183
184
185
186

        Indexer::KvIndexer(KvIndexer::new_with_frequency(
            cancellation_token,
            None, // expiration_duration for frequency tracking
            block_size,
            kv_indexer_metrics,
187
            None,
Yan Ru Pei's avatar
Yan Ru Pei committed
188
        ))
189
190
191
    }

    pub(crate) async fn find_matches(
192
193
194
195
196
        &self,
        sequence: Vec<LocalBlockHash>,
    ) -> Result<OverlapScores, KvRouterError> {
        match self {
            Indexer::KvIndexer(indexer) => indexer.find_matches(sequence).await,
Yan Ru Pei's avatar
Yan Ru Pei committed
197
            Indexer::Concurrent(tpi) => tpi.find_matches(sequence).await,
198
            Indexer::None => Ok(OverlapScores::new()),
199
200
        }
    }
201

202
    pub(crate) async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
203
204
        match self {
            Indexer::KvIndexer(indexer) => indexer.dump_events().await,
Yan Ru Pei's avatar
Yan Ru Pei committed
205
            Indexer::Concurrent(tpi) => tpi.dump_events().await,
206
207
208
209
210
            Indexer::None => {
                panic!(
                    "Cannot dump events: indexer does not exist (is overlap_score_weight set to 0?)"
                );
            }
211
212
        }
    }
213

214
    pub(crate) async fn process_routing_decision_for_request(
215
        &self,
216
        tokens_with_hashes: &mut TokensWithHashes,
217
218
219
220
221
        worker: WorkerWithDpRank,
    ) -> Result<(), KvRouterError> {
        match self {
            Indexer::KvIndexer(indexer) => {
                indexer
222
                    .process_routing_decision_for_request(tokens_with_hashes, worker)
223
224
                    .await
            }
Yan Ru Pei's avatar
Yan Ru Pei committed
225
226
227
228
            Indexer::Concurrent(tpi) => {
                tpi.process_routing_decision_for_request(tokens_with_hashes, worker)
                    .await
            }
229
230
231
            Indexer::None => Ok(()),
        }
    }
Yan Ru Pei's avatar
Yan Ru Pei committed
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273

    pub(crate) async fn apply_event(&self, event: RouterEvent) {
        match self {
            Indexer::KvIndexer(indexer) => {
                if let Err(e) = indexer.event_sender().send(event).await {
                    tracing::warn!("Failed to send event to indexer: {e}");
                }
            }
            Indexer::Concurrent(tpi) => tpi.apply_event(event).await,
            Indexer::None => {}
        }
    }

    pub(crate) async fn remove_worker(&self, worker_id: WorkerId) {
        match self {
            Indexer::KvIndexer(indexer) => {
                if let Err(e) = indexer.remove_worker_sender().send(worker_id).await {
                    tracing::warn!("Failed to send worker removal for {worker_id}: {e}");
                }
            }
            Indexer::Concurrent(tpi) => {
                KvIndexerInterface::remove_worker(tpi.as_ref(), worker_id).await;
            }
            Indexer::None => {}
        }
    }

    pub(crate) async fn get_workers(&self) -> Vec<WorkerId> {
        match self {
            Indexer::KvIndexer(indexer) => {
                let (resp_tx, resp_rx) = oneshot::channel();
                let req = GetWorkersRequest { resp: resp_tx };
                if let Err(e) = indexer.get_workers_sender().send(req).await {
                    tracing::warn!("Failed to send get_workers request: {e}");
                    return Vec::new();
                }
                resp_rx.await.unwrap_or_default()
            }
            Indexer::Concurrent(tpi) => tpi.backend().get_workers(),
            Indexer::None => Vec::new(),
        }
    }
274
275
}

276
277
/// A KvRouter only decides which worker you should use. It doesn't send you there.
/// TODO: Rename this to indicate it only selects a worker, it does not route.
278
pub struct KvRouter {
279
    indexer: Indexer,
280
    scheduler: KvScheduler,
281
    block_size: u32,
282
    kv_router_config: KvRouterConfig,
Yan Ru Pei's avatar
Yan Ru Pei committed
283
    cancellation_token: tokio_util::sync::CancellationToken,
284
    client: Client,
285
286
287
288
}

impl KvRouter {
    pub async fn new(
289
290
        endpoint: Endpoint,
        client: Client,
291
        mut workers_with_configs: RuntimeConfigWatch,
292
        block_size: u32,
293
        selector: Option<Box<WorkerSelector>>,
294
        kv_router_config: Option<KvRouterConfig>,
295
        worker_type: &'static str,
296
    ) -> Result<Self> {
297
        let kv_router_config = kv_router_config.unwrap_or_default();
298
        kv_router_config.validate()?;
299
        let component = endpoint.component();
300
        let cancellation_token = component.drt().primary_token();
301

302
        let indexer = Indexer::new(component, &kv_router_config, block_size);
303

304
        // Wait for at least one worker with a known runtime config before starting scheduler
305
306
307
308
309
310
        let _ = workers_with_configs
            .wait_for(|m| !m.is_empty())
            .await
            .map_err(|_| {
                anyhow::anyhow!("runtime config watch closed before any workers appeared")
            })?;
311

312
        let scheduler = KvScheduler::start(
313
            component.clone(),
314
            block_size,
315
            workers_with_configs.clone(),
316
            selector,
317
            &kv_router_config,
318
            worker_type,
319
320
        )
        .await?;
321

322
323
        // Start KV event subscription if needed (use_kv_events=true and overlap_score_weight>0)
        if kv_router_config.should_subscribe_to_kv_events() {
324
325
            subscriber::start_subscriber(component.clone(), &kv_router_config, indexer.clone())
                .await?;
326
        } else {
327
            tracing::info!(
328
329
330
                "Skipping KV event subscription (use_kv_events={}, overlap_score_weight={})",
                kv_router_config.use_kv_events,
                kv_router_config.overlap_score_weight,
331
            );
332
        }
333

334
        tracing::info!("KV Routing initialized");
335
        Ok(Self {
336
            indexer,
337
            scheduler,
338
            block_size,
339
            kv_router_config,
Yan Ru Pei's avatar
Yan Ru Pei committed
340
            cancellation_token,
341
            client,
342
        })
343
344
    }

345
346
347
348
349
    /// Get a reference to the client used by this KvRouter
    pub fn client(&self) -> &Client {
        &self.client
    }

350
351
352
353
354
355
356
357
    pub fn indexer(&self) -> &Indexer {
        &self.indexer
    }

    pub fn kv_router_config(&self) -> &KvRouterConfig {
        &self.kv_router_config
    }

358
    /// Give these tokens, find the worker with the best match in it's KV cache.
Yan Ru Pei's avatar
Yan Ru Pei committed
359
    /// Returns the best worker (with dp_rank) and overlap amount in number of blocks.
360
361
362
    /// Now also takes optional context_id for request tracking.
    ///
    /// When `allowed_worker_ids` is Some, only workers in that set are considered for selection.
363
    #[allow(clippy::too_many_arguments)]
Yan Ru Pei's avatar
Yan Ru Pei committed
364
    pub async fn find_best_match(
365
        &self,
Yan Ru Pei's avatar
Yan Ru Pei committed
366
        context_id: Option<&str>,
367
        tokens: &[u32],
368
        block_mm_infos: Option<&[Option<BlockExtraInfo>]>,
369
        router_config_override: Option<&RouterConfigOverride>,
370
        update_states: bool,
371
        lora_name: Option<String>,
372
        priority_jump: f64,
373
        allowed_worker_ids: Option<HashSet<WorkerId>>,
Yan Ru Pei's avatar
Yan Ru Pei committed
374
    ) -> anyhow::Result<(WorkerWithDpRank, u32)> {
375
376
        let start = Instant::now();

Yan Ru Pei's avatar
Yan Ru Pei committed
377
        if update_states && context_id.is_none() {
378
            anyhow::bail!("context_id must be provided when update_states is true");
Yan Ru Pei's avatar
Yan Ru Pei committed
379
380
        }

381
        let isl_tokens = tokens.len();
382

383
384
385
386
387
388
389
390
        let block_hashes = tracing::info_span!("kv_router.compute_block_hashes").in_scope(|| {
            compute_block_hash_for_seq(
                tokens,
                self.block_size,
                block_mm_infos,
                lora_name.as_deref(),
            )
        });
391
        let hash_elapsed = start.elapsed();
392

393
        let overlap_scores = self
394
395
396
397
            .indexer
            .find_matches(block_hashes)
            .instrument(tracing::info_span!("kv_router.find_matches"))
            .await?;
398
        let find_matches_elapsed = start.elapsed();
399

400
        // Compute seq_hashes only if scheduler needs it for active blocks tracking
401
402
403
404
405
        let maybe_seq_hashes = tracing::info_span!("kv_router.compute_seq_hashes").in_scope(|| {
            self.kv_router_config.compute_seq_hashes_for_tracking(
                tokens,
                self.block_size,
                router_config_override,
406
                lora_name.as_deref(),
407
408
            )
        });
409
        let seq_hash_elapsed = start.elapsed();
410

411
        let response = self
412
            .scheduler
413
            .schedule(
Yan Ru Pei's avatar
Yan Ru Pei committed
414
                context_id.map(|s| s.to_string()),
415
                isl_tokens,
416
                maybe_seq_hashes,
417
                overlap_scores,
418
                router_config_override,
419
                update_states,
420
                lora_name,
421
                priority_jump,
422
                allowed_worker_ids,
423
            )
424
            .instrument(tracing::info_span!("kv_router.schedule"))
425
            .await?;
426
427
        let total_elapsed = start.elapsed();

428
429
430
431
432
433
434
435
        if let Some(m) = metrics::RoutingOverheadMetrics::get() {
            m.observe(
                hash_elapsed,
                find_matches_elapsed,
                seq_hash_elapsed,
                total_elapsed,
            );
        }
436

437
        #[cfg(feature = "bench")]
438
439
440
441
442
443
444
445
446
        tracing::info!(
            isl_tokens,
            hash_us = hash_elapsed.as_micros() as u64,
            find_matches_us = (find_matches_elapsed - hash_elapsed).as_micros() as u64,
            seq_hash_us = (seq_hash_elapsed - find_matches_elapsed).as_micros() as u64,
            schedule_us = (total_elapsed - seq_hash_elapsed).as_micros() as u64,
            total_us = total_elapsed.as_micros() as u64,
            "find_best_match completed"
        );
447

448
        Ok((response.best_worker, response.overlap_blocks))
449
450
    }

451
    #[allow(clippy::too_many_arguments)]
452
453
454
455
456
    pub async fn add_request(
        &self,
        request_id: String,
        tokens: &[u32],
        overlap_blocks: u32,
457
        expected_output_tokens: Option<u32>,
Yan Ru Pei's avatar
Yan Ru Pei committed
458
        worker: WorkerWithDpRank,
459
        lora_name: Option<String>,
460
        router_config_override: Option<&RouterConfigOverride>,
461
462
    ) {
        let isl_tokens = tokens.len();
463

464
465
466
467
        let maybe_seq_hashes = self.kv_router_config.compute_seq_hashes_for_tracking(
            tokens,
            self.block_size,
            router_config_override,
468
            lora_name.as_deref(),
469
        );
470

471
472
        if let Err(e) = self
            .scheduler
473
474
475
476
477
            .add_request(SequenceRequest {
                request_id: request_id.clone(),
                token_sequence: maybe_seq_hashes,
                isl: isl_tokens,
                overlap: overlap_blocks,
478
                expected_output_tokens,
Yan Ru Pei's avatar
Yan Ru Pei committed
479
                worker,
480
                lora_name,
481
            })
482
483
484
485
            .await
        {
            tracing::warn!("Failed to add request {request_id}: {e}");
        }
486
487
    }

488
    pub async fn mark_prefill_completed(&self, request_id: &str) -> Result<(), SequenceError> {
489
        self.scheduler.mark_prefill_completed(request_id).await
490
491
    }

492
    pub async fn free(&self, request_id: &str) -> Result<(), SequenceError> {
493
        self.scheduler.free(request_id).await
494
    }
495

496
497
498
499
500
501
    /// Get the worker type for this router ("prefill" or "decode").
    /// Used for Prometheus metric labeling.
    pub fn worker_type(&self) -> &'static str {
        self.scheduler.worker_type()
    }

502
    pub fn add_output_block(
503
504
505
506
        &self,
        request_id: &str,
        decay_fraction: Option<f64>,
    ) -> Result<(), SequenceError> {
507
        self.scheduler.add_output_block(request_id, decay_fraction)
508
509
    }

510
    pub fn block_size(&self) -> u32 {
511
512
        self.block_size
    }
513

514
515
516
517
518
519
    /// Compute the overlap blocks for a given token sequence and worker.
    /// This queries the indexer to find how many blocks are already cached.
    pub async fn get_overlap_blocks(
        &self,
        tokens: &[u32],
        worker: WorkerWithDpRank,
520
        lora_name: Option<&str>,
521
    ) -> Result<u32, KvRouterError> {
522
        let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, None, lora_name);
523
524
525
526
        let overlap_scores = self.indexer.find_matches(block_hashes).await?;
        Ok(overlap_scores.scores.get(&worker).copied().unwrap_or(0))
    }

527
    /// Get potential prefill and decode loads for all workers
528
529
530
531
    pub async fn get_potential_loads(
        &self,
        tokens: &[u32],
        router_config_override: Option<&RouterConfigOverride>,
532
        lora_name: Option<&str>,
533
    ) -> Result<Vec<PotentialLoad>> {
534
        let isl_tokens = tokens.len();
535
        let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, None, lora_name);
536
        let overlap_scores = self.indexer.find_matches(block_hashes.clone()).await?;
537

538
539
540
541
        let maybe_seq_hashes = self.kv_router_config.compute_seq_hashes_for_tracking(
            tokens,
            self.block_size,
            router_config_override,
542
            lora_name,
543
        );
544

545
546
        Ok(self
            .scheduler
547
            .get_potential_loads(maybe_seq_hashes, isl_tokens, overlap_scores))
548
549
    }

550
551
552
553
    /// Dump all events from the indexer
    pub async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
        self.indexer.dump_events().await
    }
554
555
}

Michael Feil's avatar
Michael Feil committed
556
557
// NOTE: KVRouter works like a PushRouter,
// but without the reverse proxy functionality, but based on contract of 3 request types
558
559
560
561
562
563
564
#[async_trait]
impl AsyncEngine<SingleIn<RouterRequest>, ManyOut<Annotated<RouterResponse>>, Error> for KvRouter {
    async fn generate(
        &self,
        request: SingleIn<RouterRequest>,
    ) -> Result<ManyOut<Annotated<RouterResponse>>> {
        let (request, ctx) = request.into_parts();
Michael Feil's avatar
Michael Feil committed
565
566
567
        let context_id = ctx.context().id().to_string();
        // Handle different request types
        let response = match request {
568
569
570
571
            RouterRequest::New {
                tokens,
                block_mm_infos,
            } => {
Yan Ru Pei's avatar
Yan Ru Pei committed
572
                let (best_worker, overlap_blocks) = self
573
574
575
576
577
578
579
580
                    .find_best_match(
                        Some(&context_id),
                        &tokens,
                        block_mm_infos.as_deref(),
                        None,
                        true,
                        None,
                        0.0,
581
                        None,
582
                    )
Michael Feil's avatar
Michael Feil committed
583
584
585
                    .await?;

                RouterResponse::New {
Yan Ru Pei's avatar
Yan Ru Pei committed
586
587
                    worker_id: best_worker.worker_id,
                    dp_rank: best_worker.dp_rank,
Michael Feil's avatar
Michael Feil committed
588
589
590
                    overlap_blocks,
                }
            }
591
592
593
594
595
596
            RouterRequest::MarkPrefill => RouterResponse::PrefillMarked {
                success: self.mark_prefill_completed(&context_id).await.is_ok(),
            },
            RouterRequest::MarkFree => RouterResponse::FreeMarked {
                success: self.free(&context_id).await.is_ok(),
            },
Michael Feil's avatar
Michael Feil committed
597
        };
598
599
600
601
602
603

        let response = Annotated::from_data(response);
        let stream = stream::iter(vec![response]);
        Ok(ResponseStream::new(Box::pin(stream), ctx.context()))
    }
}
604

Yan Ru Pei's avatar
Yan Ru Pei committed
605
606
607
608
609
610
impl Drop for KvRouter {
    fn drop(&mut self) {
        tracing::info!("Dropping KvRouter - cancelling background tasks");
        self.cancellation_token.cancel();
    }
}