kv_router.rs 20.6 KB
Newer Older
1
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
3
// SPDX-License-Identifier: Apache-2.0

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

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

26
27
28
29
30
// 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;

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

43
pub use config::{KvRouterConfig, RouterConfigOverride};
44
pub use prefill_router::PrefillRouter;
45
pub use push_router::{DirectRoutingRouter, KvPushRouter};
46

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

63
64
// [gluo TODO] shouldn't need to be public
// this should be discovered from the component
65
66
67
68
69

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

// for metric publishing (push-based)
70
pub const KV_EVENT_SUBJECT: &str = "kv-events";
71
72
73
74
75
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";
76

77
78
79
80
// for radix tree snapshot storage
pub const RADIX_STATE_BUCKET: &str = "radix-bucket";
pub const RADIX_STATE_FILE: &str = "radix-state";

81
82
83
// for worker-local kvindexer query
pub const WORKER_KV_INDEXER_BUFFER_SIZE: usize = 1024; // store 1024 most recent events in worker buffer

84
85
86
87
88
89
/// 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}")
}

90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
// for router discovery registration
pub const KV_ROUTER_COMPONENT: &str = "kv-router";
pub const KV_ROUTER_ENDPOINT: &str = "generate";

/// Creates an EndpointId for the KV router in the given namespace.
pub fn router_endpoint_id(namespace: String) -> EndpointId {
    EndpointId {
        namespace,
        component: KV_ROUTER_COMPONENT.to_string(),
        name: KV_ROUTER_ENDPOINT.to_string(),
    }
}

/// Creates a DiscoveryQuery for the KV router in the given namespace.
pub fn router_discovery_query(namespace: String) -> DiscoveryQuery {
    DiscoveryQuery::Endpoint {
        namespace,
        component: KV_ROUTER_COMPONENT.to_string(),
        endpoint: KV_ROUTER_ENDPOINT.to_string(),
    }
}

112
113
114
115
/// A trait that users can implement to define custom selection logic
pub trait WorkerSelector {
    fn select_worker(
        &self,
116
        workers: &HashMap<protocols::WorkerId, ModelRuntimeConfig>,
117
        request: &SchedulingRequest,
118
        block_size: u32,
119
120
    ) -> Result<WorkerSelectionResult, KvSchedulerError>;
}
121

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

Yan Ru Pei's avatar
Yan Ru Pei committed
129
130
131
132
133
    /// 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>>),

134
135
136
    /// 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,
137
138
139
}

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

        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,
154
                block_size,
Yan Ru Pei's avatar
Yan Ru Pei committed
155
            )));
156
        }
Yan Ru Pei's avatar
Yan Ru Pei committed
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177

        let kv_indexer_metrics = indexer::KvIndexerMetrics::from_component(component);

        // If use_kv_events is false, enable TTL and pruning for approximate behavior
        let prune_config = if !kv_router_config.use_kv_events {
            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,
            })
        } else {
            None
        };

        Indexer::KvIndexer(KvIndexer::new_with_frequency(
            cancellation_token,
            None, // expiration_duration for frequency tracking
            block_size,
            kv_indexer_metrics,
            prune_config,
        ))
178
179
180
    }

    pub(crate) async fn find_matches(
181
182
183
184
185
        &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
186
            Indexer::Concurrent(tpi) => tpi.find_matches(sequence).await,
187
188
189
            Indexer::None => Ok(OverlapScores {
                scores: HashMap::new(),
                frequencies: Vec::new(),
190
                tree_sizes: HashMap::new(),
191
            }),
192
193
        }
    }
194

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

207
    pub(crate) async fn process_routing_decision_for_request(
208
        &self,
209
        tokens_with_hashes: &mut TokensWithHashes,
210
211
212
213
214
        worker: WorkerWithDpRank,
    ) -> Result<(), KvRouterError> {
        match self {
            Indexer::KvIndexer(indexer) => {
                indexer
215
                    .process_routing_decision_for_request(tokens_with_hashes, worker)
216
217
                    .await
            }
Yan Ru Pei's avatar
Yan Ru Pei committed
218
219
220
221
            Indexer::Concurrent(tpi) => {
                tpi.process_routing_decision_for_request(tokens_with_hashes, worker)
                    .await
            }
222
223
224
            Indexer::None => Ok(()),
        }
    }
Yan Ru Pei's avatar
Yan Ru Pei committed
225
226
227
228
229
230
231
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

    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(),
        }
    }
267
268
}

269
270
/// 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.
271
pub struct KvRouter {
272
    indexer: Indexer,
273
    scheduler: KvScheduler,
274
    block_size: u32,
275
    kv_router_config: KvRouterConfig,
Yan Ru Pei's avatar
Yan Ru Pei committed
276
    cancellation_token: tokio_util::sync::CancellationToken,
277
    client: Client,
278
279
280
}

