distributed.rs 6.13 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
Ryan Olson's avatar
Ryan Olson committed
15
16
17

pub use crate::component::Component;
use crate::{
18
    component::{self, ComponentBuilder, Endpoint, InstanceSource, Namespace},
Ryan Olson's avatar
Ryan Olson committed
19
20
21
22
23
24
    discovery::DiscoveryClient,
    service::ServiceClient,
    transports::{etcd, nats, tcp},
    ErrorContext,
};

25
use super::{error, Arc, DistributedRuntime, OnceCell, Result, Runtime, Weak, OK};
Ryan Olson's avatar
Ryan Olson committed
26
27
28

use derive_getters::Dissolve;
use figment::error;
29
30
use std::collections::HashMap;
use tokio::sync::Mutex;
31
use tokio_util::sync::CancellationToken;
Ryan Olson's avatar
Ryan Olson committed
32
33
34
35

impl DistributedRuntime {
    pub async fn new(runtime: Runtime, config: DistributedConfig) -> Result<Self> {
        let secondary = runtime.secondary();
36
        let (etcd_config, nats_config, is_static) = config.dissolve();
Ryan Olson's avatar
Ryan Olson committed
37
38
39

        let runtime_clone = runtime.clone();

40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
        let etcd_client = if is_static {
            None
        } else {
            Some(
                secondary
                    .spawn(async move {
                        let client = etcd::Client::new(etcd_config.clone(), runtime_clone)
                            .await
                            .context(format!(
                                "Failed to connect to etcd server with config {:?}",
                                etcd_config
                            ))?;
                        OK(client)
                    })
                    .await??,
            )
        };
Ryan Olson's avatar
Ryan Olson committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

        let nats_client = secondary
            .spawn(async move {
                let client = nats_config.clone().connect().await.context(format!(
                    "Failed to connect to NATS server with config {:?}",
                    nats_config
                ))?;
                anyhow::Ok(client)
            })
            .await??;

        Ok(Self {
            runtime,
            etcd_client,
            nats_client,
            tcp_server: Arc::new(OnceCell::new()),
            component_registry: component::Registry::new(),
74
            is_static,
75
            instance_sources: Arc::new(Mutex::new(HashMap::new())),
Ryan Olson's avatar
Ryan Olson committed
76
77
78
79
        })
    }

    pub async fn from_settings(runtime: Runtime) -> Result<Self> {
80
81
82
83
84
85
86
        let config = DistributedConfig::from_settings(false);
        Self::new(runtime, config).await
    }

    // Call this if you are using static workers that do not need etcd-based discovery.
    pub async fn from_settings_without_discovery(runtime: Runtime) -> Result<Self> {
        let config = DistributedConfig::from_settings(true);
Ryan Olson's avatar
Ryan Olson committed
87
88
89
90
91
92
93
        Self::new(runtime, config).await
    }

    pub fn runtime(&self) -> &Runtime {
        &self.runtime
    }

94
95
96
97
    pub fn primary_token(&self) -> CancellationToken {
        self.runtime.primary_token()
    }

98
99
100
101
    /// The etcd lease all our components will be attached to.
    /// Not available for static workers.
    pub fn primary_lease(&self) -> Option<etcd::Lease> {
        self.etcd_client.as_ref().map(|c| c.primary_lease())
Ryan Olson's avatar
Ryan Olson committed
102
103
104
105
106
107
108
109
    }

    pub fn shutdown(&self) {
        self.runtime.shutdown();
    }

    /// Create a [`Namespace`]
    pub fn namespace(&self, name: impl Into<String>) -> Result<Namespace> {
110
        Namespace::new(self.clone(), name.into(), self.is_static)
Ryan Olson's avatar
Ryan Olson committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
    }

    // /// Create a [`Component`]
    // pub fn component(
    //     &self,
    //     name: impl Into<String>,
    //     namespace: impl Into<String>,
    // ) -> Result<Component> {
    //     Ok(ComponentBuilder::from_runtime(self.clone())
    //         .name(name.into())
    //         .namespace(namespace.into())
    //         .build()?)
    // }

    pub(crate) fn discovery_client(&self, namespace: impl Into<String>) -> DiscoveryClient {
126
127
128
129
130
131
        DiscoveryClient::new(
            namespace.into(),
            self.etcd_client
                .clone()
                .expect("Attempt to get discovery_client on static DistributedRuntime"),
        )
Ryan Olson's avatar
Ryan Olson committed
132
133
134
135
136
137
    }

    pub(crate) fn service_client(&self) -> ServiceClient {
        ServiceClient::new(self.nats_client.clone())
    }

138
    pub async fn tcp_server(&self) -> Result<Arc<tcp::server::TcpStreamServer>> {
Ryan Olson's avatar
Ryan Olson committed
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
        Ok(self
            .tcp_server
            .get_or_try_init(async move {
                let options = tcp::server::ServerOptions::default();
                let server = tcp::server::TcpStreamServer::new(options).await?;
                OK(server)
            })
            .await?
            .clone())
    }

    pub fn nats_client(&self) -> nats::Client {
        self.nats_client.clone()
    }

154
    // todo(ryan): deprecate this as we move to Discovery traits and Component Identifiers
155
    pub fn etcd_client(&self) -> Option<etcd::Client> {
Ryan Olson's avatar
Ryan Olson committed
156
157
        self.etcd_client.clone()
    }
158
159
160
161

    pub fn child_token(&self) -> CancellationToken {
        self.runtime.child_token()
    }
162
163
164
165

    pub fn instance_sources(&self) -> Arc<Mutex<HashMap<Endpoint, Weak<InstanceSource>>>> {
        self.instance_sources.clone()
    }
Ryan Olson's avatar
Ryan Olson committed
166
167
168
169
170
171
}

#[derive(Dissolve)]
pub struct DistributedConfig {
    pub etcd_config: etcd::ClientOptions,
    pub nats_config: nats::ClientOptions,
172
    pub is_static: bool,
Ryan Olson's avatar
Ryan Olson committed
173
174
175
}

impl DistributedConfig {
176
    pub fn from_settings(is_static: bool) -> DistributedConfig {
Ryan Olson's avatar
Ryan Olson committed
177
178
179
        DistributedConfig {
            etcd_config: etcd::ClientOptions::default(),
            nats_config: nats::ClientOptions::default(),
180
            is_static,
Ryan Olson's avatar
Ryan Olson committed
181
182
        }
    }
Ryan Olson's avatar
Ryan Olson committed
183
184
185
186
187

    pub fn for_cli() -> DistributedConfig {
        let mut config = DistributedConfig {
            etcd_config: etcd::ClientOptions::default(),
            nats_config: nats::ClientOptions::default(),
188
            is_static: false,
Ryan Olson's avatar
Ryan Olson committed
189
190
191
192
193
194
        };

        config.etcd_config.attach_lease = false;

        config
    }
Ryan Olson's avatar
Ryan Olson committed
195
}