main.rs

// TODO: Look into how to properly resolve `clippy::result_large_err`.
// This will need changes in our and upstream error types.
#![allow(clippy::result_large_err)]
use std::sync::Arc;

use anyhow::anyhow;
use clap::Parser;
use connect::crd::{CONNECT_FULL_CONTROLLER_NAME, SparkConnectServer};
use futures::{FutureExt, StreamExt, TryFutureExt};
use history::history_controller;
use product_config::ProductConfigManager;
use stackable_operator::{
    YamlSchema,
    cli::{Command, RunArguments},
    eos::EndOfSupportChecker,
    k8s_openapi::api::{
        apps::v1::StatefulSet,
        core::v1::{ConfigMap, Pod, Service},
    },
    kube::{
        CustomResourceExt as _,
        core::DeserializeGuard,
        runtime::{
            Controller,
            events::{Recorder, Reporter},
            watcher,
        },
    },
    logging::controller::report_controller_reconciled,
    shared::yaml::SerializeOptions,
    telemetry::Tracing,
    utils::signal::{self, SignalWatcher},
};
use tracing::info_span;
use tracing_futures::Instrument;

use crate::{
    connect::crd::SparkConnectServerVersion,
    crd::{
        SparkApplication,
        constants::{
            HISTORY_FULL_CONTROLLER_NAME, OPERATOR_NAME, POD_DRIVER_FULL_CONTROLLER_NAME,
            SPARK_CONTROLLER_NAME, SPARK_FULL_CONTROLLER_NAME,
        },
        history::SparkHistoryServer,
        template_spec::{SparkApplicationTemplate, SparkApplicationTemplateVersion},
    },
    webhooks::conversion::create_webhook_server,
};

mod config;
mod connect;
mod crd;
mod history;
mod pod_driver_controller;
mod product_logging;
mod spark_k8s_controller;
mod webhooks;

mod built_info {
    include!(concat!(env!("OUT_DIR"), "/built.rs"));
}

#[derive(Parser)]
#[clap(about, author)]
struct Opts {
    #[clap(subcommand)]
    cmd: Command,
}

