client.rs 15.3 KB
Newer Older
1
2
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
Ryan Olson's avatar
Ryan Olson committed
3

4
5
6
7
8
9
10
11
use std::sync::Arc;
use std::{collections::HashMap, time::Duration};

use anyhow::Result;
use arc_swap::ArcSwap;
use futures::StreamExt;
use tokio::net::unix::pipe::Receiver;

12
use crate::discovery::{DiscoveryEvent, DiscoveryInstance};
13
use crate::{
14
15
    component::{Endpoint, Instance},
    pipeline::async_trait,
16
17
18
19
    pipeline::{
        AddressedPushRouter, AddressedRequest, AsyncEngine, Data, ManyOut, PushRouter, RouterMode,
        SingleIn,
    },
20
21
    traits::DistributedRuntimeProvider,
    transports::etcd::Client as EtcdClient,
Ryan Olson's avatar
Ryan Olson committed
22
23
};

24
25
26
/// Default interval for periodic reconciliation of instance_avail with instance_source
const DEFAULT_RECONCILE_INTERVAL: Duration = Duration::from_secs(5);

27
28
29
30
#[derive(Clone, Debug)]
pub struct Client {
    // This is me
    pub endpoint: Endpoint,
31
32
    // These are the remotes I know about from watching key-value store
    pub instance_source: Arc<tokio::sync::watch::Receiver<Vec<Instance>>>,
33
    // These are the instance source ids less those reported as down from sending rpc
34
    instance_avail: Arc<ArcSwap<Vec<u64>>>,
35
    // These are the instance source ids less those reported as busy (above threshold)
36
    instance_free: Arc<ArcSwap<Vec<u64>>>,
37
38
39
40
    // Watch sender for available instance IDs (for sending updates)
    instance_avail_tx: Arc<tokio::sync::watch::Sender<Vec<u64>>>,
    // Watch receiver for available instance IDs (for cloning to external subscribers)
    instance_avail_rx: tokio::sync::watch::Receiver<Vec<u64>>,
41
42
43
    /// Interval for periodic reconciliation of instance_avail with instance_source.
    /// This ensures instances removed via `report_instance_down` are eventually restored.
    reconcile_interval: Duration,
44
45
}

46
impl Client {
47
48
    // Client with auto-discover instances using key-value store
    pub(crate) async fn new(endpoint: Endpoint) -> Result<Self> {
49
50
51
52
53
54
55
56
57
58
        Self::with_reconcile_interval(endpoint, DEFAULT_RECONCILE_INTERVAL).await
    }

