Skip to content

fix: implement discovery interface with KeyValueStoreManager #3944

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mohammedabdulwahhab wants to merge 3 commits into from
Closed

fix: implement discovery interface with KeyValueStoreManager #3944

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
aaa5166
fix: introduce service discovery interface
mohammedabdulwahhab Oct 28, 2025
e262cd3
fix: resolve merge conflicts
mohammedabdulwahhab Oct 28, 2025
2a2a344
fix: define key value store
mohammedabdulwahhab Oct 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 77 additions & 42 deletions lib/runtime/src/component/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ use std::sync::Arc;
use tokio::net::unix::pipe::Receiver;

use crate::{
discovery::{DiscoveryClient, DiscoveryEvent, DiscoveryKey},
pipeline::async_trait,
transports::etcd::{Client as EtcdClient, WatchEvent},
};
Expand Down Expand Up @@ -65,17 +66,15 @@ impl Client {
})
}

// Client with auto-discover instances using etcd
// Client with auto-discover instances using discovery client
pub(crate) async fn new_dynamic(endpoint: Endpoint) -> Result<Self> {
const INSTANCE_REFRESH_PERIOD: Duration = Duration::from_secs(1);

@coderabbitai coderabbitai Bot Oct 28, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🔴 Critical

Compilation blocker: Duration not in scope and const is unused.

This const won’t compile (missing use std::time::Duration;) and it’s unused. Remove it or import and use it.

Apply this minimal fix (remove the unused const):

-        const INSTANCE_REFRESH_PERIOD: Duration = Duration::from_secs(1);
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
const INSTANCE_REFRESH_PERIOD: Duration = Duration::from_secs(1);
🤖 Prompt for AI Agents 
In lib/runtime/src/component/client.rs around line 71, the const declaration
`INSTANCE_REFRESH_PERIOD: Duration = Duration::from_secs(1);` is unused and
`Duration` is not imported; remove this unused const from the file (delete the
line) so the code compiles without adding an unnecessary import.


// create live endpoint watcher
let Some(etcd_client) = &endpoint.component.drt.etcd_client else {
anyhow::bail!("Attempt to create a dynamic client on a static endpoint");
};
// Get the discovery client from DRT
let discovery_client = endpoint.component.drt.discovery_client().await?;

let instance_source =
Self::get_or_create_dynamic_instance_source(etcd_client, &endpoint).await?;
Self::get_or_create_dynamic_instance_source(discovery_client, &endpoint).await?;

let client = Client {
endpoint,
Expand Down Expand Up @@ -194,7 +193,7 @@ impl Client {
}

async fn get_or_create_dynamic_instance_source(
etcd_client: &EtcdClient,
discovery_client: Arc<dyn DiscoveryClient>,
endpoint: &Endpoint,
) -> Result<Arc<InstanceSource>> {
let drt = endpoint.drt();
Expand All @@ -209,72 +208,108 @@ impl Client {
}
}

let prefix_watcher = etcd_client
.kv_get_and_watch_prefix(endpoint.etcd_root())
.await?;
// Create discovery key from endpoint info
let discovery_key = DiscoveryKey::Endpoint {
namespace: endpoint.component.namespace.name.clone(),
component: endpoint.component.name.clone(),
endpoint: endpoint.name.clone(),
};

let (prefix, _watcher, mut kv_event_rx) = prefix_watcher.dissolve();
// Start list_and_watch on the discovery client
let mut discovery_stream = discovery_client
.list_and_watch(discovery_key.clone())
.await?;

let (watch_tx, watch_rx) = tokio::sync::watch::channel(vec![]);

let secondary = endpoint.component.drt.runtime.secondary().clone();
let endpoint_path = endpoint.path();

// this task should be included in the registry
// currently this is created once per client, but this object/task should only be instantiated
// once per worker/instance
// Spawn task to handle discovery events
secondary.spawn(async move {
tracing::debug!("Starting endpoint watcher for prefix: {}", prefix);
tracing::debug!("Starting discovery watcher for endpoint: {}", endpoint_path);
let mut map = HashMap::new();

// Use StreamExt to iterate over the discovery stream
use futures::StreamExt;

loop {
let kv_event = tokio::select! {
let discovery_event = tokio::select! {
_ = watch_tx.closed() => {
tracing::debug!("all watchers have closed; shutting down endpoint watcher for prefix: {prefix}");
tracing::debug!("all watchers have closed; shutting down discovery watcher for endpoint: {endpoint_path}");
break;
}
kv_event = kv_event_rx.recv() => {
match kv_event {
Some(kv_event) => kv_event,
event = discovery_stream.next() => {
match event {
Some(Ok(event)) => event,
Some(Err(e)) => {
tracing::error!("discovery stream error: {}; shutting down discovery watcher for endpoint: {endpoint_path}", e);
break;
}
None => {
tracing::debug!("watch stream has closed; shutting down endpoint watcher for prefix: {prefix}");
tracing::debug!("discovery stream closed; shutting down discovery watcher for endpoint: {endpoint_path}");
break;
}
}
}
};
Comment on lines 225 to 255

@coderabbitai coderabbitai Bot Oct 28, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🛠️ Refactor suggestion | 🟠 Major

Add runtime cancellation to the discovery watcher to avoid stranded tasks.

The spawned watcher exits only on receiver close or stream end. It should also listen to the runtime’s primary cancellation token.

Apply:

         let secondary = endpoint.component.drt.runtime.secondary().clone();
+        let cancel_token = endpoint.drt().primary_token();
         let endpoint_path = endpoint.path();

         // Spawn task to handle discovery events
         secondary.spawn(async move {
             tracing::debug!("Starting discovery watcher for endpoint: {}", endpoint_path);
             let mut map = HashMap::new();

             // Use StreamExt to iterate over the discovery stream
             use futures::StreamExt;

             loop {
-                let discovery_event = tokio::select! {
+                let discovery_event = tokio::select! {
+                    _ = cancel_token.cancelled() => {
+                        tracing::debug!("runtime cancelled; shutting down discovery watcher for endpoint: {endpoint_path}");
+                        break;
+                    }
                     _ = watch_tx.closed() => {
                         tracing::debug!("all watchers have closed; shutting down discovery watcher for endpoint: {endpoint_path}");
                         break;
                     }
                     event = discovery_stream.next() => {

Committable suggestion skipped: line range outside the PR's diff.

🤖 Prompt for AI Agents 
lib/runtime/src/component/client.rs around lines 225 to 255: the discovery
watcher currently only stops on watch_tx closed or discovery stream end/error;
add listening for the runtime’s primary cancellation token so the spawned task
is not left stranded. Clone or capture the runtime primary cancellation handle
(e.g., let primary = endpoint.component.drt.runtime.primary().clone()) before
spawning, then add a branch to the tokio::select! inside the loop that awaits
primary.cancelled() (or the equivalent cancel future) and on that branch log a
debug about runtime cancellation for the endpoint_path and break the loop;
ensure the token is moved into the async block and that the select branch is
ordered with the other shutdown conditions.


match kv_event {
WatchEvent::Put(kv) => {
let key = String::from_utf8(kv.key().to_vec());
let val = serde_json::from_slice::<Instance>(kv.value());
if let (Ok(key), Ok(val)) = (key, val) {
map.insert(key.clone(), val);
} else {
tracing::error!("Unable to parse put endpoint event; shutting down endpoint watcher for prefix: {prefix}");
break;
}
match discovery_event {
DiscoveryEvent::Added(discovery_instance) => {
// Extract fields from discovery instance
let (namespace, component, endpoint_name, instance_id) = match &discovery_instance {
crate::discovery::DiscoveryInstance::Endpoint {
namespace,
component,
endpoint,
instance_id
} => (
namespace.clone(),
component.clone(),
endpoint.clone(),
*instance_id
),
};

// For now, create a test Instance with NatsTcp transport
// TODO: Discovery should store full Instance details, not just metadata
let instance = Instance {
component,
endpoint: endpoint_name,
namespace,
instance_id,
transport: TransportType::NatsTcp("".to_string()),
};

tracing::debug!(
instance_id = %instance.id(),
endpoint = %endpoint_path,
"Discovery: Added instance"
);

// Use instance_id as the key in the HashMap
map.insert(instance_id.to_string(), instance);
}
WatchEvent::Delete(kv) => {
match String::from_utf8(kv.key().to_vec()) {
Ok(key) => { map.remove(&key); }
Err(_) => {
tracing::error!("Unable to parse delete endpoint event; shutting down endpoint watcher for prefix: {}", prefix);
break;
}
}
DiscoveryEvent::Removed(instance_id) => {
tracing::debug!(
instance_id = %instance_id,
endpoint = %endpoint_path,
"Discovery: Removed instance"
);
map.remove(&instance_id.to_string());
}
}

// Send complete snapshot of current instances
let instances: Vec<Instance> = map.values().cloned().collect();

if watch_tx.send(instances).is_err() {
tracing::debug!("Unable to send watch updates; shutting down endpoint watcher for prefix: {}", prefix);
tracing::debug!("Unable to send watch updates; shutting down discovery watcher for endpoint: {}", endpoint_path);
break;
}

}

tracing::debug!("Completed endpoint watcher for prefix: {prefix}");
tracing::debug!("Completed discovery watcher for endpoint: {endpoint_path}");
let _ = watch_tx.send(vec![]);
});

Expand Down
49 changes: 32 additions & 17 deletions lib/runtime/src/component/endpoint.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,8 +118,6 @@ impl EndpointConfigBuilder {
let endpoint_name = endpoint.name.clone();
let system_health = endpoint.drt().system_health.clone();
let subject = endpoint.subject_to(connection_id);
let etcd_path = endpoint.etcd_path_with_lease_id(connection_id);
let etcd_client = endpoint.component.drt.etcd_client.clone();

// Register health check target in SystemHealth if provided
if let Some(health_check_payload) = &health_check_payload {
Expand Down Expand Up @@ -193,35 +191,52 @@ impl EndpointConfigBuilder {
result
});

// make the components service endpoint discovery in etcd

// client.register_service()
let info = Instance {
// Register the endpoint with the discovery service
let discovery_key = crate::discovery::DiscoveryKey::Endpoint {
namespace: namespace_name.clone(),
component: component_name.clone(),
endpoint: endpoint_name.clone(),
namespace: namespace_name.clone(),
instance_id: connection_id,
transport: TransportType::NatsTcp(subject),
};

let info = serde_json::to_vec_pretty(&info)?;
// Attempt to get discovery client and register
let discovery_client = match endpoint.component.drt.discovery_client().await {
Ok(client) => client,
Err(e) => {
tracing::error!(
component_name,
endpoint_name,
error = %e,
"Unable to get discovery client"
);
endpoint_shutdown_token.cancel();
return Err(error!(
"Unable to get discovery client. Check discovery service configuration"
));
}
};

if let Some(etcd_client) = &etcd_client
&& let Err(e) = etcd_client
.kv_create(&etcd_path, info, Some(connection_id))
.await
{
// Use the discovery client's instance_id (which is the connection_id)
let instance_id = discovery_client.instance_id();

if let Err(e) = discovery_client.serve(discovery_key).await {
tracing::error!(
component_name,
endpoint_name,
error = %e,
"Unable to register service for discovery"
"Unable to register endpoint with discovery service"
);
endpoint_shutdown_token.cancel();
return Err(error!(
"Unable to register service for discovery. Check discovery service status"
"Unable to register endpoint with discovery service. Check discovery service status"
));
}

tracing::debug!(
component_name,
endpoint_name,
instance_id = %instance_id,
"Successfully registered endpoint with discovery service"
);
Comment on lines +194 to +239

@coderabbitai coderabbitai Bot Oct 28, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

Deregister on shutdown and gate registration on readiness

  • No deregistration → stale instances accumulate (TTL=0 in KV impl). Add a deregister path (trait method or guard) so Delete events are emitted on shutdown.
  • Registration happens immediately after spawning without a readiness barrier; clients may discover before the NATS endpoint is fully serving.

Recommendations:

  • Extend DiscoveryClient with async fn deregister(&self, key: DiscoveryKey) -> Result<()> and call it when the task ends (or via a drop-guard tied to endpoint_shutdown_token).
  • If possible, register only after the endpoint is ready (or emit readiness from push_endpoint.start), then deregister before cancel/unregister.

I'll verify the core concerns by examining the codebase structure and implementation details.


---

🏁 Script executed:

```shell
#!/bin/bash
# Extract full function context around endpoint registration code
cd lib/runtime/src/component && cat endpoint.rs | head -300 | tail -150

Length of output: 5706

🏁 Script executed:

#!/bin/bash
# Search for DiscoveryClient trait definition and deregister methods
rg -nP "trait\s+DiscoveryClient|async\s+fn\s+deregister|async\s+fn\s+serve" --type=rs -A5

Length of output: 26

🏁 Script executed:

#!/bin/bash
# Search for endpoint shutdown/cleanup patterns and drop implementations
rg -nP "endpoint_shutdown_token|impl.*Drop|cleanup|deregister" --type=rs -C3 lib/runtime/src/component/endpoint.rs

Length of output: 26

🏁 Script executed:

#!/bin/bash
# Check discovery service implementation for TTL and Delete event handling
fd -e rs -path "*discovery*" -type f | head -20

Length of output: 168

🏁 Script executed:

#!/bin/bash
# Search for readiness gates or checks before discovery registration
rg -nP "ready|readiness|Ready" lib/runtime/src/component/endpoint.rs -B5 -A5

Length of output: 0


---

🏁 Script executed:

```shell
#!/bin/bash
# Search for DiscoveryClient trait definition with corrected syntax
rg "trait DiscoveryClient" -A 20

Length of output: 955

🏁 Script executed:

#!/bin/bash
# Search for serve and deregister methods in discovery client implementations
rg "async fn (serve|deregister)" -B2 -A8

Length of output: 2111

🏁 Script executed:

#!/bin/bash
# Look for PushEndpoint and its start method
rg "impl PushEndpoint|fn start\(" -B2 -A15 lib/runtime/src/component/endpoint.rs

Length of output: 470

🏁 Script executed:

#!/bin/bash
# Find discovery module files
find . -path "*/discovery*" -type f -name "*.rs" | head -15

Length of output: 256

🏁 Script executed:

#!/bin/bash
# Search for endpoint shutdown behavior and where discovery deregistration might occur
rg "endpoint_shutdown_token|\.cancel\(\)" -B2 -A2

Length of output: 38482


---

🏁 Script executed:

```shell
#!/bin/bash
# Check discovery service implementation for TTL/cleanup handling
rg "TTL|ttl|cleanup|Delete|Removed" lib/runtime/src/discovery/ -B2 -A5

Length of output: 4320

🏁 Script executed:

#!/bin/bash
# Look at DiscoveryInstance structure and lifecycle
rg "pub enum DiscoveryInstance|pub struct DiscoveryInstance" lib/runtime/src/discovery/ -A15

Length of output: 949

🏁 Script executed:

#!/bin/bash
# Check PushEndpoint.start() return type and whether it signals readiness
rg "pub async fn start" lib/runtime/src/component/endpoint.rs -B2 -A30 | head -60

Length of output: 1099

🏁 Script executed:

#!/bin/bash
# Search for any deregister or cleanup calls in discovery or component shutdown
rg "deregister|Remove|Delete.*discovery|discovery.*Remove" lib/runtime/src/ -B3 -A3

Length of output: 8388

Add deregister() trait method and call on shutdown; gate discovery registration on endpoint readiness

The DiscoveryClient trait lacks a deregister() method, and the endpoint discovery entry is never deleted on shutdown. With TTL=0 in the key-value store, stale registrations accumulate indefinitely. Additionally, discovery_client.serve() is called immediately after retrieving the client, before the spawned endpoint task may be ready to accept connections, creating a race where clients can discover and attempt to connect to a not-yet-listening endpoint.

Recommendations:

  • Extend DiscoveryClient trait with async fn deregister(&self, key: DiscoveryKey) -> Result<()> and invoke it when the endpoint task completes or during shutdown.
  • Defer discovery_client.serve() until after push_endpoint.start() signals readiness, or wait for the endpoint to become operational before registering.
🤖 Prompt for AI Agents 
In lib/runtime/src/component/endpoint.rs around lines 194 to 239, the discovery
registration is performed immediately after obtaining the DiscoveryClient and
there is no deregister() call on shutdown; this causes stale registrations when
TTL=0 and a race where clients can discover an endpoint before it is ready. Add
an async deregister(&self, key: DiscoveryKey) -> Result<()> method to the
DiscoveryClient trait and call it during endpoint shutdown or when the spawned
endpoint task completes (use the existing endpoint_shutdown_token or task join
to trigger deregistration). Also defer calling
discovery_client.serve(discovery_key).await until after the endpoint signals
readiness (e.g., wait for push_endpoint.start() readiness notification or for
the spawned listener to be bound) so registration only happens once the endpoint
can accept connections.

task.await??;

Ok(())
Expand Down
Loading
Loading