lib.rs

// 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.

//! Triton

#![allow(dead_code)]
#![allow(unused_imports)]

use std::sync::{Arc, Mutex};

pub use anyhow::{
    anyhow as error, bail as raise, Context as ErrorContext, Error, Ok as OK, Result,
};

use async_once_cell::OnceCell;

mod config;
pub use config::RuntimeConfig;

pub mod component;
pub mod discovery;
pub mod engine;
pub mod logging;
pub mod pipeline;
pub mod protocols;
pub mod runnable;
pub mod runtime;
pub mod service;
pub mod transports;
pub mod worker;

pub mod distributed;

pub use futures::stream;
pub use tokio_util::sync::CancellationToken;
pub use worker::Worker;

/// Types of Tokio runtimes that can be used to construct a Triton [Runtime].
#[derive(Clone)]
enum RuntimeType {
    Shared(Arc<tokio::runtime::Runtime>),
    External(tokio::runtime::Handle),
}

/// Local [Runtime] which provides access to shared resources local to the physical node/machine.
#[derive(Debug, Clone)]
pub struct Runtime {
    id: Arc<String>,
    primary: RuntimeType,
    secondary: Arc<tokio::runtime::Runtime>,
    cancellation_token: CancellationToken,
}

/// Distributed [Runtime] which provides access to shared resources across the cluster, this includes
/// communication protocols and transports.
#[derive(Clone)]
pub struct DistributedRuntime {
    // local runtime
    runtime: Runtime,

    // we might consider a unifed transport manager here
    etcd_client: transports::etcd::Client,
    nats_client: transports::nats::Client,
    tcp_server: Arc<OnceCell<Arc<transports::tcp::server::TcpStreamServer>>>,

    // local registry for components
    // the registry allows us to use share runtime resources across instances of the same component object.
    // take fo example two instances of a client to the same remote component. The registry allows us to use
    // a single endpoint watcher for both clients, this keeps the number background tasking watching specific
    // paths in etcd to a minimum.
    component_registry: component::Registry,
}

pub mod traits {
    use super::*;

    /// A trait for objects taht proivde access to the [Runtime]
    pub trait RuntimeProvider {
        fn rt(&self) -> &Runtime;
    }

    /// A trait for objects that provide access to the [DistributedRuntime].
    pub trait DistributedRuntimeProvider {
        fn drt(&self) -> &DistributedRuntime;
    }

    impl RuntimeProvider for DistributedRuntime {
        fn rt(&self) -> &Runtime {
            &self.runtime
        }
    }
}