    /// Create a client with a custom reconcile interval.
    /// The reconcile interval controls how often `instance_avail` is reset to match
    /// `instance_source`, restoring any instances removed via `report_instance_down`.
    pub(crate) async fn with_reconcile_interval(
        endpoint: Endpoint,
        reconcile_interval: Duration,
    ) -> Result<Self> {
59
        tracing::trace!(
60
            "Client::new_dynamic: Creating dynamic client for endpoint: {}",
61
            endpoint.id()
62
        );
63
        let instance_source = Self::get_or_create_dynamic_instance_source(&endpoint).await?;
64

65
        let (avail_tx, avail_rx) = tokio::sync::watch::channel(vec![]);
66
        let client = Client {
67
            endpoint: endpoint.clone(),
68
            instance_source: instance_source.clone(),
69
            instance_avail: Arc::new(ArcSwap::from(Arc::new(vec![]))),
70
            instance_free: Arc::new(ArcSwap::from(Arc::new(vec![]))),
71
72
            instance_avail_tx: Arc::new(avail_tx),
            instance_avail_rx: avail_rx,
73
            reconcile_interval,
74
        };
75
        client.monitor_instance_source();
76
        Ok(client)
77
78
    }

79
    /// Instances available from watching key-value store
80
    pub fn instances(&self) -> Vec<Instance> {
81
        self.instance_source.borrow().clone()
82
83
    }

84
    pub fn instance_ids(&self) -> Vec<u64> {
85
86
87
        self.instances().into_iter().map(|ep| ep.id()).collect()
    }

88
    pub fn instance_ids_avail(&self) -> arc_swap::Guard<Arc<Vec<u64>>> {
89
90
91
        self.instance_avail.load()
    }

92
    pub fn instance_ids_free(&self) -> arc_swap::Guard<Arc<Vec<u64>>> {
93
94
95
        self.instance_free.load()
    }

96
97
98
99
100
    /// Get a watcher for available instance IDs
    pub fn instance_avail_watcher(&self) -> tokio::sync::watch::Receiver<Vec<u64>> {
        self.instance_avail_rx.clone()
    }

101
102
    /// Wait for at least one Instance to be available for this Endpoint
    pub async fn wait_for_instances(&self) -> Result<Vec<Instance>> {
103
        tracing::trace!(
104
            "wait_for_instances: Starting wait for endpoint: {}",
105
            self.endpoint.id()
106
        );
107
108
109
110
111
112
113
114
115
116
        let mut rx = self.instance_source.as_ref().clone();
        // wait for there to be 1 or more endpoints
        let mut instances: Vec<Instance>;
        loop {
            instances = rx.borrow_and_update().to_vec();
            if instances.is_empty() {
                rx.changed().await?;
            } else {
                tracing::info!(
                    "wait_for_instances: Found {} instance(s) for endpoint: {}",
117
                    instances.len(),
118
                    self.endpoint.id()
119
                );
120
                break;
121
122
123
124
125
            }
        }
        Ok(instances)
    }

126
    /// Mark an instance as down/unavailable
127
    pub fn report_instance_down(&self, instance_id: u64) {
128
129
130
131
132
        let filtered = self
            .instance_ids_avail()
            .iter()
            .filter_map(|&id| if id == instance_id { None } else { Some(id) })
            .collect::<Vec<_>>();
133
134
135
136
        self.instance_avail.store(Arc::new(filtered.clone()));

        // Notify watch channel subscribers about the change
        let _ = self.instance_avail_tx.send(filtered);
137
138
139
140

        tracing::debug!("inhibiting instance {instance_id}");
    }

141
    /// Update the set of free instances based on busy instance IDs
142
    pub fn update_free_instances(&self, busy_instance_ids: &[u64]) {
143
        let all_instance_ids = self.instance_ids();
144
        let free_ids: Vec<u64> = all_instance_ids
145
146
147
148
149
150
            .into_iter()
            .filter(|id| !busy_instance_ids.contains(id))
            .collect();
        self.instance_free.store(Arc::new(free_ids));
    }

151
    /// Monitor the key-value instance source and update instance_avail.
152
153
154
155
156
    ///
    /// This function also performs periodic reconciliation: if `instance_source` hasn't
    /// changed for `reconcile_interval`, we reset `instance_avail` to match
    /// `instance_source`. This ensures instances removed via `report_instance_down`
    /// are eventually restored even if the discovery source doesn't emit updates.
157
    fn monitor_instance_source(&self) {
158
        let reconcile_interval = self.reconcile_interval;
159
160
        let cancel_token = self.endpoint.drt().primary_token();
        let client = self.clone();
161
        let endpoint_id = self.endpoint.id();
162
        tokio::task::spawn(async move {
163
            let mut rx = client.instance_source.as_ref().clone();
164
            while !cancel_token.is_cancelled() {
165
                let instance_ids: Vec<u64> = rx
166
167
168
169
                    .borrow_and_update()
                    .iter()
                    .map(|instance| instance.id())
                    .collect();
170
171
172

                // TODO: this resets both tracked available and free instances
                client.instance_avail.store(Arc::new(instance_ids.clone()));
173
                client.instance_free.store(Arc::new(instance_ids.clone()));
174

175
176
177
                // Send update to watch channel subscribers
                let _ = client.instance_avail_tx.send(instance_ids);

178
179
180
181
182
183
184
185
186
187
188
189
190
191
                tokio::select! {
                    result = rx.changed() => {
                        if let Err(err) = result {
                            tracing::error!(
                                "monitor_instance_source: The Sender is dropped: {err}, endpoint={endpoint_id}",
                            );
                            cancel_token.cancel();
                        }
                    }
                    _ = tokio::time::sleep(reconcile_interval) => {
                        tracing::trace!(
                            "monitor_instance_source: periodic reconciliation for endpoint={endpoint_id}",
                        );
                    }
192
193
194
                }
            }
        });
195
196
197
198
    }