impl KvRouter {
281
    #[allow(clippy::too_many_arguments)]
282
    pub async fn new(
283
284
        endpoint: Endpoint,
        client: Client,
285
        mut workers_with_configs: RuntimeConfigWatch,
286
        block_size: u32,
287
        selector: Option<Box<dyn WorkerSelector + Send + Sync>>,
288
        kv_router_config: Option<KvRouterConfig>,
289
        router_id: u64,
290
        worker_type: &'static str,
291
    ) -> Result<Self> {
292
        let kv_router_config = kv_router_config.unwrap_or_default();
293
        kv_router_config.validate()?;
294
        let component = endpoint.component();
295
        let cancellation_token = component.drt().primary_token();
296

297
298
299
300
301
302
        let indexer = Indexer::new(
            component,
            &kv_router_config,
            block_size,
            cancellation_token.clone(),
        );
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.router_replica_sync,
318
            router_id,
319
            worker_type,
320
            kv_router_config.router_queue_threshold,
321
322
        )
        .await?;
323

324
325
326
327
328
329
        // Start KV event subscription if needed (use_kv_events=true and overlap_score_weight>0)
        if kv_router_config.should_subscribe_to_kv_events() {
            subscriber::start_subscriber(
                component.clone(),
                &kv_router_config,
                router_id,
Yan Ru Pei's avatar
Yan Ru Pei committed
330
                indexer.clone(),
331
332
333
334
                cancellation_token.clone(),
            )
            .await?;
        } else {
335
            tracing::info!(
336
337
338
                "Skipping KV event subscription (use_kv_events={}, overlap_score_weight={})",
                kv_router_config.use_kv_events,
                kv_router_config.overlap_score_weight,
339
            );
340
        }
341

342
        tracing::info!("KV Routing initialized");
343
        Ok(Self {
344
            indexer,
345
            scheduler,
346
            block_size,
347
            kv_router_config,
Yan Ru Pei's avatar
Yan Ru Pei committed
348
            cancellation_token,
349
            client,
350
        })
351
352
    }

353
354
355
356
357
    /// Get a reference to the client used by this KvRouter
    pub fn client(&self) -> &Client {
        &self.client
    }

358
359
360
361
362
363
364
365
    pub fn indexer(&self) -> &Indexer {
        &self.indexer
    }

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

366
    /// Give these tokens, find the worker with the best match in it's KV cache.
Yan Ru Pei's avatar
Yan Ru Pei committed
367
    /// Returns the best worker (with dp_rank) and overlap amount in number of blocks.
Yan Ru Pei's avatar
Yan Ru Pei committed
368
    /// Now also takes optional context_id for request tracking
369
    #[allow(clippy::too_many_arguments)]
Yan Ru Pei's avatar
Yan Ru Pei committed
370
    pub async fn find_best_match(
371
        &self,
Yan Ru Pei's avatar
Yan Ru Pei committed
372
        context_id: Option<&str>,
373
        tokens: &[u32],
374
        block_mm_infos: Option<&[Option<BlockExtraInfo>]>,
375
        router_config_override: Option<&RouterConfigOverride>,
376
        update_states: bool,
377
        lora_name: Option<String>,
378
        priority_jump: f64,
Yan Ru Pei's avatar
Yan Ru Pei committed
379
    ) -> anyhow::Result<(WorkerWithDpRank, u32)> {
380
381
        let start = Instant::now();

Yan Ru Pei's avatar
Yan Ru Pei committed
382
        if update_states && context_id.is_none() {
383
            anyhow::bail!("context_id must be provided when update_states is true");
Yan Ru Pei's avatar
Yan Ru Pei committed
384
385
        }

386
        let isl_tokens = tokens.len();
387

388
        let block_hashes = tracing::info_span!("kv_router.compute_block_hashes")
389
            .in_scope(|| compute_block_hash_for_seq(tokens, self.block_size, block_mm_infos));
390
        let hash_elapsed = start.elapsed();
391

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

399
        // Compute seq_hashes only if scheduler needs it for active blocks tracking
400
401
402
403
404
405
406
        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,
            )
        });
407
        let seq_hash_elapsed = start.elapsed();
408

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

        metrics::ROUTING_OVERHEAD_METRICS.observe(
            hash_elapsed,
            find_matches_elapsed,
            seq_hash_elapsed,
            total_elapsed,
        );
431

432
        #[cfg(feature = "bench")]
433
434
435
436
437
438
439
440
441
        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"
        );
442

443
444
        // Note: Routing decision recording (for approximate mode) is now handled
        // by KvPushRouter::generate after select_worker returns.
445

446
447
        let overlap_amount = overlap_scores
            .scores
Yan Ru Pei's avatar
Yan Ru Pei committed
448
            .get(&best_worker)
449
450
            .copied()
            .unwrap_or(0);
Yan Ru Pei's avatar
Yan Ru Pei committed
451
        Ok((best_worker, overlap_amount))
452
453
    }

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

467
468
469
470
471
        let maybe_seq_hashes = self.kv_router_config.compute_seq_hashes_for_tracking(
            tokens,
            self.block_size,
            router_config_override,
        );
472

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

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

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

498
499
500
501
502
503
    /// 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()
    }

504
505
506
507
508
509
510
511
512
513
    pub async fn add_output_block(
        &self,
        request_id: &str,
        decay_fraction: Option<f64>,
    ) -> Result<(), SequenceError> {
        self.scheduler
            .add_output_block(request_id, decay_fraction)
            .await
    }

514
    pub fn block_size(&self) -> u32 {
515
516
        self.block_size
    }
517

518
519
520
521
522
523
524
525
526
527
528
529
    /// 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,
    ) -> Result<u32, KvRouterError> {
        let block_hashes = compute_block_hash_for_seq(tokens, self.block_size, None);
        let overlap_scores = self.indexer.find_matches(block_hashes).await?;
        Ok(overlap_scores.scores.get(&worker).copied().unwrap_or(0))
    }

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

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

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

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

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

                RouterResponse::New {
Yan Ru Pei's avatar
Yan Ru Pei committed
587
588
                    worker_id: best_worker.worker_id,
                    dp_rank: best_worker.dp_rank,
Michael Feil's avatar
Michael Feil committed
589
590
591
                    overlap_blocks,
                }
            }
592
593
594
595
596
597
            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
598
        };
599
600
601
602
603
604

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

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