const PRODUCT_CONFIG_PATHS: [&str; 2] = [
    "deploy/config-spec/properties.yaml",
    "/etc/stackable/spark-k8s-operator/config-spec/properties.yaml",
];
pub struct Ctx {
    pub client: stackable_operator::client::Client,
    pub product_config: ProductConfigManager,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let opts = Opts::parse();
    match opts.cmd {
        Command::Crd => {
            SparkApplication::merged_crd(crd::SparkApplicationVersion::V1Alpha1)?
                .print_yaml_schema(built_info::PKG_VERSION, SerializeOptions::default())?;
            SparkHistoryServer::merged_crd(crd::history::SparkHistoryServerVersion::V1Alpha1)?
                .print_yaml_schema(built_info::PKG_VERSION, SerializeOptions::default())?;
            SparkConnectServer::merged_crd(SparkConnectServerVersion::V1Alpha1)?
                .print_yaml_schema(built_info::PKG_VERSION, SerializeOptions::default())?;
            SparkApplicationTemplate::merged_crd(SparkApplicationTemplateVersion::V1Alpha1)?
                .print_yaml_schema(built_info::PKG_VERSION, SerializeOptions::default())?;
        }
        Command::Run(RunArguments {
            operator_environment,
            watch_namespace,
            product_config,
            maintenance,
            common,
        }) => {
            // NOTE (@NickLarsenNZ): Before stackable-telemetry was used:
            // - The console log level was set by `SPARK_K8S_OPERATOR_LOG`, and is now `CONSOLE_LOG` (when using Tracing::pre_configured).
            // - The file log level was set by `SPARK_K8S_OPERATOR_LOG`, and is now set via `FILE_LOG` (when using Tracing::pre_configured).
            // - The file log directory was set by `SPARK_K8S_OPERATOR_LOG_DIRECTORY`, and is now set by `ROLLING_LOGS_DIR` (or via `--rolling-logs <DIRECTORY>`).
            let _tracing_guard =
                Tracing::pre_configured(built_info::PKG_NAME, common.telemetry).init()?;

            tracing::info!(
                built_info.pkg_version = built_info::PKG_VERSION,
                built_info.git_version = built_info::GIT_VERSION,
                built_info.target = built_info::TARGET,
                built_info.built_time_utc = built_info::BUILT_TIME_UTC,
                built_info.rustc_version = built_info::RUSTC_VERSION,
                "Starting {description}",
                description = built_info::PKG_DESCRIPTION
            );

            // Watches for the SIGTERM signal and sends a signal to all receivers, which gracefully
            // shuts down all concurrent tasks below (EoS checker, controller).
            let sigterm_watcher = SignalWatcher::sigterm()?;

            let eos_checker =
                EndOfSupportChecker::new(built_info::BUILT_TIME_UTC, maintenance.end_of_support)?
                    .run(sigterm_watcher.handle())
                    .map(anyhow::Ok);

            let client = stackable_operator::client::initialize_operator(
                Some(OPERATOR_NAME.to_string()),
                &common.cluster_info,
            )
            .await?;

            let webhook_server = create_webhook_server(
                &operator_environment,
                maintenance.disable_crd_maintenance,
                client.as_kube_client(),
            )
            .await?;

            let webhook_server = webhook_server
                .run(sigterm_watcher.handle())
                .map_err(|err| anyhow!(err).context("failed to run webhook server"));

            let ctx = Ctx {
                client: client.clone(),
                product_config: product_config.load(&PRODUCT_CONFIG_PATHS)?,
            };
            let spark_event_recorder = Arc::new(Recorder::new(
                client.as_kube_client(),
                Reporter {
                    controller: SPARK_FULL_CONTROLLER_NAME.to_string(),
                    instance: None,
                },
            ));
            let app_controller = Controller::new(
                watch_namespace
                    .get_api::<DeserializeGuard<crd::v1alpha1::SparkApplication>>(&client),
                watcher::Config::default(),
            )
            .owns(
                watch_namespace.get_api::<DeserializeGuard<ConfigMap>>(&client),
                watcher::Config::default(),
            )
            .graceful_shutdown_on(sigterm_watcher.handle())
            .run(
                spark_k8s_controller::reconcile,
                spark_k8s_controller::error_policy,
                Arc::new(ctx),
            )
            .instrument(info_span!("app_controller"))
            // We can let the reporting happen in the background
            .for_each_concurrent(
                16, // concurrency limit
                |result| {
                    // The event_recorder needs to be shared across all invocations, so that
                    // events are correctly aggregated
                    let spark_event_recorder = spark_event_recorder.clone();
                    async move {
                        report_controller_reconciled(
                            &spark_event_recorder,
                            SPARK_FULL_CONTROLLER_NAME,
                            &result,
                        )
                        .await;
                    }
                },
            )
            .map(anyhow::Ok);

            let pod_driver_event_recorder = Arc::new(Recorder::new(
                client.as_kube_client(),
                Reporter {
                    controller: POD_DRIVER_FULL_CONTROLLER_NAME.to_string(),
                    instance: None,
                },
            ));
            let pod_driver_controller = Controller::new(
                watch_namespace.get_api::<DeserializeGuard<Pod>>(&client),
                watcher::Config::default()
                 .labels(&format!("app.kubernetes.io/managed-by={OPERATOR_NAME}_{SPARK_CONTROLLER_NAME},spark-role=driver")),
            )
            .owns(
                watch_namespace.get_api::<DeserializeGuard<Pod>>(&client),
                watcher::Config::default(),
            )
            .graceful_shutdown_on(sigterm_watcher.handle())
            .run(
                pod_driver_controller::reconcile,
                pod_driver_controller::error_policy,
                Arc::new(client.clone()),
            )
            .instrument(info_span!("pod_driver_controller"))
            // We can let the reporting happen in the background
            .for_each_concurrent(
                16, // concurrency limit
                |result| {
                    // The event_recorder needs to be shared across all invocations, so that
                    // events are correctly aggregated
                    let pod_driver_event_recorder = pod_driver_event_recorder.clone();
                    async move {
                        report_controller_reconciled(
                            &pod_driver_event_recorder,
                            POD_DRIVER_FULL_CONTROLLER_NAME,
                            &result,
                        )
                        .await;
                    }
                },
            ).map(anyhow::Ok);

            // Create new object because Ctx cannot be cloned
            let ctx = Ctx {
                client: client.clone(),
                product_config: product_config.load(&PRODUCT_CONFIG_PATHS)?,
            };
            let history_event_recorder = Arc::new(Recorder::new(
                client.as_kube_client(),
                Reporter {
                    controller: HISTORY_FULL_CONTROLLER_NAME.to_string(),
                    instance: None,
                },
            ));
            let history_controller = Controller::new(
                watch_namespace
                    .get_api::<DeserializeGuard<crd::history::v1alpha1::SparkHistoryServer>>(
                        &client,
                    ),
                watcher::Config::default(),
            )
            .owns(
                watch_namespace
                    .get_api::<DeserializeGuard<crd::history::v1alpha1::SparkHistoryServer>>(
                        &client,
                    ),
                watcher::Config::default(),
            )
            .owns(
                watch_namespace.get_api::<DeserializeGuard<StatefulSet>>(&client),
                watcher::Config::default(),
            )
            .owns(
                watch_namespace.get_api::<DeserializeGuard<Service>>(&client),
                watcher::Config::default(),
            )
            .owns(
                watch_namespace.get_api::<DeserializeGuard<ConfigMap>>(&client),
                watcher::Config::default(),
            )
            .graceful_shutdown_on(sigterm_watcher.handle())
            .run(
                history_controller::reconcile,
                history_controller::error_policy,
                Arc::new(ctx),
            )
            .instrument(info_span!("history_controller"))
            // We can let the reporting happen in the background
            .for_each_concurrent(
                16, // concurrency limit
                |result| {
                    // The event_recorder needs to be shared across all invocations, so that
                    // events are correctly aggregated
                    let history_event_recorder = history_event_recorder.clone();
                    async move {
                        report_controller_reconciled(
                            &history_event_recorder,
                            HISTORY_FULL_CONTROLLER_NAME,
                            &result,
                        )
                        .await;
                    }
                },
            )
            .map(anyhow::Ok);

            // ==============================
            // Create new object because Ctx cannot be cloned
            let ctx = Ctx {
                client: client.clone(),
                product_config: product_config.load(&PRODUCT_CONFIG_PATHS)?,
            };
            let connect_event_recorder = Arc::new(Recorder::new(
                client.as_kube_client(),
                Reporter {
                    controller: CONNECT_FULL_CONTROLLER_NAME.to_string(),
                    instance: None,
                },
            ));
            let connect_controller = Controller::new(
                watch_namespace
                    .get_api::<DeserializeGuard<connect::crd::v1alpha1::SparkConnectServer>>(
                        &client,
                    ),
                watcher::Config::default(),
            )
            .owns(
                watch_namespace
                    .get_api::<DeserializeGuard<connect::crd::v1alpha1::SparkConnectServer>>(
                        &client,
                    ),
                watcher::Config::default(),
            )
            .owns(
                watch_namespace.get_api::<DeserializeGuard<StatefulSet>>(&client),
                watcher::Config::default(),
            )
            .owns(
                watch_namespace.get_api::<DeserializeGuard<Service>>(&client),
                watcher::Config::default(),
            )
            .owns(
                watch_namespace.get_api::<DeserializeGuard<ConfigMap>>(&client),
                watcher::Config::default(),
            )
            .graceful_shutdown_on(sigterm_watcher.handle())
            .run(
                connect::controller::reconcile,
                connect::controller::error_policy,
                Arc::new(ctx),
            )
            .instrument(info_span!("connect_controller"))
            // We can let the reporting happen in the background
            .for_each_concurrent(
                16, // concurrency limit
                |result| {
                    // The event_recorder needs to be shared across all invocations, so that
                    // events are correctly aggregated
                    let connect_event_recorder = connect_event_recorder.clone();
                    async move {
                        report_controller_reconciled(
                            &connect_event_recorder,
                            CONNECT_FULL_CONTROLLER_NAME,
                            &result,
                        )
                        .await;
                    }
                },
            )
            .map(anyhow::Ok);

            let delayed_history_controller = async {
                signal::crd_established(
                    &client,
                    crd::history::v1alpha1::SparkHistoryServer::crd_name(),
                    None,
                )
                .await?;
                history_controller.await
            };

            let delayed_connect_controller = async {
                signal::crd_established(
                    &client,
                    connect::crd::v1alpha1::SparkConnectServer::crd_name(),
                    None,
                )
                .await?;
                connect_controller.await
            };

            let delayed_app_controller = async {
                signal::crd_established(&client, crd::v1alpha1::SparkApplication::crd_name(), None)
                    .await?;
                app_controller.await
            };

            // kube-runtime's Controller will tokio::spawn each reconciliation, so this only concerns the internal watch machinery
            futures::try_join!(
                pod_driver_controller,
                delayed_history_controller,
                delayed_connect_controller,
                delayed_app_controller,
                webhook_server,
                eos_checker,
            )?;
        }
    }
    Ok(())
}