    async fn get_or_create_dynamic_instance_source(
        endpoint: &Endpoint,
199
    ) -> Result<Arc<tokio::sync::watch::Receiver<Vec<Instance>>>> {
200
201
202
203
204
205
206
207
208
209
210
211
        let drt = endpoint.drt();
        let instance_sources = drt.instance_sources();
        let mut instance_sources = instance_sources.lock().await;

        if let Some(instance_source) = instance_sources.get(endpoint) {
            if let Some(instance_source) = instance_source.upgrade() {
                return Ok(instance_source);
            } else {
                instance_sources.remove(endpoint);
            }
        }

212
213
214
215
216
217
218
219
220
221
        let discovery = drt.discovery();
        let discovery_query = crate::discovery::DiscoveryQuery::Endpoint {
            namespace: endpoint.component.namespace.name.clone(),
            component: endpoint.component.name.clone(),
            endpoint: endpoint.name.clone(),
        };

        let mut discovery_stream = discovery
            .list_and_watch(discovery_query.clone(), None)
            .await?;
Ryan Olson's avatar
Ryan Olson committed
222
223
        let (watch_tx, watch_rx) = tokio::sync::watch::channel(vec![]);

224
        let secondary = endpoint.component.drt.runtime().secondary().clone();
Ryan Olson's avatar
Ryan Olson committed
225
226

        secondary.spawn(async move {
227
            tracing::trace!("endpoint_watcher: Starting for discovery query: {:?}", discovery_query);
228
            let mut map: HashMap<u64, Instance> = HashMap::new();
Ryan Olson's avatar
Ryan Olson committed
229
230

            loop {
231
                let discovery_event = tokio::select! {
Ryan Olson's avatar
Ryan Olson committed
232
233
234
                    _ = watch_tx.closed() => {
                        break;
                    }
235
236
237
238
239
240
241
242
243
                    discovery_event = discovery_stream.next() => {
                        match discovery_event {
                            Some(Ok(event)) => {
                                event
                            },
                            Some(Err(e)) => {
                                tracing::error!("endpoint_watcher: discovery stream error: {}; shutting down for discovery query: {:?}", e, discovery_query);
                                break;
                            }
Ryan Olson's avatar
Ryan Olson committed
244
245
246
247
248
249
250
                            None => {
                                break;
                            }
                        }
                    }
                };

251
                match discovery_event {
252
253
254
                    DiscoveryEvent::Added(discovery_instance) => {
                        if let DiscoveryInstance::Endpoint(instance) = discovery_instance {

255
                                map.insert(instance.instance_id, instance);
Ryan Olson's avatar
Ryan Olson committed
256
                        }
257
                    }
258
                    DiscoveryEvent::Removed(instance_id) => {
259
                        map.remove(&instance_id);
Ryan Olson's avatar
Ryan Olson committed
260
261
262
                    }
                }

263
264
                let instances: Vec<Instance> = map.values().cloned().collect();
                if watch_tx.send(instances).is_err() {
Ryan Olson's avatar
Ryan Olson committed
265
266
267
268
269
270
                    break;
                }
            }
            let _ = watch_tx.send(vec![]);
        });

271
        let instance_source = Arc::new(watch_rx);
272
273
        instance_sources.insert(endpoint.clone(), Arc::downgrade(&instance_source));
        Ok(instance_source)
274
    }
Ryan Olson's avatar
Ryan Olson committed
275
}
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{DistributedRuntime, Runtime, distributed::DistributedConfig};

    /// Test that instances removed via report_instance_down are restored after
    /// the reconciliation interval elapses.
    #[tokio::test]
    async fn test_instance_reconciliation() {
        const TEST_RECONCILE_INTERVAL: Duration = Duration::from_millis(100);

        let rt = Runtime::from_current().unwrap();
        // Use process_local config to avoid needing etcd/nats
        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
            .await
            .unwrap();
        let ns = drt.namespace("test_reconciliation".to_string()).unwrap();
        let component = ns.component("test_component".to_string()).unwrap();
        let endpoint = component.endpoint("test_endpoint".to_string());

        // Use a short reconcile interval for faster tests
        let client = Client::with_reconcile_interval(endpoint, TEST_RECONCILE_INTERVAL)
            .await
            .unwrap();

        // Initially, instance_avail should be empty (no registered instances)
        assert!(client.instance_ids_avail().is_empty());

        // For this test, we'll directly manipulate instance_avail and verify reconciliation
        // Store some test IDs
        client.instance_avail.store(Arc::new(vec![1, 2, 3]));

        assert_eq!(**client.instance_ids_avail(), vec![1u64, 2, 3]);

        // Simulate report_instance_down removing instance 2
        client.report_instance_down(2);
        assert_eq!(**client.instance_ids_avail(), vec![1u64, 3]);

        // Wait for reconciliation interval + buffer
        // The monitor_instance_source will reset instance_avail to match instance_source
        // Since instance_source is empty, after reconciliation instance_avail should be empty
        tokio::time::sleep(TEST_RECONCILE_INTERVAL + Duration::from_millis(50)).await;

        // After reconciliation, instance_avail should match instance_source (which is empty)
        assert!(
            client.instance_ids_avail().is_empty(),
            "After reconciliation, instance_avail should match instance_source"
        );

        rt.shutdown();
    }

    /// Test that report_instance_down correctly removes an instance from instance_avail.
    #[tokio::test]
    async fn test_report_instance_down() {
        let rt = Runtime::from_current().unwrap();
        // Use process_local config to avoid needing etcd/nats
        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
            .await
            .unwrap();
        let ns = drt.namespace("test_report_down".to_string()).unwrap();
        let component = ns.component("test_component".to_string()).unwrap();
        let endpoint = component.endpoint("test_endpoint".to_string());

        let client = endpoint.client().await.unwrap();

        // Manually set up instance_avail with test instances
        client.instance_avail.store(Arc::new(vec![1, 2, 3]));
        assert_eq!(**client.instance_ids_avail(), vec![1u64, 2, 3]);

        // Report instance 2 as down
        client.report_instance_down(2);

        // Verify instance 2 is removed
        let avail = client.instance_ids_avail();
        assert!(avail.contains(&1), "Instance 1 should still be available");
        assert!(
            !avail.contains(&2),
            "Instance 2 should be removed after report_instance_down"
        );
        assert!(avail.contains(&3), "Instance 3 should still be available");

        rt.shutdown();
    }

    /// Test that instance_avail_watcher receives updates when instances change.
    #[tokio::test]
    async fn test_instance_avail_watcher() {
        let rt = Runtime::from_current().unwrap();
        // Use process_local config to avoid needing etcd/nats
        let drt = DistributedRuntime::new(rt.clone(), DistributedConfig::process_local())
            .await
            .unwrap();
        let ns = drt.namespace("test_watcher".to_string()).unwrap();
        let component = ns.component("test_component".to_string()).unwrap();
        let endpoint = component.endpoint("test_endpoint".to_string());

        let client = endpoint.client().await.unwrap();
        let watcher = client.instance_avail_watcher();

        // Set initial instances
        client.instance_avail.store(Arc::new(vec![1, 2, 3]));

        // Report instance down - this should notify the watcher
        client.report_instance_down(2);

        // The watcher should receive the update
        // Note: We need to check if changed() was signaled
        let current = watcher.borrow().clone();
        assert_eq!(current, vec![1, 3]);

        rt.shutdown();
    }
}