dynamo_runtime/discovery/
kv_store.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::HashMap;
5use std::pin::Pin;
6use std::sync::Arc;
7
8use anyhow::Result;
9use async_trait::async_trait;
10use futures::{Stream, StreamExt};
11use tokio_util::sync::CancellationToken;
12
13use super::{
14    Discovery, DiscoveryEvent, DiscoveryInstance, DiscoveryInstanceId, DiscoveryQuery,
15    DiscoverySpec, DiscoveryStream, EndpointInstanceId, EventChannelInstanceId, EventScope,
16    EventSourceInstanceId, ModelCardInstanceId, encode_event_segment,
17    validate_event_source_reregistration,
18};
19use crate::storage::kv;
20
21const INSTANCES_BUCKET: &str = "v1/instances";
22const MODELS_BUCKET: &str = "v1/mdc";
23const EVENT_CHANNELS_BUCKET: &str = "v1/event_channels";
24const EVENT_SOURCES_BUCKET: &str = "v1/event_sources";
25
26/// Discovery implementation backed by a kv::Store
27pub struct KVStoreDiscovery {
28    store: Arc<kv::Manager>,
29    cancel_token: CancellationToken,
30}
31
32impl KVStoreDiscovery {
33    pub fn new(store: kv::Manager, cancel_token: CancellationToken) -> Self {
34        Self {
35            store: Arc::new(store),
36            cancel_token,
37        }
38    }
39
40    /// Build the key path for an endpoint (relative to bucket, not absolute)
41    fn endpoint_key(instance: &crate::component::Instance) -> String {
42        instance.endpoint_instance_id().to_path()
43    }
44
45    /// Build the key path for a model (relative to bucket, not absolute)
46    fn model_key(namespace: &str, component: &str, endpoint: &str, instance_id: u64) -> String {
47        format!("{}/{}/{}/{:x}", namespace, component, endpoint, instance_id)
48    }
49
50    /// Build the key path for an event channel relative to bucket, not absolute)
51    fn event_channel_key(scope: &EventScope, topic: &str, instance_id: u64) -> String {
52        format!(
53            "{}/topic/{}/{:x}",
54            scope.path_prefix(),
55            encode_event_segment(topic),
56            instance_id
57        )
58    }
59
60    /// Build the key path for an event source relative to its bucket.
61    fn event_source_key(scope: &EventScope, topic: &str, publisher_id: u64) -> String {
62        EventSourceInstanceId {
63            scope: scope.clone(),
64            topic: topic.to_string(),
65            publisher_id,
66        }
67        .to_path()
68    }
69
70    /// Extract prefix for querying based on discovery query
71    fn query_prefix(query: &DiscoveryQuery) -> String {
72        match query {
73            DiscoveryQuery::AllEndpoints => INSTANCES_BUCKET.to_string(),
74            DiscoveryQuery::NamespacedEndpoints { namespace } => {
75                format!("{}/{}", INSTANCES_BUCKET, namespace)
76            }
77            DiscoveryQuery::ComponentEndpoints {
78                namespace,
79                component,
80            } => {
81                format!("{}/{}/{}", INSTANCES_BUCKET, namespace, component)
82            }
83            DiscoveryQuery::Endpoint {
84                namespace,
85                component,
86                endpoint,
87            } => {
88                format!(
89                    "{}/{}/{}/{}",
90                    INSTANCES_BUCKET, namespace, component, endpoint
91                )
92            }
93            DiscoveryQuery::AllModels => MODELS_BUCKET.to_string(),
94            DiscoveryQuery::NamespacedModels { namespace } => {
95                format!("{}/{}", MODELS_BUCKET, namespace)
96            }
97            DiscoveryQuery::ComponentModels {
98                namespace,
99                component,
100            } => {
101                format!("{}/{}/{}", MODELS_BUCKET, namespace, component)
102            }
103            DiscoveryQuery::EndpointModels {
104                namespace,
105                component,
106                endpoint,
107            } => {
108                format!("{}/{}/{}/{}", MODELS_BUCKET, namespace, component, endpoint)
109            }
110            DiscoveryQuery::EventChannels(query) => {
111                let mut path = EVENT_CHANNELS_BUCKET.to_string();
112                if let Some(scope) = &query.scope {
113                    path.push('/');
114                    path.push_str(&scope.path_prefix());
115                    if let Some(topic) = &query.topic {
116                        path.push_str("/topic/");
117                        path.push_str(&encode_event_segment(topic));
118                    }
119                }
120                path
121            }
122            DiscoveryQuery::EventSources(query) => {
123                let mut path = EVENT_SOURCES_BUCKET.to_string();
124                if let Some(scope) = &query.scope {
125                    path.push('/');
126                    path.push_str(&scope.path_prefix());
127                    if let Some(topic) = &query.topic {
128                        path.push_str("/topic/");
129                        path.push_str(&encode_event_segment(topic));
130                    }
131                }
132                path
133            }
134        }
135    }
136
137    /// Strip bucket prefix from a key if present, returning the relative path within the bucket
138    /// For example: "v1/instances/ns/comp/ep" -> "ns/comp/ep"
139    /// Or if already relative: "ns/comp/ep" -> "ns/comp/ep"
140    fn strip_bucket_prefix<'a>(key: &'a str, bucket_name: &str) -> &'a str {
141        // Try to strip "bucket_name/" from the beginning
142        if let Some(stripped) = key.strip_prefix(bucket_name) {
143            // Strip the leading slash if present
144            stripped.strip_prefix('/').unwrap_or(stripped)
145        } else {
146            // Key is already relative to bucket
147            key
148        }
149    }
150
151    /// Check if a key matches the given prefix, handling both absolute and relative key formats
152    /// This works regardless of whether keys include the bucket prefix (etcd) or not (memory)
153    fn matches_prefix(key_str: &str, prefix: &str, bucket_name: &str) -> bool {
154        // Normalize both the key and prefix to relative paths (without bucket prefix)
155        let relative_key = Self::strip_bucket_prefix(key_str, bucket_name);
156        let relative_prefix = Self::strip_bucket_prefix(prefix, bucket_name);
157
158        // Empty prefix matches everything in the bucket
159        if relative_prefix.is_empty() {
160            return true;
161        }
162
163        relative_key == relative_prefix
164            || relative_key
165                .strip_prefix(relative_prefix)
166                .is_some_and(|suffix| suffix.starts_with('/'))
167    }
168
169    fn bucket_for_prefix(prefix: &str) -> &'static str {
170        if prefix == INSTANCES_BUCKET
171            || prefix
172                .strip_prefix(INSTANCES_BUCKET)
173                .is_some_and(|suffix| suffix.starts_with('/'))
174        {
175            INSTANCES_BUCKET
176        } else if prefix == EVENT_CHANNELS_BUCKET
177            || prefix
178                .strip_prefix(EVENT_CHANNELS_BUCKET)
179                .is_some_and(|suffix| suffix.starts_with('/'))
180        {
181            EVENT_CHANNELS_BUCKET
182        } else if prefix == EVENT_SOURCES_BUCKET
183            || prefix
184                .strip_prefix(EVENT_SOURCES_BUCKET)
185                .is_some_and(|suffix| suffix.starts_with('/'))
186        {
187            EVENT_SOURCES_BUCKET
188        } else {
189            MODELS_BUCKET
190        }
191    }
192
193    /// Parse and deserialize a discovery instance from KV store entry
194    fn parse_instance(value: &[u8]) -> Result<DiscoveryInstance> {
195        let instance: DiscoveryInstance = serde_json::from_slice(value)?;
196        Ok(instance)
197    }
198
199    fn parse_instance_id_from_key(key_str: &str, bucket_name: &str) -> Option<DiscoveryInstanceId> {
200        let relative_key = Self::strip_bucket_prefix(key_str, bucket_name);
201        let parsed = match bucket_name {
202            INSTANCES_BUCKET => {
203                EndpointInstanceId::from_path(relative_key).map(DiscoveryInstanceId::Endpoint)
204            }
205            MODELS_BUCKET => {
206                ModelCardInstanceId::from_path(relative_key).map(DiscoveryInstanceId::Model)
207            }
208            EVENT_CHANNELS_BUCKET => EventChannelInstanceId::from_path(relative_key)
209                .map(DiscoveryInstanceId::EventChannel),
210            EVENT_SOURCES_BUCKET => {
211                EventSourceInstanceId::from_path(relative_key).map(DiscoveryInstanceId::EventSource)
212            }
213            _ => {
214                tracing::warn!(
215                    key = %key_str,
216                    bucket = bucket_name,
217                    "Unknown discovery bucket for delete/resync key"
218                );
219                return None;
220            }
221        };
222
223        parsed
224            .inspect_err(|err| {
225                tracing::warn!(
226                    key = %key_str,
227                    relative_key = %relative_key,
228                    bucket = bucket_name,
229                    error = %err,
230                    "Failed to parse discovery instance id from key"
231                );
232            })
233            .ok()
234    }
235
236    fn discovery_events_from_watch_event(
237        event: kv::WatchEvent,
238        prefix: &str,
239        bucket_name: &str,
240        known_instances: &mut HashMap<DiscoveryInstanceId, DiscoveryInstance>,
241    ) -> Vec<DiscoveryEvent> {
242        match event {
243            kv::WatchEvent::Put(kv) => {
244                if !Self::matches_prefix(kv.key_str(), prefix, bucket_name) {
245                    return vec![];
246                }
247
248                match Self::parse_instance(kv.value()) {
249                    Ok(instance) => {
250                        known_instances.insert(instance.id(), instance.clone());
251                        vec![DiscoveryEvent::Added(instance)]
252                    }
253                    Err(e) => {
254                        tracing::warn!(
255                            key = %kv.key_str(),
256                            error = %e,
257                            "Failed to parse discovery instance from watch event"
258                        );
259                        vec![]
260                    }
261                }
262            }
263            kv::WatchEvent::Delete(kv) => {
264                let key_str = kv.as_ref();
265                if !Self::matches_prefix(key_str, prefix, bucket_name) {
266                    return vec![];
267                }
268
269                let Some(id) = Self::parse_instance_id_from_key(key_str, bucket_name) else {
270                    return vec![];
271                };
272
273                known_instances.remove(&id);
274                tracing::debug!(
275                    "KVStoreDiscovery::list_and_watch: Emitting Removed event for {:?}, key={}",
276                    id,
277                    key_str
278                );
279                vec![DiscoveryEvent::Removed(id)]
280            }
281            kv::WatchEvent::Resync(snapshot) => {
282                let mut next_instances = HashMap::<DiscoveryInstanceId, DiscoveryInstance>::new();
283
284                for (key, value) in snapshot {
285                    let key_str = key.as_ref();
286                    if !Self::matches_prefix(key_str, prefix, bucket_name) {
287                        continue;
288                    }
289
290                    match Self::parse_instance(value.as_ref()) {
291                        Ok(instance) => {
292                            next_instances.insert(instance.id(), instance);
293                        }
294                        Err(e) => {
295                            tracing::warn!(
296                                key = %key_str,
297                                error = %e,
298                                "Failed to parse discovery instance from resync event"
299                            );
300                            // The key is still present in the authoritative snapshot; keep
301                            // the previous value if only this local parse failed.
302                            if let Some(id) = Self::parse_instance_id_from_key(key_str, bucket_name)
303                                && let Some(existing) = known_instances.get(&id)
304                            {
305                                next_instances.insert(id, existing.clone());
306                            }
307                        }
308                    }
309                }
310
311                let mut events = Vec::new();
312                for id in known_instances.keys() {
313                    if !next_instances.contains_key(id) {
314                        events.push(DiscoveryEvent::Removed(id.clone()));
315                    }
316                }
317
318                for (id, instance) in &next_instances {
319                    if known_instances.get(id) != Some(instance) {
320                        // Added is an upsert event here: a resync can discover
321                        // either a new instance or changed data for an existing id.
322                        events.push(DiscoveryEvent::Added(instance.clone()));
323                    }
324                }
325
326                tracing::warn!(
327                    old_count = known_instances.len(),
328                    new_count = next_instances.len(),
329                    emitted_events = events.len(),
330                    "KVStoreDiscovery::list_and_watch resynced discovery state"
331                );
332
333                *known_instances = next_instances;
334                events
335            }
336        }
337    }
338}
339
340#[async_trait]
341impl Discovery for KVStoreDiscovery {
342    fn instance_id(&self) -> u64 {
343        self.store.connection_id()
344    }
345
346    async fn register_internal(&self, spec: DiscoverySpec) -> Result<DiscoveryInstance> {
347        let instance = spec.into_instance(self.instance_id());
348        let instance_id = instance.instance_id();
349        let is_event_source = matches!(&instance, DiscoveryInstance::EventSource { .. });
350
351        let (bucket_name, key_path) = match &instance {
352            DiscoveryInstance::Endpoint(inst) => {
353                let key = Self::endpoint_key(inst);
354                tracing::debug!(
355                    "KVStoreDiscovery::register: Registering endpoint instance_id={}, namespace={}, component={}, endpoint={}, key={}",
356                    inst.instance_id,
357                    inst.namespace,
358                    inst.component,
359                    inst.endpoint,
360                    key
361                );
362                (INSTANCES_BUCKET, key)
363            }
364            DiscoveryInstance::Model {
365                namespace,
366                component,
367                endpoint,
368                instance_id,
369                model_suffix,
370                ..
371            } => {
372                let mut key = Self::model_key(namespace, component, endpoint, *instance_id);
373
374                // If there's a model_suffix (e.g., for LoRA adapters), append it after the instance_id
375                // Key format: {namespace}/{component}/{endpoint}/{instance_id:x}/{model_suffix}
376                if let Some(suffix) = model_suffix
377                    && !suffix.is_empty()
378                {
379                    key = format!("{}/{}", key, suffix);
380                    tracing::debug!(
381                        "KVStoreDiscovery::register: Registering LoRA model with suffix={}, instance_id={}, namespace={}, component={}, endpoint={}, key={}",
382                        suffix,
383                        instance_id,
384                        namespace,
385                        component,
386                        endpoint,
387                        key
388                    );
389                }
390
391                // Log for base models (no suffix or empty suffix)
392                if model_suffix.as_ref().is_none_or(|s| s.is_empty()) {
393                    tracing::debug!(
394                        "KVStoreDiscovery::register: Registering base model instance_id={}, namespace={}, component={}, endpoint={}, key={}",
395                        instance_id,
396                        namespace,
397                        component,
398                        endpoint,
399                        key
400                    );
401                }
402                (MODELS_BUCKET, key)
403            }
404            DiscoveryInstance::EventChannel {
405                scope,
406                topic,
407                instance_id,
408                ..
409            } => {
410                let key = Self::event_channel_key(scope, topic, *instance_id);
411                // TODO: bis - remove this info log
412                tracing::info!(
413                    "KVStoreDiscovery::register: EventChannel bucket={}, key={}",
414                    EVENT_CHANNELS_BUCKET,
415                    key
416                );
417                tracing::debug!(
418                    "KVStoreDiscovery::register: Registering event channel instance_id={}, scope={:?}, topic={}, key={}",
419                    instance_id,
420                    scope,
421                    topic,
422                    key
423                );
424                (EVENT_CHANNELS_BUCKET, key)
425            }
426            DiscoveryInstance::EventSource {
427                scope,
428                topic,
429                publisher_id,
430                ..
431            } => {
432                let key = Self::event_source_key(scope, topic, *publisher_id);
433                tracing::debug!(
434                    "KVStoreDiscovery::register: Registering event source publisher_id={}, scope={:?}, topic={}, key={}",
435                    publisher_id,
436                    scope,
437                    topic,
438                    key
439                );
440                (EVENT_SOURCES_BUCKET, key)
441            }
442        };
443
444        // Serialize the instance
445        let instance_json = serde_json::to_vec(&instance)?;
446        tracing::debug!(
447            "KVStoreDiscovery::register: Serialized instance to {} bytes for key={}",
448            instance_json.len(),
449            key_path
450        );
451
452        // Store in the KV store with no TTL (instances persist until explicitly removed)
453        tracing::debug!(
454            "KVStoreDiscovery::register: Getting/creating bucket={} for key={}",
455            bucket_name,
456            key_path
457        );
458        let bucket = self.store.get_or_create_bucket(bucket_name, None).await?;
459        let key = kv::Key::new(key_path.clone());
460
461        if is_event_source && let Some(existing) = bucket.get(&key).await? {
462            let existing: DiscoveryInstance = serde_json::from_slice(existing.as_ref())?;
463            validate_event_source_reregistration(&existing, &instance)?;
464            return Ok(existing);
465        }
466
467        tracing::debug!(
468            "KVStoreDiscovery::register: Inserting into bucket={}, key={}",
469            bucket_name,
470            key_path
471        );
472        // Use revision 0 for initial registration
473        let outcome = match bucket.insert(&key, instance_json.into(), 0).await {
474            Ok(outcome) => outcome,
475            Err(error) if is_event_source => {
476                let Some(existing) = bucket.get(&key).await? else {
477                    return Err(error.into());
478                };
479                let existing: DiscoveryInstance = serde_json::from_slice(existing.as_ref())?;
480                validate_event_source_reregistration(&existing, &instance)?;
481                return Ok(existing);
482            }
483            Err(error) => return Err(error.into()),
484        };
485        tracing::debug!(
486            "KVStoreDiscovery::register: Registration insert completed instance_id={}, key={}, outcome={:?}",
487            instance_id,
488            key_path,
489            outcome
490        );
491
492        Ok(instance)
493    }
494
495    async fn unregister(&self, instance: DiscoveryInstance) -> Result<()> {
496        let (bucket_name, key_path) = match &instance {
497            DiscoveryInstance::Endpoint(inst) => {
498                let key = Self::endpoint_key(inst);
499                tracing::debug!(
500                    "Unregistering endpoint instance_id={}, namespace={}, component={}, endpoint={}, key={}",
501                    inst.instance_id,
502                    inst.namespace,
503                    inst.component,
504                    inst.endpoint,
505                    key
506                );
507                (INSTANCES_BUCKET, key)
508            }
509            DiscoveryInstance::Model {
510                namespace,
511                component,
512                endpoint,
513                instance_id,
514                model_suffix,
515                ..
516            } => {
517                let mut key = Self::model_key(namespace, component, endpoint, *instance_id);
518
519                // If there's a model_suffix (e.g., for LoRA adapters), append it after the instance_id
520                if let Some(suffix) = model_suffix
521                    && !suffix.is_empty()
522                {
523                    key = format!("{}/{}", key, suffix);
524                    tracing::debug!(
525                        "KVStoreDiscovery::unregister: Unregistering LoRA model with suffix={}, instance_id={}, namespace={}, component={}, endpoint={}, key={}",
526                        suffix,
527                        instance_id,
528                        namespace,
529                        component,
530                        endpoint,
531                        key
532                    );
533                }
534
535                // Log for base models (no suffix or empty suffix)
536                if model_suffix.as_ref().is_none_or(|s| s.is_empty()) {
537                    tracing::debug!(
538                        "Unregistering base model instance_id={}, namespace={}, component={}, endpoint={}, key={}",
539                        instance_id,
540                        namespace,
541                        component,
542                        endpoint,
543                        key
544                    );
545                }
546                (MODELS_BUCKET, key)
547            }
548            DiscoveryInstance::EventChannel {
549                scope,
550                topic,
551                instance_id,
552                ..
553            } => {
554                let key = Self::event_channel_key(scope, topic, *instance_id);
555                tracing::debug!(
556                    "KVStoreDiscovery::unregister: Unregistering event channel instance_id={}, scope={:?}, topic={}, key={}",
557                    instance_id,
558                    scope,
559                    topic,
560                    key
561                );
562                (EVENT_CHANNELS_BUCKET, key)
563            }
564            DiscoveryInstance::EventSource {
565                scope,
566                topic,
567                publisher_id,
568                ..
569            } => {
570                let key = Self::event_source_key(scope, topic, *publisher_id);
571                tracing::debug!(
572                    "KVStoreDiscovery::unregister: Unregistering event source publisher_id={}, scope={:?}, topic={}, key={}",
573                    publisher_id,
574                    scope,
575                    topic,
576                    key
577                );
578                (EVENT_SOURCES_BUCKET, key)
579            }
580        };
581
582        // Get the bucket - if it doesn't exist, the instance is already removed from the KV store
583        let Some(bucket) = self.store.get_bucket(bucket_name).await? else {
584            tracing::warn!(
585                "Bucket {} does not exist, instance already removed",
586                bucket_name
587            );
588            return Ok(());
589        };
590
591        let key = kv::Key::new(key_path.clone());
592
593        // Delete the entry from the bucket
594        bucket.delete(&key).await?;
595
596        Ok(())
597    }
598
599    async fn list(&self, query: DiscoveryQuery) -> Result<Vec<DiscoveryInstance>> {
600        let prefix = Self::query_prefix(&query);
601        let bucket_name = Self::bucket_for_prefix(&prefix);
602
603        // Get bucket - if it doesn't exist, return empty list
604        let Some(bucket) = self.store.get_bucket(bucket_name).await? else {
605            tracing::debug!(
606                "KVStoreDiscovery::list: bucket missing for query={:?}, prefix={}, bucket={}",
607                query,
608                prefix,
609                bucket_name
610            );
611            return Ok(Vec::new());
612        };
613
614        // Get all entries from the bucket
615        let entries = bucket.entries().await?;
616        tracing::debug!(
617            "KVStoreDiscovery::list: query={:?}, prefix={}, bucket={}, entries={}",
618            query,
619            prefix,
620            bucket_name,
621            entries.len()
622        );
623
624        // Filter by prefix and deserialize
625        let mut instances = Vec::new();
626        for (key, value) in entries {
627            if Self::matches_prefix(key.as_ref(), &prefix, bucket_name) {
628                match Self::parse_instance(&value) {
629                    Ok(instance) => instances.push(instance),
630                    Err(e) => {
631                        tracing::warn!(%key, error = %e, "Failed to parse discovery instance");
632                    }
633                }
634            }
635        }
636
637        Ok(instances)
638    }
639
640    async fn list_and_watch(
641        &self,
642        query: DiscoveryQuery,
643        cancel_token: Option<CancellationToken>,
644    ) -> Result<DiscoveryStream> {
645        let prefix = Self::query_prefix(&query);
646        let bucket_name = Self::bucket_for_prefix(&prefix);
647
648        tracing::trace!(
649            "KVStoreDiscovery::list_and_watch: Starting watch for query={:?}, prefix={}, bucket={}",
650            query,
651            prefix,
652            bucket_name
653        );
654
655        // Use the provided cancellation token, or fall back to the default token
656        let cancel_token = cancel_token.unwrap_or_else(|| self.cancel_token.clone());
657
658        // Use the kv::Manager's watch mechanism
659        let (_, mut rx) = self.store.clone().watch(
660            bucket_name,
661            None, // No TTL
662            cancel_token,
663        );
664
665        // Create a stream that filters and transforms WatchEvents to DiscoveryEvents
666        let stream = async_stream::stream! {
667            let mut known_instances = HashMap::<DiscoveryInstanceId, DiscoveryInstance>::new();
668
669            while let Some(event) = rx.recv().await {
670                let discovery_events = Self::discovery_events_from_watch_event(
671                    event,
672                    &prefix,
673                    bucket_name,
674                    &mut known_instances,
675                );
676
677                for event in discovery_events {
678                    yield Ok(event);
679                }
680            }
681        };
682        Ok(Box::pin(stream))
683    }
684
685    fn shutdown(&self) {
686        self.store.shutdown();
687    }
688}
689
690#[cfg(test)]
691mod tests {
692    use super::*;
693    use crate::component::TransportType;
694    use crate::discovery::{EventChannelQuery, EventSourceQuery, EventTransport};
695    use crate::protocols::EndpointId;
696
697    fn endpoint_instance(instance_id: u64) -> DiscoveryInstance {
698        DiscoveryInstance::Endpoint(crate::component::Instance {
699            namespace: "ns".to_string(),
700            component: "component".to_string(),
701            endpoint: "endpoint".to_string(),
702            instance_id,
703            transport: TransportType::Nats("nats://127.0.0.1:4222".to_string()),
704            device_type: None,
705            request_plane_codec: None,
706        })
707    }
708
709    fn endpoint_kv(instance_id: u64) -> kv::KeyValue {
710        let instance = endpoint_instance(instance_id);
711        kv::KeyValue::new(
712            kv::Key::new(format!(
713                "{}/{}/{}/{:x}",
714                "ns", "component", "endpoint", instance_id
715            )),
716            serde_json::to_vec(&instance).unwrap().into(),
717        )
718    }
719
720    #[test]
721    fn test_resync_removes_missing_discovery_instances() {
722        let prefix = format!("{}/{}/{}", INSTANCES_BUCKET, "ns", "component");
723        let mut known_instances = HashMap::new();
724
725        let first = endpoint_instance(1);
726        let second = endpoint_instance(2);
727        let third = endpoint_instance(3);
728        known_instances.insert(first.id(), first);
729        known_instances.insert(second.id(), second.clone());
730
731        let mut snapshot = HashMap::new();
732        let second_kv = endpoint_kv(2);
733        snapshot.insert(
734            kv::Key::new(second_kv.key()),
735            second_kv.value().to_vec().into(),
736        );
737        let third_kv = endpoint_kv(3);
738        snapshot.insert(
739            kv::Key::new(third_kv.key()),
740            third_kv.value().to_vec().into(),
741        );
742
743        let events = KVStoreDiscovery::discovery_events_from_watch_event(
744            kv::WatchEvent::Resync(snapshot),
745            &prefix,
746            INSTANCES_BUCKET,
747            &mut known_instances,
748        );
749
750        assert!(!events.contains(&DiscoveryEvent::Added(second)));
751        assert_eq!(
752            events,
753            vec![
754                DiscoveryEvent::Removed(endpoint_instance(1).id()),
755                DiscoveryEvent::Added(third),
756            ]
757        );
758        assert_eq!(known_instances.len(), 2);
759        assert!(known_instances.contains_key(&endpoint_instance(2).id()));
760        assert!(known_instances.contains_key(&endpoint_instance(3).id()));
761    }
762
763    #[test]
764    fn test_resync_retains_known_instance_on_parse_failure() {
765        let prefix = format!("{}/{}/{}", INSTANCES_BUCKET, "ns", "component");
766        let mut known_instances = HashMap::new();
767
768        let first = endpoint_instance(1);
769        known_instances.insert(first.id(), first.clone());
770
771        let mut snapshot = HashMap::new();
772        snapshot.insert(
773            kv::Key::new(format!("ns/component/endpoint/{:x}", 1)),
774            bytes::Bytes::from_static(b"not json"),
775        );
776
777        let events = KVStoreDiscovery::discovery_events_from_watch_event(
778            kv::WatchEvent::Resync(snapshot),
779            &prefix,
780            INSTANCES_BUCKET,
781            &mut known_instances,
782        );
783
784        assert!(events.is_empty());
785        assert_eq!(known_instances.len(), 1);
786        assert_eq!(known_instances.get(&first.id()), Some(&first));
787    }
788
789    #[test]
790    fn test_matches_prefix_requires_path_boundary() {
791        let prefix = format!("{}/{}/{}", INSTANCES_BUCKET, "ns", "component");
792
793        assert!(KVStoreDiscovery::matches_prefix(
794            "ns/component/endpoint/1",
795            &prefix,
796            INSTANCES_BUCKET
797        ));
798        assert!(KVStoreDiscovery::matches_prefix(
799            "ns/component",
800            &prefix,
801            INSTANCES_BUCKET
802        ));
803        assert!(!KVStoreDiscovery::matches_prefix(
804            "ns/component2/endpoint/1",
805            &prefix,
806            INSTANCES_BUCKET
807        ));
808    }
809
810    #[test]
811    fn test_bucket_for_prefix_requires_path_boundary() {
812        assert_eq!(
813            KVStoreDiscovery::bucket_for_prefix("v1/instances/ns/component"),
814            INSTANCES_BUCKET
815        );
816        assert_eq!(
817            KVStoreDiscovery::bucket_for_prefix("v1/event_channels/ns/component/topic"),
818            EVENT_CHANNELS_BUCKET
819        );
820        assert_eq!(
821            KVStoreDiscovery::bucket_for_prefix("v1/event_sources/ns/component/topic"),
822            EVENT_SOURCES_BUCKET
823        );
824        assert_eq!(
825            KVStoreDiscovery::bucket_for_prefix("v1/instances2/ns/component"),
826            MODELS_BUCKET
827        );
828    }
829
830    #[tokio::test]
831    async fn event_channel_keys_and_queries_preserve_exact_endpoint_scope() {
832        let store = kv::Manager::memory();
833        let client = KVStoreDiscovery::new(store, CancellationToken::new());
834        let endpoint_a = EndpointId {
835            namespace: "ns/one".to_string(),
836            component: "worker.component".to_string(),
837            name: "a/*".to_string(),
838        };
839        let endpoint_b = EndpointId {
840            name: "b/>".to_string(),
841            ..endpoint_a.clone()
842        };
843
844        for (publisher_id, endpoint) in [(1, endpoint_a.clone()), (2, endpoint_b.clone())] {
845            client
846                .register(DiscoverySpec::EventChannel {
847                    scope: EventScope::Endpoint { endpoint },
848                    topic: "kv/events".to_string(),
849                    publisher_id,
850                    transport: EventTransport::zmq(format!(
851                        "tcp://127.0.0.1:{}",
852                        5000 + publisher_id
853                    )),
854                })
855                .await
856                .unwrap();
857        }
858
859        let mut a = client
860            .list(DiscoveryQuery::EventChannels(
861                EventChannelQuery::endpoint_topic(endpoint_a.clone(), "kv/events"),
862            ))
863            .await
864            .unwrap();
865        assert_eq!(a.len(), 1);
866        assert_eq!(a[0].instance_id(), 1);
867        client.unregister(a.pop().unwrap()).await.unwrap();
868        assert!(
869            client
870                .list(DiscoveryQuery::EventChannels(
871                    EventChannelQuery::endpoint_topic(endpoint_a, "kv/events"),
872                ))
873                .await
874                .unwrap()
875                .is_empty()
876        );
877
878        let b = client
879            .list(DiscoveryQuery::EventChannels(
880                EventChannelQuery::endpoint_topic(endpoint_b, "kv/events"),
881            ))
882            .await
883            .unwrap();
884        assert_eq!(b.len(), 1);
885        assert_eq!(b[0].instance_id(), 2);
886    }
887
888    async fn assert_event_source_lifecycle(store: kv::Manager) {
889        let client = KVStoreDiscovery::new(store, CancellationToken::new());
890        let endpoint = EndpointId {
891            namespace: "ns/one".to_string(),
892            component: "worker.component".to_string(),
893            name: "decode/*".to_string(),
894        };
895        let query = DiscoveryQuery::EventSources(EventSourceQuery::endpoint_topic(
896            endpoint.clone(),
897            "kv/events",
898        ));
899        let spec = |publisher_id, worker_id| DiscoverySpec::EventSource {
900            scope: EventScope::Endpoint {
901                endpoint: endpoint.clone(),
902            },
903            topic: "kv/events".to_string(),
904            publisher_id,
905            metadata: serde_json::json!({"worker_id": worker_id, "dp_rank": 0}),
906        };
907
908        let first = client.register(spec(100, 7)).await.unwrap();
909        assert_eq!(client.register(spec(100, 7)).await.unwrap(), first);
910        assert!(client.register(spec(100, 8)).await.is_err());
911        assert_eq!(
912            client.list(query.clone()).await.unwrap(),
913            vec![first.clone()]
914        );
915
916        let second = client.register(spec(205, 7)).await.unwrap();
917        assert_eq!(client.list(query.clone()).await.unwrap().len(), 2);
918
919        client.unregister(first).await.unwrap();
920        assert_eq!(client.list(query).await.unwrap(), vec![second]);
921    }
922
923    #[tokio::test]
924    async fn event_source_lifecycle_round_trips_through_memory_kv_discovery() {
925        assert_event_source_lifecycle(kv::Manager::memory()).await;
926    }
927
928    #[tokio::test]
929    async fn event_source_lifecycle_round_trips_through_file_kv_discovery() {
930        let tempdir = tempfile::tempdir().unwrap();
931        let store_cancel = CancellationToken::new();
932        let store = kv::Manager::file(store_cancel.clone(), tempdir.path());
933        assert_event_source_lifecycle(store).await;
934        store_cancel.cancel();
935    }
936
937    #[tokio::test]
938    async fn event_source_watch_removes_exact_publisher_incarnation() {
939        let client = KVStoreDiscovery::new(kv::Manager::memory(), CancellationToken::new());
940        let endpoint = EndpointId {
941            namespace: "ns".to_string(),
942            component: "worker".to_string(),
943            name: "decode".to_string(),
944        };
945        let query = DiscoveryQuery::EventSources(EventSourceQuery::endpoint_topic(
946            endpoint.clone(),
947            "kv-events",
948        ));
949        let mut stream = client.list_and_watch(query, None).await.unwrap();
950        let spec = |publisher_id| DiscoverySpec::EventSource {
951            scope: EventScope::Endpoint {
952                endpoint: endpoint.clone(),
953            },
954            topic: "kv-events".to_string(),
955            publisher_id,
956            metadata: serde_json::json!({"dp_rank": 0}),
957        };
958
959        let first = client.register(spec(100)).await.unwrap();
960        let second = client.register(spec(205)).await.unwrap();
961        let mut added = std::collections::HashSet::new();
962        for _ in 0..2 {
963            let DiscoveryEvent::Added(instance) = stream.next().await.unwrap().unwrap() else {
964                panic!("expected source addition");
965            };
966            added.insert(instance.id());
967        }
968        assert_eq!(
969            added,
970            std::collections::HashSet::from([first.id(), second.id()])
971        );
972
973        client.unregister(first).await.unwrap();
974        let removed = tokio::time::timeout(tokio::time::Duration::from_secs(1), async {
975            loop {
976                if let DiscoveryEvent::Removed(id) = stream.next().await.unwrap().unwrap() {
977                    break id;
978                }
979            }
980        })
981        .await
982        .unwrap();
983        assert_eq!(
984            removed,
985            DiscoveryInstanceId::EventSource(EventSourceInstanceId {
986                scope: EventScope::Endpoint { endpoint },
987                topic: "kv-events".to_string(),
988                publisher_id: 100,
989            })
990        );
991        assert_eq!(
992            client
993                .list(DiscoveryQuery::EventSources(EventSourceQuery::all()))
994                .await
995                .unwrap(),
996            vec![second]
997        );
998    }
999
1000    #[tokio::test]
1001    async fn test_kv_store_discovery_list() {
1002        let store = kv::Manager::memory();
1003        let cancel_token = CancellationToken::new();
1004        let client = KVStoreDiscovery::new(store, cancel_token);
1005
1006        // Register multiple endpoints
1007        let spec1 = DiscoverySpec::Endpoint {
1008            namespace: "ns1".to_string(),
1009            component: "comp1".to_string(),
1010            endpoint: "ep1".to_string(),
1011            device_type: None,
1012            request_plane_codec: None,
1013            transport: TransportType::Nats("nats://localhost:4222".to_string()),
1014        };
1015        client.register(spec1).await.unwrap();
1016
1017        let spec2 = DiscoverySpec::Endpoint {
1018            namespace: "ns1".to_string(),
1019            component: "comp1".to_string(),
1020            device_type: None,
1021            request_plane_codec: None,
1022            endpoint: "ep2".to_string(),
1023            transport: TransportType::Nats("nats://localhost:4222".to_string()),
1024        };
1025        client.register(spec2).await.unwrap();
1026
1027        let spec3 = DiscoverySpec::Endpoint {
1028            namespace: "ns2".to_string(),
1029            device_type: None,
1030            request_plane_codec: None,
1031            component: "comp2".to_string(),
1032            endpoint: "ep1".to_string(),
1033            transport: TransportType::Nats("nats://localhost:4222".to_string()),
1034        };
1035        client.register(spec3).await.unwrap();
1036
1037        // List all endpoints
1038        let all = client.list(DiscoveryQuery::AllEndpoints).await.unwrap();
1039        assert_eq!(all.len(), 3);
1040
1041        // List namespaced endpoints
1042        let ns1 = client
1043            .list(DiscoveryQuery::NamespacedEndpoints {
1044                namespace: "ns1".to_string(),
1045            })
1046            .await
1047            .unwrap();
1048        assert_eq!(ns1.len(), 2);
1049
1050        // List component endpoints
1051        let comp1 = client
1052            .list(DiscoveryQuery::ComponentEndpoints {
1053                namespace: "ns1".to_string(),
1054                component: "comp1".to_string(),
1055            })
1056            .await
1057            .unwrap();
1058        assert_eq!(comp1.len(), 2);
1059    }
1060
1061    #[tokio::test]
1062    async fn test_kv_store_discovery_watch() {
1063        let store = kv::Manager::memory();
1064        let cancel_token = CancellationToken::new();
1065        let client = Arc::new(KVStoreDiscovery::new(store, cancel_token.clone()));
1066
1067        // Start watching before registering
1068        let mut stream = client
1069            .list_and_watch(DiscoveryQuery::AllEndpoints, None)
1070            .await
1071            .unwrap();
1072
1073        let client_clone = client.clone();
1074        let register_task = tokio::spawn(async move {
1075            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
1076
1077            let spec = DiscoverySpec::Endpoint {
1078                device_type: None,
1079                request_plane_codec: None,
1080                namespace: "test".to_string(),
1081                component: "comp1".to_string(),
1082                endpoint: "ep1".to_string(),
1083                transport: TransportType::Nats("nats://localhost:4222".to_string()),
1084            };
1085            client_clone.register(spec).await.unwrap();
1086        });
1087
1088        // Wait for the added event
1089        let event = stream.next().await.unwrap().unwrap();
1090        match event {
1091            DiscoveryEvent::Added(instance) => match instance {
1092                DiscoveryInstance::Endpoint(inst) => {
1093                    assert_eq!(inst.namespace, "test");
1094                    assert_eq!(inst.component, "comp1");
1095                    assert_eq!(inst.endpoint, "ep1");
1096                }
1097                _ => panic!("Expected Endpoint instance"),
1098            },
1099            _ => panic!("Expected Added event"),
1100        }
1101
1102        register_task.await.unwrap();
1103        cancel_token.cancel();
1104    }
1105}
dynamo_runtime/discovery/kv_store.rs

dynamo_runtime/discovery/
kv_store.rs
