From 75a28a0f5f9874dc4b985bdbe9c30b7963c8be2f Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Wed, 8 Jul 2026 15:21:29 -0400 Subject: [PATCH] Refactoring... --- components/spider-client/src/client.rs | 21 ++- components/spider-client/src/grpc/job.rs | 37 +++--- .../spider-client/src/grpc/resource_group.rs | 18 ++- components/spider-core/src/compression.rs | 3 +- components/spider-core/src/session.rs | 7 +- .../spider-core/src/{task.rs => task/mod.rs} | 0 components/spider-core/src/task/task_graph.rs | 42 +++--- .../spider-core/src/task/type_descriptor.rs | 3 +- components/spider-core/src/types/id.rs | 18 ++- components/spider-core/src/types/io.rs | 21 +-- .../src/{types.rs => types/mod.rs} | 0 components/spider-core/src/types/scheduler.rs | 7 +- components/spider-derive/src/lib.rs | 3 +- components/spider-derive/src/mysql.rs | 4 +- .../src/bin/execution_manager.rs | 16 ++- .../src/client/grpc/liveness.rs | 28 ++-- .../src/client/grpc/scheduler.rs | 21 +-- .../src/client/grpc/storage.rs | 29 ++-- .../src/client/liveness.rs | 6 +- .../src/{client.rs => client/mod.rs} | 13 +- .../src/client/scheduler.rs | 7 +- .../src/client/storage.rs | 10 +- .../spider-execution-manager/src/config.rs | 11 +- .../spider-execution-manager/src/liveness.rs | 43 +++--- .../src/process_pool.rs | 42 +++--- .../spider-execution-manager/src/runtime.rs | 53 ++++---- components/spider-proto-rust/build.rs | 4 +- .../spider-proto-rust/src/assignment.rs | 37 +++--- components/spider-proto-rust/src/id.rs | 3 +- components/spider-proto-rust/src/io.rs | 10 +- components/spider-proto-rust/src/job.rs | 6 +- components/spider-proto-rust/src/lib.rs | 26 ++++ .../src/scheduler_registration.rs | 3 +- .../spider-proto-rust/src/unpack/common.rs | 3 +- .../spider-proto-rust/src/unpack/mod.rs | 3 +- .../spider-proto-rust/src/unpack/scheduler.rs | 20 ++- .../spider-proto-rust/src/unpack/storage.rs | 49 ++++--- .../spider-scheduler/src/bin/grpc_server.rs | 19 +-- components/spider-scheduler/src/config.rs | 12 +- components/spider-scheduler/src/core.rs | 11 +- .../core_impl/round_robin/implementation.rs | 35 ++--- .../src/core_impl/round_robin/mod.rs | 3 +- .../src/core_impl/round_robin/tests.rs | 55 ++++---- .../spider-scheduler/src/dispatch_queue.rs | 31 +++-- components/spider-scheduler/src/error.rs | 3 +- .../src/execution_manager_registry.rs | 34 ++--- components/spider-scheduler/src/grpc.rs | 42 +++--- components/spider-scheduler/src/lib.rs | 34 ++--- components/spider-scheduler/src/runtime.rs | 41 +++--- components/spider-scheduler/src/service.rs | 72 +++++----- .../src/storage_client/grpc.rs | 49 +++---- .../src/storage_client/mod.rs | 11 +- components/spider-scheduler/src/types.rs | 4 +- .../spider-storage/src/bin/grpc_server.rs | 27 ++-- components/spider-storage/src/cache/error.rs | 4 +- components/spider-storage/src/cache/io.rs | 8 +- components/spider-storage/src/cache/job.rs | 53 ++++---- .../src/{cache.rs => cache/mod.rs} | 0 components/spider-storage/src/cache/sync.rs | 4 +- components/spider-storage/src/cache/task.rs | 79 +++++------ components/spider-storage/src/config.rs | 3 +- components/spider-storage/src/db.rs | 16 --- components/spider-storage/src/db/error.rs | 8 +- components/spider-storage/src/db/mariadb.rs | 51 ++++--- components/spider-storage/src/db/mod.rs | 14 ++ components/spider-storage/src/db/protocol.rs | 19 +-- components/spider-storage/src/grpc.rs | 50 +++---- .../spider-storage/src/job_submission.rs | 23 ++-- components/spider-storage/src/lib.rs | 3 +- components/spider-storage/src/ready_queue.rs | 10 +- components/spider-storage/src/state.rs | 14 -- components/spider-storage/src/state/error.rs | 6 +- .../spider-storage/src/state/job_cache.rs | 56 ++++---- .../spider-storage/src/state/job_cache_gc.rs | 41 +++--- components/spider-storage/src/state/mod.rs | 18 +++ .../spider-storage/src/state/runtime.rs | 67 +++++----- .../spider-storage/src/state/service.rs | 124 +++++++++--------- .../spider-storage/src/state/test_utils.rs | 66 +++++----- .../spider-storage/src/task_instance_pool.rs | 70 +++++----- components/spider-storage/tests/jcb_test.rs | 40 +++--- .../spider-storage/tests/mariadb_infra.rs | 7 +- .../spider-storage/tests/mariadb_test.rs | 50 ++++--- .../tests/runtime_recovery_test.rs | 59 +++++---- .../spider-storage/tests/scheduling_infra.rs | 70 +++++----- .../tests/task_graph_builder.rs | 28 ++-- .../src/bin/spider_task_executor.rs | 30 +++-- components/spider-task-executor/src/error.rs | 3 +- components/spider-task-executor/src/lib.rs | 3 +- .../spider-task-executor/src/manager.rs | 16 ++- .../spider-task-executor/src/protocol.rs | 3 +- components/spider-tdl-derive/src/lib.rs | 3 +- .../spider-tdl-derive/src/task_macro.rs | 26 ++-- components/spider-tdl/src/ffi.rs | 7 +- components/spider-tdl/src/lib.rs | 5 +- components/spider-tdl/src/task.rs | 23 ++-- components/spider-tdl/src/task_context.rs | 9 +- .../spider-tdl/tests/test_task_macro.rs | 21 ++- components/spider-utils/src/config.rs | 14 +- components/spider-utils/src/grpc/client.rs | 10 +- components/spider-utils/src/logging.rs | 3 +- components/spider-utils/src/wire.rs | 9 +- examples/huntsman/complex/tasks/src/lib.rs | 7 +- examples/huntsman/complex/types/src/lib.rs | 6 +- .../huntsman/em-runtime/tests/test_runtime.rs | 62 ++++----- .../integration-test-tasks/src/lib.rs | 7 +- .../tests/overhead_instrument.rs | 22 ++-- .../task-executor/tests/test_executor.rs | 19 ++- .../task-executor/tests/test_process_pool.rs | 38 +++--- .../huntsman/tdl-integration/tests/complex.rs | 20 ++- tests/huntsman/test-utils/src/executor.rs | 32 +++-- tests/huntsman/test-utils/src/mock.rs | 52 ++++---- tools/yscope-dev-utils | 2 +- 112 files changed, 1374 insertions(+), 1249 deletions(-) rename components/spider-core/src/{task.rs => task/mod.rs} (100%) rename components/spider-core/src/{types.rs => types/mod.rs} (100%) rename components/spider-execution-manager/src/{client.rs => client/mod.rs} (61%) rename components/spider-storage/src/{cache.rs => cache/mod.rs} (100%) delete mode 100644 components/spider-storage/src/db.rs create mode 100644 components/spider-storage/src/db/mod.rs delete mode 100644 components/spider-storage/src/state.rs create mode 100644 components/spider-storage/src/state/mod.rs diff --git a/components/spider-client/src/client.rs b/components/spider-client/src/client.rs index 42affea0e..e54820f1a 100644 --- a/components/spider-client/src/client.rs +++ b/components/spider-client/src/client.rs @@ -2,20 +2,17 @@ use std::num::NonZeroUsize; -use spider_core::{ - job::JobState, - task::TaskGraph, - types::{ - id::{JobId, ResourceGroupId}, - io::{TaskInput, TaskOutput}, - }, -}; +use spider_core::job::JobState; +use spider_core::task::TaskGraph; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::io::TaskInput; +use spider_core::types::io::TaskOutput; use tonic::transport::Endpoint; -use crate::{ - error::ClientError, - grpc::{job::JobOrchestrationClient, resource_group::ResourceGroupManagementClient}, -}; +use crate::error::ClientError; +use crate::grpc::job::JobOrchestrationClient; +use crate::grpc::resource_group::ResourceGroupManagementClient; /// User-facing client for the Spider storage gRPC services. #[derive(Debug, Clone)] diff --git a/components/spider-client/src/grpc/job.rs b/components/spider-client/src/grpc/job.rs index a430ced28..d6b353cfe 100644 --- a/components/spider-client/src/grpc/job.rs +++ b/components/spider-client/src/grpc/job.rs @@ -2,27 +2,26 @@ use std::num::NonZeroUsize; -use spider_core::{ - compression::encode_zstd_bytes, - job::JobState, - task::TaskGraph, - types::{ - id::{JobId, ResourceGroupId}, - io::{SerializedTaskOutputs, TaskInput, TaskInputsSerializer, TaskOutput}, - }, -}; -use spider_proto_rust::{ - error::Error as ProtoError, - storage::{self, job_orchestration_service_client::JobOrchestrationServiceClient}, -}; +use spider_core::compression::encode_zstd_bytes; +use spider_core::job::JobState; +use spider_core::task::TaskGraph; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::io::SerializedTaskOutputs; +use spider_core::types::io::TaskInput; +use spider_core::types::io::TaskInputsSerializer; +use spider_core::types::io::TaskOutput; +use spider_proto_rust::error::Error as ProtoError; +use spider_proto_rust::storage::JobOrchestrationServiceClient; +use spider_proto_rust::storage::{self}; use spider_utils::grpc::client::ConnectionPool; -use tonic::{ - Code, - Status, - transport::{Channel, Endpoint}, -}; +use tonic::Code; +use tonic::Status; +use tonic::transport::Channel; +use tonic::transport::Endpoint; -use crate::error::{ClientError, to_transport_error}; +use crate::error::ClientError; +use crate::error::to_transport_error; /// gRPC client for the storage server's job-orchestration service. #[derive(Debug, Clone)] diff --git a/components/spider-client/src/grpc/resource_group.rs b/components/spider-client/src/grpc/resource_group.rs index 671f2b8a9..102f6151b 100644 --- a/components/spider-client/src/grpc/resource_group.rs +++ b/components/spider-client/src/grpc/resource_group.rs @@ -3,18 +3,16 @@ use std::num::NonZeroUsize; use spider_core::types::id::ResourceGroupId; -use spider_proto_rust::storage::{ - self, - resource_group_management_service_client::ResourceGroupManagementServiceClient, -}; +use spider_proto_rust::storage::ResourceGroupManagementServiceClient; +use spider_proto_rust::storage::{self}; use spider_utils::grpc::client::ConnectionPool; -use tonic::{ - Code, - Status, - transport::{Channel, Endpoint}, -}; +use tonic::Code; +use tonic::Status; +use tonic::transport::Channel; +use tonic::transport::Endpoint; -use crate::error::{ClientError, to_transport_error}; +use crate::error::ClientError; +use crate::error::to_transport_error; /// gRPC client for the storage server's resource-group-management service. #[derive(Debug, Clone)] diff --git a/components/spider-core/src/compression.rs b/components/spider-core/src/compression.rs index df5faaeea..fbc005706 100644 --- a/components/spider-core/src/compression.rs +++ b/components/spider-core/src/compression.rs @@ -32,7 +32,8 @@ pub fn decode_zstd_bytes(compressed: &[u8]) -> Result, Error> { #[cfg(test)] mod tests { - use super::{decode_zstd_bytes, encode_zstd_bytes}; + use super::decode_zstd_bytes; + use super::encode_zstd_bytes; #[test] fn zstd_bytes_round_trip() { diff --git a/components/spider-core/src/session.rs b/components/spider-core/src/session.rs index a428e0016..0867e96a8 100644 --- a/components/spider-core/src/session.rs +++ b/components/spider-core/src/session.rs @@ -4,10 +4,9 @@ //! execution manager and the scheduler) can observe and advance a shared view of storage's current //! session id. -use std::sync::{ - Arc, - atomic::{AtomicU64, Ordering}, -}; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use crate::types::id::SessionId; diff --git a/components/spider-core/src/task.rs b/components/spider-core/src/task/mod.rs similarity index 100% rename from components/spider-core/src/task.rs rename to components/spider-core/src/task/mod.rs diff --git a/components/spider-core/src/task/task_graph.rs b/components/spider-core/src/task/task_graph.rs index 785e23d3e..7eb7fd29f 100644 --- a/components/spider-core/src/task/task_graph.rs +++ b/components/spider-core/src/task/task_graph.rs @@ -1,16 +1,22 @@ -use semver::{Version, VersionReq}; -use serde::{ - Deserialize, - Serialize, - de::{self, MapAccess, Visitor}, - ser::{SerializeMap, SerializeSeq, Serializer}, -}; -use strum::{AsRefStr, EnumCount, EnumIter, IntoEnumIterator}; - -use crate::{ - compression::{decode_zstd_bytes, encode_zstd_bytes}, - task::{DataTypeDescriptor, Error}, -}; +use semver::Version; +use semver::VersionReq; +use serde::Deserialize; +use serde::Serialize; +use serde::de::MapAccess; +use serde::de::Visitor; +use serde::de::{self}; +use serde::ser::SerializeMap; +use serde::ser::SerializeSeq; +use serde::ser::Serializer; +use strum::AsRefStr; +use strum::EnumCount; +use strum::EnumIter; +use strum::IntoEnumIterator; + +use crate::compression::decode_zstd_bytes; +use crate::compression::encode_zstd_bytes; +use crate::task::DataTypeDescriptor; +use crate::task::Error; /// A unique identifier for a task within a task graph, assigned based on insertion order. /// @@ -1078,12 +1084,10 @@ mod tests { use std::fmt::Debug; use super::*; - use crate::task::{ - DataTypeDescriptor, - IntTypeDescriptor, - MapKeyTypeDescriptor, - ValueTypeDescriptor, - }; + use crate::task::DataTypeDescriptor; + use crate::task::IntTypeDescriptor; + use crate::task::MapKeyTypeDescriptor; + use crate::task::ValueTypeDescriptor; const TEST_PACKAGE: &str = "test_pkg"; diff --git a/components/spider-core/src/task/type_descriptor.rs b/components/spider-core/src/task/type_descriptor.rs index ebe861346..ca8332db2 100644 --- a/components/spider-core/src/task/type_descriptor.rs +++ b/components/spider-core/src/task/type_descriptor.rs @@ -1,5 +1,6 @@ use non_empty_string::NonEmptyString; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; +use serde::Serialize; use crate::task::Error; diff --git a/components/spider-core/src/types/id.rs b/components/spider-core/src/types/id.rs index e500c2ce7..2e86e1d2c 100644 --- a/components/spider-core/src/types/id.rs +++ b/components/spider-core/src/types/id.rs @@ -1,10 +1,13 @@ -use std::{ - fmt::{Debug, Display}, - marker::PhantomData, -}; +use std::fmt::Debug; +use std::fmt::Display; +use std::marker::PhantomData; -use serde::{Deserialize, Deserializer, Serialize, Serializer}; -use sqlx::{Database, encode::IsNull}; +use serde::Deserialize; +use serde::Deserializer; +use serde::Serialize; +use serde::Serializer; +use sqlx::Database; +use sqlx::encode::IsNull; use crate::task::TaskIndex; @@ -224,7 +227,8 @@ pub type SignedJobId = SignedId; #[cfg(test)] mod tests { - use super::{JobId, ResourceGroupId}; + use super::JobId; + use super::ResourceGroupId; #[test] fn id_serializes_as_u64() { diff --git a/components/spider-core/src/types/io.rs b/components/spider-core/src/types/io.rs index 671dd1803..fc4ae717c 100644 --- a/components/spider-core/src/types/io.rs +++ b/components/spider-core/src/types/io.rs @@ -1,11 +1,16 @@ -use serde::{Deserialize, Serialize}; -use spider_utils::wire::{self, StreamDeserializer, WireError, WireFrameBuilder}; - -use crate::{ - compression::{self, decode_zstd_bytes, encode_zstd_bytes}, - task::{TdlContext, TimeoutPolicy}, - types::id::TaskInstanceId, -}; +use serde::Deserialize; +use serde::Serialize; +use spider_utils::wire::StreamDeserializer; +use spider_utils::wire::WireError; +use spider_utils::wire::WireFrameBuilder; +use spider_utils::wire::{self}; + +use crate::compression::decode_zstd_bytes; +use crate::compression::encode_zstd_bytes; +use crate::compression::{self}; +use crate::task::TdlContext; +use crate::task::TimeoutPolicy; +use crate::types::id::TaskInstanceId; /// Represents an input of a task. #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] diff --git a/components/spider-core/src/types.rs b/components/spider-core/src/types/mod.rs similarity index 100% rename from components/spider-core/src/types.rs rename to components/spider-core/src/types/mod.rs diff --git a/components/spider-core/src/types/scheduler.rs b/components/spider-core/src/types/scheduler.rs index 9df747015..a86f3a403 100644 --- a/components/spider-core/src/types/scheduler.rs +++ b/components/spider-core/src/types/scheduler.rs @@ -1,6 +1,11 @@ use std::net::IpAddr; -use crate::types::id::{JobId, ResourceGroupId, SchedulerId, SessionId, TaskAssignmentId, TaskId}; +use crate::types::id::JobId; +use crate::types::id::ResourceGroupId; +use crate::types::id::SchedulerId; +use crate::types::id::SessionId; +use crate::types::id::TaskAssignmentId; +use crate::types::id::TaskId; /// The currently registered scheduler endpoint. #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/components/spider-derive/src/lib.rs b/components/spider-derive/src/lib.rs index 53156591a..96d50297f 100644 --- a/components/spider-derive/src/lib.rs +++ b/components/spider-derive/src/lib.rs @@ -1,7 +1,8 @@ mod mysql; use proc_macro::TokenStream; -use syn::{DeriveInput, parse_macro_input}; +use syn::DeriveInput; +use syn::parse_macro_input; /// Derives MySQL-compatible enum support for use with [`sqlx`]. /// diff --git a/components/spider-derive/src/mysql.rs b/components/spider-derive/src/mysql.rs index eec9571fa..4e1e5f5fc 100644 --- a/components/spider-derive/src/mysql.rs +++ b/components/spider-derive/src/mysql.rs @@ -1,6 +1,8 @@ use proc_macro2::TokenStream; use quote::quote; -use syn::{Data, DataEnum, DeriveInput}; +use syn::Data; +use syn::DataEnum; +use syn::DeriveInput; pub fn derive_mysql_enum(input: &DeriveInput) -> syn::Result { let enum_type_name = &input.ident; diff --git a/components/spider-execution-manager/src/bin/execution_manager.rs b/components/spider-execution-manager/src/bin/execution_manager.rs index 442fafd6a..6d5a98251 100644 --- a/components/spider-execution-manager/src/bin/execution_manager.rs +++ b/components/spider-execution-manager/src/bin/execution_manager.rs @@ -1,14 +1,16 @@ //! Command-line entrypoint for the execution manager. -use std::{error::Error, path::PathBuf}; +use std::error::Error; +use std::path::PathBuf; use clap::Parser; -use spider_execution_manager::{ - Config, - client::grpc::{GrpcLivenessClient, GrpcSchedulerClient, GrpcStorageClient}, - runtime::Runtime, -}; -use spider_utils::{config::YamlConfig, logging::set_up_logging}; +use spider_execution_manager::Config; +use spider_execution_manager::client::grpc::GrpcLivenessClient; +use spider_execution_manager::client::grpc::GrpcSchedulerClient; +use spider_execution_manager::client::grpc::GrpcStorageClient; +use spider_execution_manager::runtime::Runtime; +use spider_utils::config::YamlConfig; +use spider_utils::logging::set_up_logging; /// Command-line arguments for the execution manager. #[derive(Debug, Parser)] diff --git a/components/spider-execution-manager/src/client/grpc/liveness.rs b/components/spider-execution-manager/src/client/grpc/liveness.rs index 612d6255b..cc8fdab03 100644 --- a/components/spider-execution-manager/src/client/grpc/liveness.rs +++ b/components/spider-execution-manager/src/client/grpc/liveness.rs @@ -1,21 +1,22 @@ //! gRPC-backed [`LivenessClient`] implementation. -use std::{net::IpAddr, num::NonZeroUsize}; +use std::net::IpAddr; +use std::num::NonZeroUsize; use async_trait::async_trait; -use spider_core::types::id::{ExecutionManagerId, SessionId}; -use spider_proto_rust::storage::{ - self, - execution_manager_liveness_service_client::ExecutionManagerLivenessServiceClient, -}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::SessionId; +use spider_proto_rust::storage::ExecutionManagerLivenessServiceClient; +use spider_proto_rust::storage::{self}; use spider_utils::grpc::client::ConnectionPool; -use tonic::{ - Code, - Status, - transport::{Channel, Endpoint}, -}; +use tonic::Code; +use tonic::Status; +use tonic::transport::Channel; +use tonic::transport::Endpoint; -use crate::client::liveness::{LivenessClient, LivenessResponseError, RegistrationResponse}; +use crate::client::liveness::LivenessClient; +use crate::client::liveness::LivenessResponseError; +use crate::client::liveness::RegistrationResponse; /// gRPC-backed [`LivenessClient`] implementation. #[derive(Debug, Clone)] @@ -143,7 +144,8 @@ mod tests { use spider_core::types::id::ExecutionManagerId; use super::*; - use crate::client::{LivenessResponseError, RegistrationResponse}; + use crate::client::LivenessResponseError; + use crate::client::RegistrationResponse; #[test] fn register_response_to_result_returns_registration() { diff --git a/components/spider-execution-manager/src/client/grpc/scheduler.rs b/components/spider-execution-manager/src/client/grpc/scheduler.rs index d2df06566..c018717b3 100644 --- a/components/spider-execution-manager/src/client/grpc/scheduler.rs +++ b/components/spider-execution-manager/src/client/grpc/scheduler.rs @@ -3,16 +3,19 @@ use std::num::NonZeroUsize; use async_trait::async_trait; -use spider_core::types::{id::ExecutionManagerId, scheduler::TaskAssignmentRecord}; -use spider_proto_rust::scheduler::{self, scheduler_service_client::SchedulerServiceClient}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::scheduler::TaskAssignmentRecord; +use spider_proto_rust::scheduler::SchedulerServiceClient; +use spider_proto_rust::scheduler::{self}; use spider_utils::grpc::client::ConnectionPool; -use tonic::{ - Code, - Status, - transport::{Channel, Endpoint}, -}; - -use crate::client::{SchedulerClient, SchedulerError, SchedulerResponse}; +use tonic::Code; +use tonic::Status; +use tonic::transport::Channel; +use tonic::transport::Endpoint; + +use crate::client::SchedulerClient; +use crate::client::SchedulerError; +use crate::client::SchedulerResponse; /// gRPC-backed [`SchedulerClient`] implementation. #[derive(Debug, Clone)] diff --git a/components/spider-execution-manager/src/client/grpc/storage.rs b/components/spider-execution-manager/src/client/grpc/storage.rs index 74b46fddb..119d03f0b 100644 --- a/components/spider-execution-manager/src/client/grpc/storage.rs +++ b/components/spider-execution-manager/src/client/grpc/storage.rs @@ -6,22 +6,23 @@ use std::num::NonZeroUsize; use async_trait::async_trait; -use spider_core::types::{ - id::{ExecutionManagerId, JobId, SessionId, TaskId, TaskInstanceId}, - io::ExecutionContext, -}; -use spider_proto_rust::{ - common, - storage::{self, task_instance_management_service_client::TaskInstanceManagementServiceClient}, -}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::ExecutionContext; +use spider_proto_rust::common; +use spider_proto_rust::storage::TaskInstanceManagementServiceClient; +use spider_proto_rust::storage::{self}; use spider_utils::grpc::client::ConnectionPool; -use tonic::{ - Code, - Status, - transport::{Channel, Endpoint}, -}; +use tonic::Code; +use tonic::Status; +use tonic::transport::Channel; +use tonic::transport::Endpoint; -use crate::client::storage::{StorageClient, StorageResponseError}; +use crate::client::storage::StorageClient; +use crate::client::storage::StorageResponseError; /// gRPC-backed [`StorageClient`] implementation. #[derive(Debug, Clone)] diff --git a/components/spider-execution-manager/src/client/liveness.rs b/components/spider-execution-manager/src/client/liveness.rs index 6a1adb2d3..9a5e955b1 100644 --- a/components/spider-execution-manager/src/client/liveness.rs +++ b/components/spider-execution-manager/src/client/liveness.rs @@ -3,10 +3,12 @@ //! The execution manager registers itself with storage at boot, then sends a periodic heartbeat. //! Each heartbeat both keeps the EM marked alive and returns storage's current session id. -use std::{net::IpAddr, sync::Arc}; +use std::net::IpAddr; +use std::sync::Arc; use async_trait::async_trait; -use spider_core::types::id::{ExecutionManagerId, SessionId}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::SessionId; /// The execution manager's identity and the storage session at registration time. #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/components/spider-execution-manager/src/client.rs b/components/spider-execution-manager/src/client/mod.rs similarity index 61% rename from components/spider-execution-manager/src/client.rs rename to components/spider-execution-manager/src/client/mod.rs index 6a37b79d9..52ab377f9 100644 --- a/components/spider-execution-manager/src/client.rs +++ b/components/spider-execution-manager/src/client/mod.rs @@ -11,8 +11,13 @@ pub mod liveness; pub mod scheduler; pub mod storage; -pub use grpc::{GrpcLivenessClient, GrpcStorageClient}; -pub use liveness::{LivenessClient, LivenessResponseError, RegistrationResponse}; -pub use scheduler::{SchedulerClient, SchedulerError}; +pub use grpc::GrpcLivenessClient; +pub use grpc::GrpcStorageClient; +pub use liveness::LivenessClient; +pub use liveness::LivenessResponseError; +pub use liveness::RegistrationResponse; +pub use scheduler::SchedulerClient; +pub use scheduler::SchedulerError; pub use spider_core::types::scheduler::SchedulerResponse; -pub use storage::{StorageClient, StorageResponseError}; +pub use storage::StorageClient; +pub use storage::StorageResponseError; diff --git a/components/spider-execution-manager/src/client/scheduler.rs b/components/spider-execution-manager/src/client/scheduler.rs index de03a77e1..1167fba80 100644 --- a/components/spider-execution-manager/src/client/scheduler.rs +++ b/components/spider-execution-manager/src/client/scheduler.rs @@ -3,10 +3,9 @@ //! The execution manager acquires tasks from the scheduler through [`SchedulerClient`]. use async_trait::async_trait; -use spider_core::types::{ - id::ExecutionManagerId, - scheduler::{SchedulerResponse, TaskAssignmentRecord}, -}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::scheduler::SchedulerResponse; +use spider_core::types::scheduler::TaskAssignmentRecord; /// Errors returned by [`SchedulerClient::next_task`]. #[derive(Debug, thiserror::Error)] diff --git a/components/spider-execution-manager/src/client/storage.rs b/components/spider-execution-manager/src/client/storage.rs index aa900bdc0..0228ab807 100644 --- a/components/spider-execution-manager/src/client/storage.rs +++ b/components/spider-execution-manager/src/client/storage.rs @@ -4,10 +4,12 @@ //! instance, fetch its [`ExecutionContext`], and report success or failure. use async_trait::async_trait; -use spider_core::types::{ - id::{ExecutionManagerId, JobId, SessionId, TaskId, TaskInstanceId}, - io::ExecutionContext, -}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::ExecutionContext; /// Errors returned by [`StorageClient`] operations. /// diff --git a/components/spider-execution-manager/src/config.rs b/components/spider-execution-manager/src/config.rs index e34e0a9af..2e94c2a46 100644 --- a/components/spider-execution-manager/src/config.rs +++ b/components/spider-execution-manager/src/config.rs @@ -1,9 +1,8 @@ -use std::{ - net::IpAddr, - num::{NonZeroU64, NonZeroUsize}, - path::PathBuf, - time::Duration, -}; +use std::net::IpAddr; +use std::num::NonZeroU64; +use std::num::NonZeroUsize; +use std::path::PathBuf; +use std::time::Duration; use serde::Deserialize; use spider_utils::config::EndpointConfig; diff --git a/components/spider-execution-manager/src/liveness.rs b/components/spider-execution-manager/src/liveness.rs index e66d91a74..600b84b52 100644 --- a/components/spider-execution-manager/src/liveness.rs +++ b/components/spider-execution-manager/src/liveness.rs @@ -9,15 +9,16 @@ use std::time::Duration; -use spider_core::{session::SessionTracker, types::id::ExecutionManagerId}; -use tokio::{ - sync::mpsc, - task::JoinHandle, - time::{Interval, MissedTickBehavior}, -}; +use spider_core::session::SessionTracker; +use spider_core::types::id::ExecutionManagerId; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tokio::time::Interval; +use tokio::time::MissedTickBehavior; use tokio_util::sync::CancellationToken; -use crate::client::{LivenessClient, LivenessResponseError}; +use crate::client::LivenessClient; +use crate::client::LivenessResponseError; /// Commands the runtime sends to the actor. #[derive(Debug)] @@ -172,23 +173,25 @@ impl LivenessActor>, diff --git a/components/spider-execution-manager/src/process_pool.rs b/components/spider-execution-manager/src/process_pool.rs index da55168ba..8251dc530 100644 --- a/components/spider-execution-manager/src/process_pool.rs +++ b/components/spider-execution-manager/src/process_pool.rs @@ -1,27 +1,33 @@ //! Process supervisor for `spider-task-executor` subprocesses. -use std::{ - fs::File, - path::PathBuf, - process::Stdio, - sync::atomic::{AtomicU64, Ordering}, - time::Duration, -}; +use std::fs::File; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; use bytes::Bytes; -use futures_util::{SinkExt, StreamExt}; -use spider_core::types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, TaskId}, - io::ExecutionContext, -}; -use spider_task_executor::protocol::{ExecutorOutcome, Request, Response}; +use futures_util::SinkExt; +use futures_util::StreamExt; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::io::ExecutionContext; +use spider_task_executor::protocol::ExecutorOutcome; +use spider_task_executor::protocol::Request; +use spider_task_executor::protocol::Response; use spider_tdl::TaskContext; use spider_utils::wire::WireError; -use tokio::{ - process::{Child, ChildStdin, ChildStdout, Command}, - sync::Mutex, -}; -use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec}; +use tokio::process::Child; +use tokio::process::ChildStdin; +use tokio::process::ChildStdout; +use tokio::process::Command; +use tokio::sync::Mutex; +use tokio_util::codec::FramedRead; +use tokio_util::codec::FramedWrite; +use tokio_util::codec::LengthDelimitedCodec; /// Pool configuration. Supplied once at construction time and never mutated. #[derive(Debug, Clone)] diff --git a/components/spider-execution-manager/src/runtime.rs b/components/spider-execution-manager/src/runtime.rs index b3e7c0d5a..1f9f72c79 100644 --- a/components/spider-execution-manager/src/runtime.rs +++ b/components/spider-execution-manager/src/runtime.rs @@ -1,30 +1,35 @@ //! Runtime — the execution manager's main loop. -use std::{collections::VecDeque, net::IpAddr, path::PathBuf, time::Duration}; - -use spider_core::{ - session::SessionTracker, - types::{ - id::{ExecutionManagerId, JobId, SessionId, TaskId, TaskInstanceId}, - io::ExecutionContext, - scheduler::TaskAssignmentRecord, - }, -}; +use std::collections::VecDeque; +use std::net::IpAddr; +use std::path::PathBuf; +use std::time::Duration; + +use spider_core::session::SessionTracker; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::ExecutionContext; +use spider_core::types::scheduler::TaskAssignmentRecord; use tokio::task::JoinHandle; -use tokio_util::sync::{CancellationToken, DropGuard}; - -use crate::{ - client::{ - LivenessClient, - LivenessResponseError, - SchedulerClient, - SchedulerResponse, - StorageClient, - StorageResponseError, - }, - liveness::{self, LivenessHandle}, - process_pool::{self, ExecuteRequest, Outcome, ProcessPool, ProcessPoolConfig}, -}; +use tokio_util::sync::CancellationToken; +use tokio_util::sync::DropGuard; + +use crate::client::LivenessClient; +use crate::client::LivenessResponseError; +use crate::client::SchedulerClient; +use crate::client::SchedulerResponse; +use crate::client::StorageClient; +use crate::client::StorageResponseError; +use crate::liveness::LivenessHandle; +use crate::liveness::{self}; +use crate::process_pool::ExecuteRequest; +use crate::process_pool::Outcome; +use crate::process_pool::ProcessPool; +use crate::process_pool::ProcessPoolConfig; +use crate::process_pool::{self}; /// Static configuration for a [`Runtime`]. Supplied once at bootstrap and never mutated. #[derive(Debug, Clone)] diff --git a/components/spider-proto-rust/build.rs b/components/spider-proto-rust/build.rs index 0462223ac..ce1534334 100644 --- a/components/spider-proto-rust/build.rs +++ b/components/spider-proto-rust/build.rs @@ -1,4 +1,6 @@ -use std::{env, fs, path::PathBuf}; +use std::env; +use std::fs; +use std::path::PathBuf; /// The environment variable that, if set, forces the build script to regenerate the protobuf code. const SPIDER_PROTO_RUST_GENERATE_FROM_SOURCE: &str = "SPIDER_PROTO_RUST_GENERATE_FROM_SOURCE"; diff --git a/components/spider-proto-rust/src/assignment.rs b/components/spider-proto-rust/src/assignment.rs index e3df93b6e..88dded6f7 100644 --- a/components/spider-proto-rust/src/assignment.rs +++ b/components/spider-proto-rust/src/assignment.rs @@ -1,19 +1,19 @@ //! Conversions between protobuf task-assignment messages and their Spider core representations. -use spider_core::types::{ - id::{JobId, ResourceGroupId, SchedulerId, TaskAssignmentId, TaskId}, - scheduler::{SchedulerResponse, TaskAssignment, TaskAssignmentRecord}, -}; - -use crate::{ - common, - error::Error, - scheduler::{ - NextTaskResponse, - TaskAssignmentRecord as ProtoTaskAssignmentRecord, - next_task_response, - }, -}; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::TaskAssignmentId; +use spider_core::types::id::TaskId; +use spider_core::types::scheduler::SchedulerResponse; +use spider_core::types::scheduler::TaskAssignment; +use spider_core::types::scheduler::TaskAssignmentRecord; + +use crate::common; +use crate::error::Error; +use crate::scheduler::NextTaskResponse; +use crate::scheduler::TaskAssignmentRecord as ProtoTaskAssignmentRecord; +use crate::scheduler::next_task_response; impl From for ProtoTaskAssignmentRecord { fn from(record: TaskAssignmentRecord) -> Self { @@ -60,10 +60,11 @@ impl TryFrom for Option { #[cfg(test)] mod tests { use super::*; - use crate::{ - common::{TaskId as ProtoTaskId, Void, task_id::Kind as ProtoTaskIdKind}, - scheduler::{SchedulerAssignment, next_task_response}, - }; + use crate::common::TaskId as ProtoTaskId; + use crate::common::Void; + use crate::common::task_id::Kind as ProtoTaskIdKind; + use crate::scheduler::SchedulerAssignment; + use crate::scheduler::next_task_response; #[test] fn assignment_record_converts_to_protocol() { diff --git a/components/spider-proto-rust/src/id.rs b/components/spider-proto-rust/src/id.rs index 9a561c180..c20dd4b31 100644 --- a/components/spider-proto-rust/src/id.rs +++ b/components/spider-proto-rust/src/id.rs @@ -2,7 +2,8 @@ use spider_core::types::id::TaskId; -use crate::{common, error::Error}; +use crate::common; +use crate::error::Error; impl From for common::TaskId { fn from(task_id: TaskId) -> Self { diff --git a/components/spider-proto-rust/src/io.rs b/components/spider-proto-rust/src/io.rs index 0a2a8d154..a7596c2f0 100644 --- a/components/spider-proto-rust/src/io.rs +++ b/components/spider-proto-rust/src/io.rs @@ -1,11 +1,11 @@ //! Conversions between protobuf I/O messages and their Spider core representations. -use spider_core::{ - task::{TdlContext, TimeoutPolicy}, - types::io::ExecutionContext, -}; +use spider_core::task::TdlContext; +use spider_core::task::TimeoutPolicy; +use spider_core::types::io::ExecutionContext; -use crate::{error::Error, storage}; +use crate::error::Error; +use crate::storage; impl TryFrom for ExecutionContext { type Error = Error; diff --git a/components/spider-proto-rust/src/job.rs b/components/spider-proto-rust/src/job.rs index 93d167cb9..48b50deb9 100644 --- a/components/spider-proto-rust/src/job.rs +++ b/components/spider-proto-rust/src/job.rs @@ -2,7 +2,8 @@ use spider_core::job::JobState; -use crate::{error::Error, storage}; +use crate::error::Error; +use crate::storage; impl From for storage::JobState { fn from(state: JobState) -> Self { @@ -39,7 +40,8 @@ impl TryFrom for JobState { mod tests { use spider_core::job::JobState; - use crate::{error::Error, storage::JobState as ProtocolJobState}; + use crate::error::Error; + use crate::storage::JobState as ProtocolJobState; #[test] fn job_state_to_protocol_converts_succeeded() { diff --git a/components/spider-proto-rust/src/lib.rs b/components/spider-proto-rust/src/lib.rs index e19242bb7..1383ef4c2 100644 --- a/components/spider-proto-rust/src/lib.rs +++ b/components/spider-proto-rust/src/lib.rs @@ -16,9 +16,35 @@ pub mod common { #[allow(clippy::all, clippy::nursery, clippy::pedantic)] pub mod scheduler { include!("generated/scheduler.rs"); + + pub use scheduler_service_client::SchedulerServiceClient; + pub use scheduler_service_server::SchedulerService; + pub use scheduler_service_server::SchedulerServiceServer; } #[allow(clippy::all, clippy::nursery, clippy::pedantic)] pub mod storage { include!("generated/storage.rs"); + + pub use execution_manager_liveness_service_client::ExecutionManagerLivenessServiceClient; + pub use execution_manager_liveness_service_server::ExecutionManagerLivenessService; + pub use execution_manager_liveness_service_server::ExecutionManagerLivenessServiceServer; + pub use inbound_queue_service_client::InboundQueueServiceClient; + pub use inbound_queue_service_server::InboundQueueService; + pub use inbound_queue_service_server::InboundQueueServiceServer; + pub use job_orchestration_service_client::JobOrchestrationServiceClient; + pub use job_orchestration_service_server::JobOrchestrationService; + pub use job_orchestration_service_server::JobOrchestrationServiceServer; + pub use resource_group_management_service_client::ResourceGroupManagementServiceClient; + pub use resource_group_management_service_server::ResourceGroupManagementService; + pub use resource_group_management_service_server::ResourceGroupManagementServiceServer; + pub use scheduler_registration_service_client::SchedulerRegistrationServiceClient; + pub use scheduler_registration_service_server::SchedulerRegistrationService; + pub use scheduler_registration_service_server::SchedulerRegistrationServiceServer; + pub use session_management_service_client::SessionManagementServiceClient; + pub use session_management_service_server::SessionManagementService; + pub use session_management_service_server::SessionManagementServiceServer; + pub use task_instance_management_service_client::TaskInstanceManagementServiceClient; + pub use task_instance_management_service_server::TaskInstanceManagementService; + pub use task_instance_management_service_server::TaskInstanceManagementServiceServer; } diff --git a/components/spider-proto-rust/src/scheduler_registration.rs b/components/spider-proto-rust/src/scheduler_registration.rs index e01403e0b..13a568171 100644 --- a/components/spider-proto-rust/src/scheduler_registration.rs +++ b/components/spider-proto-rust/src/scheduler_registration.rs @@ -18,7 +18,8 @@ impl From for storage::Scheduler { mod tests { use std::net::IpAddr; - use spider_core::types::{id::SchedulerId, scheduler::RegisteredScheduler}; + use spider_core::types::id::SchedulerId; + use spider_core::types::scheduler::RegisteredScheduler; use crate::storage; diff --git a/components/spider-proto-rust/src/unpack/common.rs b/components/spider-proto-rust/src/unpack/common.rs index ae5cd31c0..ad8bbffe8 100644 --- a/components/spider-proto-rust/src/unpack/common.rs +++ b/components/spider-proto-rust/src/unpack/common.rs @@ -7,7 +7,8 @@ use spider_core::types::id::TaskId; use tonic::Code; -use crate::{common, unpack::UnpackError}; +use crate::common; +use crate::unpack::UnpackError; /// Converts a protobuf [`common::TaskId`] into a core [`TaskId`]. /// diff --git a/components/spider-proto-rust/src/unpack/mod.rs b/components/spider-proto-rust/src/unpack/mod.rs index 8f3425c00..5d9846aaa 100644 --- a/components/spider-proto-rust/src/unpack/mod.rs +++ b/components/spider-proto-rust/src/unpack/mod.rs @@ -11,7 +11,8 @@ mod common; mod scheduler; mod storage; -use tonic::{Code, Status}; +use tonic::Code; +use tonic::Status; /// A lightweight version of [`Status`] containing only the error code and message. pub struct UnpackError { diff --git a/components/spider-proto-rust/src/unpack/scheduler.rs b/components/spider-proto-rust/src/unpack/scheduler.rs index f167ccae4..85f82b8e4 100644 --- a/components/spider-proto-rust/src/unpack/scheduler.rs +++ b/components/spider-proto-rust/src/unpack/scheduler.rs @@ -2,17 +2,15 @@ use std::time::Duration; -use spider_core::types::{id::ExecutionManagerId, scheduler::TaskAssignmentRecord}; - -use crate::{ - scheduler::{ - HeartbeatRequest, - NextTaskRequest, - ShutdownRequest, - TaskAssignmentRecord as ProtoTaskAssignmentRecord, - }, - unpack::{RequestUnpack, UnpackError}, -}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::scheduler::TaskAssignmentRecord; + +use crate::scheduler::HeartbeatRequest; +use crate::scheduler::NextTaskRequest; +use crate::scheduler::ShutdownRequest; +use crate::scheduler::TaskAssignmentRecord as ProtoTaskAssignmentRecord; +use crate::unpack::RequestUnpack; +use crate::unpack::UnpackError; /// Unpacks [`NextTaskRequest`] into a tuple containing: /// diff --git a/components/spider-proto-rust/src/unpack/storage.rs b/components/spider-proto-rust/src/unpack/storage.rs index ebeba4910..e9806a1bc 100644 --- a/components/spider-proto-rust/src/unpack/storage.rs +++ b/components/spider-proto-rust/src/unpack/storage.rs @@ -1,33 +1,30 @@ //! [`RequestUnpack`] implementations for `storage.proto` requests. -use std::{net::IpAddr, time::Duration}; - -use spider_core::types::id::{ - ExecutionManagerId, - JobId, - ResourceGroupId, - SessionId, - TaskId, - TaskInstanceId, -}; +use std::net::IpAddr; +use std::time::Duration; + +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; use tonic::Code; -use crate::{ - storage::{ - AddResourceGroupRequest, - ExecutionManagerIdRequest, - JobIdRequest, - PollReadyTasksRequest, - RegisterExecutionManagerRequest, - RegisterJobRequest, - RegisterSchedulerRequest, - RegisterTaskInstanceRequest, - ReportTaskFailureRequest, - ReportTaskSuccessRequest, - VerifyResourceGroupRequest, - }, - unpack::{RequestUnpack, UnpackError, common::unpack_task_id}, -}; +use crate::storage::AddResourceGroupRequest; +use crate::storage::ExecutionManagerIdRequest; +use crate::storage::JobIdRequest; +use crate::storage::PollReadyTasksRequest; +use crate::storage::RegisterExecutionManagerRequest; +use crate::storage::RegisterJobRequest; +use crate::storage::RegisterSchedulerRequest; +use crate::storage::RegisterTaskInstanceRequest; +use crate::storage::ReportTaskFailureRequest; +use crate::storage::ReportTaskSuccessRequest; +use crate::storage::VerifyResourceGroupRequest; +use crate::unpack::RequestUnpack; +use crate::unpack::UnpackError; +use crate::unpack::common::unpack_task_id; /// Unpacks [`RegisterJobRequest`] into a tuple containing: /// diff --git a/components/spider-scheduler/src/bin/grpc_server.rs b/components/spider-scheduler/src/bin/grpc_server.rs index 7cb2936fe..793f5454e 100644 --- a/components/spider-scheduler/src/bin/grpc_server.rs +++ b/components/spider-scheduler/src/bin/grpc_server.rs @@ -1,16 +1,17 @@ //! Command-line entrypoint for the scheduler gRPC server. -use std::{error::Error, net::SocketAddr, path::PathBuf}; +use std::error::Error; +use std::net::SocketAddr; +use std::path::PathBuf; use clap::Parser; -use spider_proto_rust::scheduler::scheduler_service_server::SchedulerServiceServer; -use spider_scheduler::{ - GrpcSchedulerStorageClient, - ServerConfig, - create_runtime, - grpc::GrpcSchedulerService, -}; -use spider_utils::{config::YamlConfig, logging::set_up_logging}; +use spider_proto_rust::scheduler::SchedulerServiceServer; +use spider_scheduler::GrpcSchedulerStorageClient; +use spider_scheduler::ServerConfig; +use spider_scheduler::create_runtime; +use spider_scheduler::grpc::GrpcSchedulerService; +use spider_utils::config::YamlConfig; +use spider_utils::logging::set_up_logging; use tokio::select; use tonic::transport::Server; diff --git a/components/spider-scheduler/src/config.rs b/components/spider-scheduler/src/config.rs index 967f13c3c..bf1caa330 100644 --- a/components/spider-scheduler/src/config.rs +++ b/components/spider-scheduler/src/config.rs @@ -5,13 +5,11 @@ use std::num::NonZeroUsize; use serde::Deserialize; use spider_utils::config::EndpointConfig; -use crate::{ - core::SchedulerCore, - core_impl::RoundRobinConfig, - dispatch_queue::DispatchQueueSink, - runtime::RuntimeConfig, - storage_client::SchedulerStorageClient, -}; +use crate::core::SchedulerCore; +use crate::core_impl::RoundRobinConfig; +use crate::dispatch_queue::DispatchQueueSink; +use crate::runtime::RuntimeConfig; +use crate::storage_client::SchedulerStorageClient; /// Top-level configuration for the scheduler gRPC server. #[derive(Clone, Debug, Deserialize)] diff --git a/components/spider-scheduler/src/core.rs b/components/spider-scheduler/src/core.rs index 13778a757..2d93783fd 100644 --- a/components/spider-scheduler/src/core.rs +++ b/components/spider-scheduler/src/core.rs @@ -1,15 +1,14 @@ //! The abstract core of a Spider scheduler. -use std::sync::{Arc, atomic::AtomicU64}; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; use async_trait::async_trait; use spider_core::types::id::TaskAssignmentId; -use crate::{ - dispatch_queue::DispatchQueueSink, - error::SchedulerError, - storage_client::SchedulerStorageClient, -}; +use crate::dispatch_queue::DispatchQueueSink; +use crate::error::SchedulerError; +use crate::storage_client::SchedulerStorageClient; /// Single-source ID issuer for creating globally unique IDs for task assignments. pub struct TaskAssignmentIdIssuer { diff --git a/components/spider-scheduler/src/core_impl/round_robin/implementation.rs b/components/spider-scheduler/src/core_impl/round_robin/implementation.rs index d0d189a39..8c02406f7 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/implementation.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/implementation.rs @@ -1,28 +1,31 @@ //! The implementation of the round-robin scheduler core. See the parent module's documentation for //! the scheduling policy and configuration. -use std::{ - collections::{HashMap, HashSet, VecDeque}, - num::{NonZeroU64, NonZeroUsize}, - time::{Duration, Instant}, -}; +use std::collections::HashMap; +use std::collections::HashSet; +use std::collections::VecDeque; +use std::num::NonZeroU64; +use std::num::NonZeroUsize; +use std::time::Duration; +use std::time::Instant; use async_trait::async_trait; use serde::Deserialize; -use spider_core::types::id::{JobId, ResourceGroupId, SessionId, TaskId}; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; use tokio::select; use tokio_util::sync::CancellationToken; -use crate::{ - DispatchQueueSink, - InboundEntry, - SchedulerCore, - SchedulerError, - SchedulerStorageClient, - StorageClientError, - TaskAssignment, - core::TaskAssignmentIdIssuer, -}; +use crate::DispatchQueueSink; +use crate::InboundEntry; +use crate::SchedulerCore; +use crate::SchedulerError; +use crate::SchedulerStorageClient; +use crate::StorageClientError; +use crate::TaskAssignment; +use crate::core::TaskAssignmentIdIssuer; /// The configuration of the round-robin scheduler core. #[derive(Clone, Debug, Deserialize)] diff --git a/components/spider-scheduler/src/core_impl/round_robin/mod.rs b/components/spider-scheduler/src/core_impl/round_robin/mod.rs index 21af05f7b..1c2893989 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/mod.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/mod.rs @@ -68,4 +68,5 @@ mod implementation; #[cfg(test)] mod tests; -pub use implementation::{RoundRobinConfig, RoundRobinCore}; +pub use implementation::RoundRobinConfig; +pub use implementation::RoundRobinCore; diff --git a/components/spider-scheduler/src/core_impl/round_robin/tests.rs b/components/spider-scheduler/src/core_impl/round_robin/tests.rs index a98228013..8a9d8fc6a 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/tests.rs @@ -1,36 +1,39 @@ //! Unit tests for the round-robin scheduler core. -use std::{ - collections::{HashMap, HashSet, VecDeque}, - num::{NonZeroU64, NonZeroUsize}, - sync::{ - Arc, - Mutex, - atomic::{AtomicU64, Ordering}, - }, - time::Duration, -}; +use std::collections::HashMap; +use std::collections::HashSet; +use std::collections::VecDeque; +use std::num::NonZeroU64; +use std::num::NonZeroUsize; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; use anyhow::bail; use async_trait::async_trait; -use spider_core::{ - job::JobState, - types::id::{JobId, ResourceGroupId, SchedulerId, SessionId, TaskId}, -}; +use spider_core::job::JobState; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; use tokio_util::sync::CancellationToken; -use super::{RoundRobinConfig, implementation::RoundRobin}; -use crate::{ - DispatchQueueSource, - InboundEntry, - SchedulerCore, - SchedulerError, - SchedulerStorageClient, - StorageClientError, - TaskAssignment, - core::TaskAssignmentIdIssuer, - dispatch_queue::{DispatchQueueReader, DispatchQueueWriter, create_dispatch_queue}, -}; +use super::RoundRobinConfig; +use super::implementation::RoundRobin; +use crate::DispatchQueueSource; +use crate::InboundEntry; +use crate::SchedulerCore; +use crate::SchedulerError; +use crate::SchedulerStorageClient; +use crate::StorageClientError; +use crate::TaskAssignment; +use crate::core::TaskAssignmentIdIssuer; +use crate::dispatch_queue::DispatchQueueReader; +use crate::dispatch_queue::DispatchQueueWriter; +use crate::dispatch_queue::create_dispatch_queue; /// The session used by tests that never bump the session. const DEFAULT_SESSION_ID: SessionId = 0; diff --git a/components/spider-scheduler/src/dispatch_queue.rs b/components/spider-scheduler/src/dispatch_queue.rs index da7cbb367..9f0e61431 100644 --- a/components/spider-scheduler/src/dispatch_queue.rs +++ b/components/spider-scheduler/src/dispatch_queue.rs @@ -1,13 +1,15 @@ //! The dispatching queue that decouples the scheduler core's placement decisions from the //! execution-manager-facing service. -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; +use std::time::Duration; use async_trait::async_trait; use spider_core::types::id::SessionId; use tokio::sync::RwLock; -use crate::{error::SchedulerError, types::TaskAssignment}; +use crate::error::SchedulerError; +use crate::types::TaskAssignment; /// The writer side of the dispatching queue used by the scheduler core. #[async_trait] @@ -197,22 +199,25 @@ struct DispatchQueueReaderInner { #[cfg(test)] mod tests { - use std::{ - collections::HashMap, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, - time::Duration, - }; + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::time::Duration; use anyhow::Result; - use dashmap::{DashMap, DashSet}; - use spider_core::types::id::{JobId, ResourceGroupId, SessionId, TaskAssignmentId, TaskId}; + use dashmap::DashMap; + use dashmap::DashSet; + use spider_core::types::id::JobId; + use spider_core::types::id::ResourceGroupId; + use spider_core::types::id::SessionId; + use spider_core::types::id::TaskAssignmentId; + use spider_core::types::id::TaskId; use tokio_util::task::TaskTracker; use super::*; - use crate::{error::SchedulerError, types::TaskAssignment}; + use crate::error::SchedulerError; + use crate::types::TaskAssignment; /// Generates a [`TaskId`] backed by a module-local monotonic counter. /// diff --git a/components/spider-scheduler/src/error.rs b/components/spider-scheduler/src/error.rs index c0d61762a..84b0f8464 100644 --- a/components/spider-scheduler/src/error.rs +++ b/components/spider-scheduler/src/error.rs @@ -1,6 +1,7 @@ //! The error types used in this crate. -use spider_core::types::id::{JobId, SessionId}; +use spider_core::types::id::JobId; +use spider_core::types::id::SessionId; use crate::execution_manager_registry::ExecutionManagerRegistryError; diff --git a/components/spider-scheduler/src/execution_manager_registry.rs b/components/spider-scheduler/src/execution_manager_registry.rs index 4937380f1..0b51238f8 100644 --- a/components/spider-scheduler/src/execution_manager_registry.rs +++ b/components/spider-scheduler/src/execution_manager_registry.rs @@ -1,18 +1,19 @@ //! Execution manager registry service. -use std::{ - collections::{HashMap, hash_map::Entry}, - num::NonZeroU64, - sync::Arc, - time::{Duration, Instant}, -}; +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::num::NonZeroU64; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; use serde::Deserialize; -use spider_core::types::id::{ExecutionManagerId, TaskAssignmentId}; -use tokio::{ - select, - sync::{Mutex, RwLock, mpsc::UnboundedSender}, -}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::TaskAssignmentId; +use tokio::select; +use tokio::sync::Mutex; +use tokio::sync::RwLock; +use tokio::sync::mpsc::UnboundedSender; use tokio_util::sync::CancellationToken; use crate::TaskAssignment; @@ -351,11 +352,12 @@ impl ExecutionManagerStateInner { mod tests { use std::collections::HashSet; - use spider_core::types::id::{JobId, ResourceGroupId, TaskId}; - use tokio::{ - sync::mpsc::{self, UnboundedReceiver}, - time::timeout, - }; + use spider_core::types::id::JobId; + use spider_core::types::id::ResourceGroupId; + use spider_core::types::id::TaskId; + use tokio::sync::mpsc::UnboundedReceiver; + use tokio::sync::mpsc::{self}; + use tokio::time::timeout; use tokio_util::task::TaskTracker; use super::*; diff --git a/components/spider-scheduler/src/grpc.rs b/components/spider-scheduler/src/grpc.rs index ce8083dee..4d519e664 100644 --- a/components/spider-scheduler/src/grpc.rs +++ b/components/spider-scheduler/src/grpc.rs @@ -1,30 +1,26 @@ //! gRPC service adapter for the scheduler service. use async_trait::async_trait; -use spider_core::types::{ - id::{SchedulerId, SessionId}, - scheduler::TaskAssignment, -}; -use spider_proto_rust::{ - common, - scheduler::{ - self, - NextTaskResponse, - SchedulerAssignment, - next_task_response, - scheduler_service_server::SchedulerService, - }, - unpack::RequestUnpack, -}; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; +use spider_core::types::scheduler::TaskAssignment; +use spider_proto_rust::common; +use spider_proto_rust::scheduler::NextTaskResponse; +use spider_proto_rust::scheduler::SchedulerAssignment; +use spider_proto_rust::scheduler::SchedulerService; +use spider_proto_rust::scheduler::next_task_response; +use spider_proto_rust::scheduler::{self}; +use spider_proto_rust::unpack::RequestUnpack; use tokio_util::sync::CancellationToken; -use tonic::{Request, Response, Status}; - -use crate::{ - dispatch_queue::DispatchQueueSource, - error::{SchedulerError, SchedulerServiceError}, - execution_manager_registry::ExecutionManagerRegistryError, - service::SchedulerServiceState, -}; +use tonic::Request; +use tonic::Response; +use tonic::Status; + +use crate::dispatch_queue::DispatchQueueSource; +use crate::error::SchedulerError; +use crate::error::SchedulerServiceError; +use crate::execution_manager_registry::ExecutionManagerRegistryError; +use crate::service::SchedulerServiceState; /// gRPC adapter over a [`SchedulerServiceState`]. /// diff --git a/components/spider-scheduler/src/lib.rs b/components/spider-scheduler/src/lib.rs index b5edc09be..d7aac238e 100644 --- a/components/spider-scheduler/src/lib.rs +++ b/components/spider-scheduler/src/lib.rs @@ -43,18 +43,22 @@ pub mod service; pub mod storage_client; pub mod types; -pub use crate::{ - config::{SchedulerConfig, ServerConfig}, - core::SchedulerCore, - dispatch_queue::{ - DispatchQueueReader, - DispatchQueueSink, - DispatchQueueSource, - DispatchQueueWriter, - }, - error::{SchedulerError, SchedulerRuntimeError, SchedulerServiceError, StorageClientError}, - runtime::{Runtime, RuntimeConfig, create_runtime}, - service::SchedulerServiceState, - storage_client::{GrpcSchedulerStorageClient, SchedulerStorageClient}, - types::{InboundEntry, TaskAssignment}, -}; +pub use crate::config::SchedulerConfig; +pub use crate::config::ServerConfig; +pub use crate::core::SchedulerCore; +pub use crate::dispatch_queue::DispatchQueueReader; +pub use crate::dispatch_queue::DispatchQueueSink; +pub use crate::dispatch_queue::DispatchQueueSource; +pub use crate::dispatch_queue::DispatchQueueWriter; +pub use crate::error::SchedulerError; +pub use crate::error::SchedulerRuntimeError; +pub use crate::error::SchedulerServiceError; +pub use crate::error::StorageClientError; +pub use crate::runtime::Runtime; +pub use crate::runtime::RuntimeConfig; +pub use crate::runtime::create_runtime; +pub use crate::service::SchedulerServiceState; +pub use crate::storage_client::GrpcSchedulerStorageClient; +pub use crate::storage_client::SchedulerStorageClient; +pub use crate::types::InboundEntry; +pub use crate::types::TaskAssignment; diff --git a/components/spider-scheduler/src/runtime.rs b/components/spider-scheduler/src/runtime.rs index 5a68a8e66..e8fe37cd0 100644 --- a/components/spider-scheduler/src/runtime.rs +++ b/components/spider-scheduler/src/runtime.rs @@ -11,16 +11,18 @@ use serde::Deserialize; use spider_core::types::id::SessionId; use tokio_util::sync::CancellationToken; -use crate::{ - config::SchedulerConfig, - core::TaskAssignmentIdIssuer, - dispatch_queue::{DispatchQueueReader, DispatchQueueWriter, create_dispatch_queue}, - error::{SchedulerError, SchedulerRuntimeError}, - execution_manager_registry::{ExecutionManagerRegistry, ExecutionManagerRegistryConfig}, - service::SchedulerServiceState, - storage_client::SchedulerStorageClient, - types::TaskAssignment, -}; +use crate::config::SchedulerConfig; +use crate::core::TaskAssignmentIdIssuer; +use crate::dispatch_queue::DispatchQueueReader; +use crate::dispatch_queue::DispatchQueueWriter; +use crate::dispatch_queue::create_dispatch_queue; +use crate::error::SchedulerError; +use crate::error::SchedulerRuntimeError; +use crate::execution_manager_registry::ExecutionManagerRegistry; +use crate::execution_manager_registry::ExecutionManagerRegistryConfig; +use crate::service::SchedulerServiceState; +use crate::storage_client::SchedulerStorageClient; +use crate::types::TaskAssignment; /// Runtime configuration for the scheduler service. #[derive(Clone, Debug, Deserialize)] @@ -175,19 +177,20 @@ const fn default_stop_timeout_sec() -> u64 { #[cfg(test)] mod tests { - use std::{ - net::{IpAddr, Ipv4Addr}, - num::{NonZeroU64, NonZeroUsize}, - }; + use std::net::IpAddr; + use std::net::Ipv4Addr; + use std::num::NonZeroU64; + use std::num::NonZeroUsize; use async_trait::async_trait; - use spider_core::{ - job::JobState, - types::id::{JobId, SchedulerId}, - }; + use spider_core::job::JobState; + use spider_core::types::id::JobId; + use spider_core::types::id::SchedulerId; use super::*; - use crate::{core_impl::RoundRobinConfig, error::StorageClientError, types::InboundEntry}; + use crate::core_impl::RoundRobinConfig; + use crate::error::StorageClientError; + use crate::types::InboundEntry; /// The scheduler identifier the mock storage client hands back from registration. const SCHEDULER_ID: u64 = 7; diff --git a/components/spider-scheduler/src/service.rs b/components/spider-scheduler/src/service.rs index 0513cf8fa..e835b175a 100644 --- a/components/spider-scheduler/src/service.rs +++ b/components/spider-scheduler/src/service.rs @@ -6,19 +6,18 @@ //! is generic over its dispatch source, so the runtime can drive it with the real dispatch queue in //! production or a mock in tests, while the [`ExecutionManagerRegistry`] is shared by value. -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; +use std::time::Duration; -use spider_core::types::{ - id::{ExecutionManagerId, SchedulerId, SessionId}, - scheduler::TaskAssignmentRecord, -}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; +use spider_core::types::scheduler::TaskAssignmentRecord; -use crate::{ - dispatch_queue::DispatchQueueSource, - error::SchedulerServiceError, - execution_manager_registry::ExecutionManagerRegistry, - types::TaskAssignment, -}; +use crate::dispatch_queue::DispatchQueueSource; +use crate::error::SchedulerServiceError; +use crate::execution_manager_registry::ExecutionManagerRegistry; +use crate::types::TaskAssignment; /// The execution-manager-facing scheduler service. /// @@ -219,41 +218,32 @@ struct SchedulerServiceStateInner #[cfg(test)] mod tests { - use std::{ - num::NonZeroU64, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, - time::Duration, - }; + use std::num::NonZeroU64; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::time::Duration; use async_trait::async_trait; - use spider_core::types::{ - id::{ - ExecutionManagerId, - JobId, - ResourceGroupId, - SchedulerId, - SessionId, - TaskAssignmentId, - TaskId, - }, - scheduler::TaskAssignmentRecord, - }; - use tokio::{ - sync::mpsc::{self, UnboundedReceiver}, - time::timeout, - }; + use spider_core::types::id::ExecutionManagerId; + use spider_core::types::id::JobId; + use spider_core::types::id::ResourceGroupId; + use spider_core::types::id::SchedulerId; + use spider_core::types::id::SessionId; + use spider_core::types::id::TaskAssignmentId; + use spider_core::types::id::TaskId; + use spider_core::types::scheduler::TaskAssignmentRecord; + use tokio::sync::mpsc::UnboundedReceiver; + use tokio::sync::mpsc::{self}; + use tokio::time::timeout; use tokio_util::sync::CancellationToken; use super::SchedulerServiceState; - use crate::{ - dispatch_queue::DispatchQueueSource, - error::SchedulerError, - execution_manager_registry::{ExecutionManagerRegistry, ExecutionManagerRegistryConfig}, - types::TaskAssignment, - }; + use crate::dispatch_queue::DispatchQueueSource; + use crate::error::SchedulerError; + use crate::execution_manager_registry::ExecutionManagerRegistry; + use crate::execution_manager_registry::ExecutionManagerRegistryConfig; + use crate::types::TaskAssignment; /// The storage session the mock dispatch source pairs with every assignment it returns. const SESSION_ID: SessionId = 7; diff --git a/components/spider-scheduler/src/storage_client/grpc.rs b/components/spider-scheduler/src/storage_client/grpc.rs index ed1d799b8..dfdcdbf92 100644 --- a/components/spider-scheduler/src/storage_client/grpc.rs +++ b/components/spider-scheduler/src/storage_client/grpc.rs @@ -1,30 +1,28 @@ //! gRPC-backed [`SchedulerStorageClient`] implementation. -use std::{num::NonZeroUsize, time::Duration}; +use std::num::NonZeroUsize; +use std::time::Duration; use async_trait::async_trait; -use spider_core::{ - job::JobState, - types::id::{JobId, ResourceGroupId, SchedulerId, SessionId, TaskId}, -}; -use spider_proto_rust::storage::{ - self, - inbound_queue_service_client::InboundQueueServiceClient, - job_orchestration_service_client::JobOrchestrationServiceClient, - scheduler_registration_service_client::SchedulerRegistrationServiceClient, -}; +use spider_core::job::JobState; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; +use spider_proto_rust::storage::InboundQueueServiceClient; +use spider_proto_rust::storage::JobOrchestrationServiceClient; +use spider_proto_rust::storage::SchedulerRegistrationServiceClient; +use spider_proto_rust::storage::{self}; use spider_utils::grpc::client::ConnectionPool; -use tonic::{ - Code, - Status, - transport::{Channel, Endpoint}, -}; - -use crate::{ - error::StorageClientError, - storage_client::SchedulerStorageClient, - types::InboundEntry, -}; +use tonic::Code; +use tonic::Status; +use tonic::transport::Channel; +use tonic::transport::Endpoint; + +use crate::error::StorageClientError; +use crate::storage_client::SchedulerStorageClient; +use crate::types::InboundEntry; /// gRPC-backed [`SchedulerStorageClient`] implementation. #[derive(Debug, Clone)] @@ -301,8 +299,11 @@ fn to_invalid_input_error(error: impl std::fmt::Display) -> StorageClientError { #[cfg(test)] mod tests { - use spider_core::types::id::{JobId, ResourceGroupId, TaskId}; - use spider_proto_rust::{common, storage}; + use spider_core::types::id::JobId; + use spider_core::types::id::ResourceGroupId; + use spider_core::types::id::TaskId; + use spider_proto_rust::common; + use spider_proto_rust::storage; use super::*; diff --git a/components/spider-scheduler/src/storage_client/mod.rs b/components/spider-scheduler/src/storage_client/mod.rs index 3f3f75246..c1a6bbaa0 100644 --- a/components/spider-scheduler/src/storage_client/mod.rs +++ b/components/spider-scheduler/src/storage_client/mod.rs @@ -3,12 +3,13 @@ use std::time::Duration; use async_trait::async_trait; -use spider_core::{ - job::JobState, - types::id::{JobId, SchedulerId, SessionId}, -}; +use spider_core::job::JobState; +use spider_core::types::id::JobId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; -use crate::{error::StorageClientError, types::InboundEntry}; +use crate::error::StorageClientError; +use crate::types::InboundEntry; pub mod grpc; diff --git a/components/spider-scheduler/src/types.rs b/components/spider-scheduler/src/types.rs index 4236ecec2..f4574ace3 100644 --- a/components/spider-scheduler/src/types.rs +++ b/components/spider-scheduler/src/types.rs @@ -1,6 +1,8 @@ //! The data types the scheduler exchanges with the storage layer and execution managers. -use spider_core::types::id::{JobId, ResourceGroupId, TaskId}; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; pub use spider_core::types::scheduler::TaskAssignment; /// A ready task drained from the storage-owned inbound queue. diff --git a/components/spider-storage/src/bin/grpc_server.rs b/components/spider-storage/src/bin/grpc_server.rs index 87e5cb8cc..79074b138 100644 --- a/components/spider-storage/src/bin/grpc_server.rs +++ b/components/spider-storage/src/bin/grpc_server.rs @@ -1,19 +1,22 @@ //! Command-line entrypoint for the storage gRPC server. -use std::{error::Error, net::SocketAddr, path::PathBuf}; +use std::error::Error; +use std::net::SocketAddr; +use std::path::PathBuf; use clap::Parser; -use spider_proto_rust::storage::{ - execution_manager_liveness_service_server::ExecutionManagerLivenessServiceServer, - inbound_queue_service_server::InboundQueueServiceServer, - job_orchestration_service_server::JobOrchestrationServiceServer, - resource_group_management_service_server::ResourceGroupManagementServiceServer, - scheduler_registration_service_server::SchedulerRegistrationServiceServer, - session_management_service_server::SessionManagementServiceServer, - task_instance_management_service_server::TaskInstanceManagementServiceServer, -}; -use spider_storage::{ServerConfig, grpc::GrpcServiceState, state::runtime::create_runtime}; -use spider_utils::{config::YamlConfig, logging::set_up_logging}; +use spider_proto_rust::storage::ExecutionManagerLivenessServiceServer; +use spider_proto_rust::storage::InboundQueueServiceServer; +use spider_proto_rust::storage::JobOrchestrationServiceServer; +use spider_proto_rust::storage::ResourceGroupManagementServiceServer; +use spider_proto_rust::storage::SchedulerRegistrationServiceServer; +use spider_proto_rust::storage::SessionManagementServiceServer; +use spider_proto_rust::storage::TaskInstanceManagementServiceServer; +use spider_storage::ServerConfig; +use spider_storage::grpc::GrpcServiceState; +use spider_storage::state::runtime::create_runtime; +use spider_utils::config::YamlConfig; +use spider_utils::logging::set_up_logging; use tokio::select; use tonic::transport::Server; diff --git a/components/spider-storage/src/cache/error.rs b/components/spider-storage/src/cache/error.rs index 9cf3f4cac..25a6e6ac3 100644 --- a/components/spider-storage/src/cache/error.rs +++ b/components/spider-storage/src/cache/error.rs @@ -1,4 +1,6 @@ -use spider_core::{job::JobState, task::TaskState, types::id::JobId}; +use spider_core::job::JobState; +use spider_core::task::TaskState; +use spider_core::types::id::JobId; use spider_utils::wire::WireError; use crate::db::DbError; diff --git a/components/spider-storage/src/cache/io.rs b/components/spider-storage/src/cache/io.rs index bd01daeb1..c48fda2b9 100644 --- a/components/spider-storage/src/cache/io.rs +++ b/components/spider-storage/src/cache/io.rs @@ -1,9 +1,9 @@ use spider_core::types::io::TaskInput; -use crate::cache::{ - error::{CacheError, InternalError}, - sync::{Reader, Writer}, -}; +use crate::cache::error::CacheError; +use crate::cache::error::InternalError; +use crate::cache::sync::Reader; +use crate::cache::sync::Writer; /// Spider channel implementation placeholder. pub struct Channel {} diff --git a/components/spider-storage/src/cache/job.rs b/components/spider-storage/src/cache/job.rs index c5373297d..10d6043b6 100644 --- a/components/spider-storage/src/cache/job.rs +++ b/components/spider-storage/src/cache/job.rs @@ -1,31 +1,32 @@ -use std::{ - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, - time::{Duration, SystemTime}, -}; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::SystemTime; -use spider_core::{ - job::JobState, - task::{TaskIndex, TaskState}, - types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, TaskId, TaskInstanceId}, - io::{ExecutionContext, TaskOutput}, - }, -}; -use tokio::sync::{RwLockReadGuard, RwLockWriteGuard}; +use spider_core::job::JobState; +use spider_core::task::TaskIndex; +use spider_core::task::TaskState; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::ExecutionContext; +use spider_core::types::io::TaskOutput; +use tokio::sync::RwLockReadGuard; +use tokio::sync::RwLockWriteGuard; -use crate::{ - cache::{ - error::{CacheError, InternalError, StaleStateError}, - task::TaskGraph, - }, - db::{InternalJobOrchestration, RecoverableJobContext}, - job_submission::ValidatedJobSubmission, - ready_queue::ReadyQueueSender, - task_instance_pool::{TaskInstanceMetadata, TaskInstancePoolConnector}, -}; +use crate::cache::error::CacheError; +use crate::cache::error::InternalError; +use crate::cache::error::StaleStateError; +use crate::cache::task::TaskGraph; +use crate::db::InternalJobOrchestration; +use crate::db::RecoverableJobContext; +use crate::job_submission::ValidatedJobSubmission; +use crate::ready_queue::ReadyQueueSender; +use crate::task_instance_pool::TaskInstanceMetadata; +use crate::task_instance_pool::TaskInstancePoolConnector; /// A shareable control block for a job. /// diff --git a/components/spider-storage/src/cache.rs b/components/spider-storage/src/cache/mod.rs similarity index 100% rename from components/spider-storage/src/cache.rs rename to components/spider-storage/src/cache/mod.rs diff --git a/components/spider-storage/src/cache/sync.rs b/components/spider-storage/src/cache/sync.rs index 0fc034488..6d99ebd74 100644 --- a/components/spider-storage/src/cache/sync.rs +++ b/components/spider-storage/src/cache/sync.rs @@ -1,6 +1,8 @@ use std::sync::Arc; -use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}; +use tokio::sync::RwLock; +use tokio::sync::RwLockReadGuard; +use tokio::sync::RwLockWriteGuard; /// Reader-writer lock for shared data in the cache. pub type SharedRw = Arc>; diff --git a/components/spider-storage/src/cache/task.rs b/components/spider-storage/src/cache/task.rs index 87676a9f9..26efd0535 100644 --- a/components/spider-storage/src/cache/task.rs +++ b/components/spider-storage/src/cache/task.rs @@ -1,26 +1,32 @@ -use std::{ - cmp::max, - collections::{HashMap, HashSet}, - sync::Arc, -}; - -use spider_core::{ - task::{Task, TaskIndex, TaskState, TdlContext, TerminationTaskDescriptor, TimeoutPolicy}, - types::{ - id::TaskInstanceId, - io::{ExecutionContext, TaskInput, TaskInputsSerializer, TaskOutput}, - }, -}; +use std::cmp::max; +use std::collections::HashMap; +use std::collections::HashSet; +use std::sync::Arc; + +use spider_core::task::Task; +use spider_core::task::TaskIndex; +use spider_core::task::TaskState; +use spider_core::task::TdlContext; +use spider_core::task::TerminationTaskDescriptor; +use spider_core::task::TimeoutPolicy; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::ExecutionContext; +use spider_core::types::io::TaskInput; +use spider_core::types::io::TaskInputsSerializer; +use spider_core::types::io::TaskOutput; use tokio::sync::RwLock; -use crate::{ - cache::{ - error::{CacheError, InternalError, StaleStateError}, - io::{InputReader, OutputReader, OutputWriter, ValuePayload}, - sync::{Reader, SharedRw, Writer}, - }, - job_submission::ValidatedJobSubmission, -}; +use crate::cache::error::CacheError; +use crate::cache::error::InternalError; +use crate::cache::error::StaleStateError; +use crate::cache::io::InputReader; +use crate::cache::io::OutputReader; +use crate::cache::io::OutputWriter; +use crate::cache::io::ValuePayload; +use crate::cache::sync::Reader; +use crate::cache::sync::SharedRw; +use crate::cache::sync::Writer; +use crate::job_submission::ValidatedJobSubmission; /// Represents the task graph in the cache as a collection of TCBs. pub struct TaskGraph { @@ -961,23 +967,20 @@ struct TerminationTaskControlBlock { #[cfg(test)] mod tests { - use std::{ - future::Future, - hash::{BuildHasher, Hasher, RandomState}, - sync::{ - Arc, - atomic::{AtomicU64, Ordering}, - }, - }; - - use spider_core::task::{ - DataTypeDescriptor, - ExecutionPolicy, - TaskDescriptor, - TaskGraph as SubmittedTaskGraph, - TerminationTaskDescriptor, - ValueTypeDescriptor, - }; + use std::future::Future; + use std::hash::BuildHasher; + use std::hash::Hasher; + use std::hash::RandomState; + use std::sync::Arc; + use std::sync::atomic::AtomicU64; + use std::sync::atomic::Ordering; + + use spider_core::task::DataTypeDescriptor; + use spider_core::task::ExecutionPolicy; + use spider_core::task::TaskDescriptor; + use spider_core::task::TaskGraph as SubmittedTaskGraph; + use spider_core::task::TerminationTaskDescriptor; + use spider_core::task::ValueTypeDescriptor; use spider_utils::wire::unframe; use super::*; diff --git a/components/spider-storage/src/config.rs b/components/spider-storage/src/config.rs index fad46b4eb..759930e31 100644 --- a/components/spider-storage/src/config.rs +++ b/components/spider-storage/src/config.rs @@ -1,7 +1,8 @@ use std::net::IpAddr; use secrecy::SecretString; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; +use serde::Serialize; use crate::state::runtime::RuntimeConfig; diff --git a/components/spider-storage/src/db.rs b/components/spider-storage/src/db.rs deleted file mode 100644 index d56c05005..000000000 --- a/components/spider-storage/src/db.rs +++ /dev/null @@ -1,16 +0,0 @@ -mod error; -mod mariadb; -mod protocol; - -pub use error::DbError; -pub use mariadb::MariaDbStorageConnector; -pub use protocol::{ - DbStorage, - ExecutionManagerLivenessManagement, - ExternalJobOrchestration, - InternalJobOrchestration, - RecoverableJobContext, - ResourceGroupManagement, - SchedulerRegistrationManagement, - SessionManagement, -}; diff --git a/components/spider-storage/src/db/error.rs b/components/spider-storage/src/db/error.rs index 3bce5386e..bb5193391 100644 --- a/components/spider-storage/src/db/error.rs +++ b/components/spider-storage/src/db/error.rs @@ -1,9 +1,9 @@ use std::fmt::Display; -use spider_core::{ - job::JobState, - types::id::{ExecutionManagerId, JobId, ResourceGroupId}, -}; +use spider_core::job::JobState; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; #[derive(thiserror::Error, Debug)] pub enum DbError { diff --git a/components/spider-storage/src/db/mariadb.rs b/components/spider-storage/src/db/mariadb.rs index 513516681..913a1b119 100644 --- a/components/spider-storage/src/db/mariadb.rs +++ b/components/spider-storage/src/db/mariadb.rs @@ -3,33 +3,32 @@ use std::net::IpAddr; use async_trait::async_trait; use const_format::formatcp; use secrecy::ExposeSecret; -use spider_core::{ - job::JobState, - types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, SchedulerId, SessionId}, - io::{SerializedTaskOutputs, TaskOutput}, - scheduler::RegisteredScheduler, - }, -}; +use spider_core::job::JobState; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; +use spider_core::types::io::SerializedTaskOutputs; +use spider_core::types::io::TaskOutput; +use spider_core::types::scheduler::RegisteredScheduler; use spider_derive::MySqlEnum; -use sqlx::{Connection, MySqlPool, mysql::MySqlDatabaseError}; - -use crate::{ - config::DatabaseConfig, - db::{ - DbError, - DbStorage, - ExecutionManagerLivenessManagement, - ExternalJobOrchestration, - InternalJobOrchestration, - RecoverableJobContext, - ResourceGroupManagement, - SchedulerRegistrationManagement, - SessionManagement, - error::ExpectedStates, - }, - job_submission::ValidatedJobSubmission, -}; +use sqlx::Connection; +use sqlx::MySqlPool; +use sqlx::mysql::MySqlDatabaseError; + +use crate::config::DatabaseConfig; +use crate::db::DbError; +use crate::db::DbStorage; +use crate::db::ExecutionManagerLivenessManagement; +use crate::db::ExternalJobOrchestration; +use crate::db::InternalJobOrchestration; +use crate::db::RecoverableJobContext; +use crate::db::ResourceGroupManagement; +use crate::db::SchedulerRegistrationManagement; +use crate::db::SessionManagement; +use crate::db::error::ExpectedStates; +use crate::job_submission::ValidatedJobSubmission; /// A cloneable storage connector for `MariaDB` database that implements Spider's DB protocols. #[derive(Clone)] diff --git a/components/spider-storage/src/db/mod.rs b/components/spider-storage/src/db/mod.rs new file mode 100644 index 000000000..2ddb6cda4 --- /dev/null +++ b/components/spider-storage/src/db/mod.rs @@ -0,0 +1,14 @@ +mod error; +mod mariadb; +mod protocol; + +pub use error::DbError; +pub use mariadb::MariaDbStorageConnector; +pub use protocol::DbStorage; +pub use protocol::ExecutionManagerLivenessManagement; +pub use protocol::ExternalJobOrchestration; +pub use protocol::InternalJobOrchestration; +pub use protocol::RecoverableJobContext; +pub use protocol::ResourceGroupManagement; +pub use protocol::SchedulerRegistrationManagement; +pub use protocol::SessionManagement; diff --git a/components/spider-storage/src/db/protocol.rs b/components/spider-storage/src/db/protocol.rs index 01fbc60fb..c5bb09331 100644 --- a/components/spider-storage/src/db/protocol.rs +++ b/components/spider-storage/src/db/protocol.rs @@ -1,16 +1,17 @@ use std::net::IpAddr; use async_trait::async_trait; -use spider_core::{ - job::JobState, - types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, SchedulerId, SessionId}, - io::TaskOutput, - scheduler::RegisteredScheduler, - }, -}; +use spider_core::job::JobState; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; +use spider_core::types::io::TaskOutput; +use spider_core::types::scheduler::RegisteredScheduler; -use crate::{db::error::DbError, job_submission::ValidatedJobSubmission}; +use crate::db::error::DbError; +use crate::job_submission::ValidatedJobSubmission; /// A job persisted in the database that should be rebuilt in the storage cache on startup. /// diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index 7d44a152c..03ff724e8 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -1,32 +1,32 @@ //! gRPC service adapters for the storage service. use async_trait::async_trait; -use spider_core::types::{id::TaskId, io::SerializedTaskOutputs}; -use spider_proto_rust::{ - common, - storage::{ - self, - SchedulerRegistration, - execution_manager_liveness_service_server::ExecutionManagerLivenessService, - inbound_queue_service_server::InboundQueueService, - job_orchestration_service_server::JobOrchestrationService, - resource_group_management_service_server::ResourceGroupManagementService, - scheduler_registration_service_server::SchedulerRegistrationService, - session_management_service_server::SessionManagementService, - task_instance_management_service_server::TaskInstanceManagementService, - }, - unpack::RequestUnpack, -}; +use spider_core::types::id::TaskId; +use spider_core::types::io::SerializedTaskOutputs; +use spider_proto_rust::common; +use spider_proto_rust::storage::ExecutionManagerLivenessService; +use spider_proto_rust::storage::InboundQueueService; +use spider_proto_rust::storage::JobOrchestrationService; +use spider_proto_rust::storage::ResourceGroupManagementService; +use spider_proto_rust::storage::SchedulerRegistration; +use spider_proto_rust::storage::SchedulerRegistrationService; +use spider_proto_rust::storage::SessionManagementService; +use spider_proto_rust::storage::TaskInstanceManagementService; +use spider_proto_rust::storage::{self}; +use spider_proto_rust::unpack::RequestUnpack; use tokio_util::sync::CancellationToken; -use tonic::{Request, Response, Status}; - -use crate::{ - cache::error::CacheError, - db::{DbError, DbStorage}, - ready_queue::{ReadyQueueEntry, ReadyQueueSender}, - state::{ServiceState, StorageServerError}, - task_instance_pool::TaskInstancePoolConnector, -}; +use tonic::Request; +use tonic::Response; +use tonic::Status; + +use crate::cache::error::CacheError; +use crate::db::DbError; +use crate::db::DbStorage; +use crate::ready_queue::ReadyQueueEntry; +use crate::ready_queue::ReadyQueueSender; +use crate::state::ServiceState; +use crate::state::StorageServerError; +use crate::task_instance_pool::TaskInstancePoolConnector; /// gRPC adapter over a storage [`ServiceState`]. /// diff --git a/components/spider-storage/src/job_submission.rs b/components/spider-storage/src/job_submission.rs index 90faa9c76..1f746558f 100644 --- a/components/spider-storage/src/job_submission.rs +++ b/components/spider-storage/src/job_submission.rs @@ -1,8 +1,7 @@ -use spider_core::{ - compression::decode_zstd_bytes, - task::{self, TaskGraph}, - types::io::TaskInput, -}; +use spider_core::compression::decode_zstd_bytes; +use spider_core::task::TaskGraph; +use spider_core::task::{self}; +use spider_core::types::io::TaskInput; use spider_utils::wire; /// Errors produced while constructing a [`ValidatedJobSubmission`] from its compressed @@ -207,14 +206,12 @@ pub fn create_validated_submission( #[cfg(test)] mod tests { - use spider_core::task::{ - DataTypeDescriptor, - ExecutionPolicy, - TaskDescriptor, - TaskGraph as SubmittedTaskGraph, - TdlContext, - ValueTypeDescriptor, - }; + use spider_core::task::DataTypeDescriptor; + use spider_core::task::ExecutionPolicy; + use spider_core::task::TaskDescriptor; + use spider_core::task::TaskGraph as SubmittedTaskGraph; + use spider_core::task::TdlContext; + use spider_core::task::ValueTypeDescriptor; use super::*; diff --git a/components/spider-storage/src/lib.rs b/components/spider-storage/src/lib.rs index 523e770b9..6921b1e04 100644 --- a/components/spider-storage/src/lib.rs +++ b/components/spider-storage/src/lib.rs @@ -7,4 +7,5 @@ pub mod ready_queue; pub mod state; pub mod task_instance_pool; -pub use config::{DatabaseConfig, ServerConfig}; +pub use config::DatabaseConfig; +pub use config::ServerConfig; diff --git a/components/spider-storage/src/ready_queue.rs b/components/spider-storage/src/ready_queue.rs index db453fb71..14fc3f612 100644 --- a/components/spider-storage/src/ready_queue.rs +++ b/components/spider-storage/src/ready_queue.rs @@ -11,13 +11,13 @@ use std::time::Duration; -use async_channel::{Receiver, Sender}; +use async_channel::Receiver; +use async_channel::Sender; use async_trait::async_trait; use serde::Deserialize; -use spider_core::{ - task::TaskIndex, - types::id::{JobId, ResourceGroupId}, -}; +use spider_core::task::TaskIndex; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; use crate::cache::error::InternalError; diff --git a/components/spider-storage/src/state.rs b/components/spider-storage/src/state.rs deleted file mode 100644 index 7e83c191c..000000000 --- a/components/spider-storage/src/state.rs +++ /dev/null @@ -1,14 +0,0 @@ -pub mod error; -pub mod job_cache; -pub mod job_cache_gc; -pub mod runtime; -pub mod service; - -pub use error::StorageServerError; -pub use job_cache::JobCache; -pub use job_cache_gc::{JobCacheGcConfig, JobCacheGcHandle, create_job_cache_gc}; -pub use runtime::{Runtime, create_runtime}; -pub use service::{ServiceState, ServiceStateParams}; - -#[cfg(test)] -mod test_utils; diff --git a/components/spider-storage/src/state/error.rs b/components/spider-storage/src/state/error.rs index fa63b8e7e..b9641aeea 100644 --- a/components/spider-storage/src/state/error.rs +++ b/components/spider-storage/src/state/error.rs @@ -1,7 +1,9 @@ -use spider_core::types::id::{JobId, SessionId}; +use spider_core::types::id::JobId; +use spider_core::types::id::SessionId; use spider_tdl::error::TdlError; -use crate::{cache::error::CacheError, db::DbError}; +use crate::cache::error::CacheError; +use crate::db::DbError; /// Errors that can occur during storage server operations. #[derive(thiserror::Error, Debug)] diff --git a/components/spider-storage/src/state/job_cache.rs b/components/spider-storage/src/state/job_cache.rs index c01a2d74c..298b641c0 100644 --- a/components/spider-storage/src/state/job_cache.rs +++ b/components/spider-storage/src/state/job_cache.rs @@ -1,21 +1,17 @@ -use std::{ - collections::{HashMap, hash_map::Entry}, - sync::Arc, -}; +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::sync::Arc; use spider_core::types::id::JobId; use tokio::sync::RwLock; -use crate::{ - cache::{ - error::{CacheError, InternalError}, - job::SharedJobControlBlock, - }, - db::InternalJobOrchestration, - ready_queue::ReadyQueueSender, - state::StorageServerError, - task_instance_pool::TaskInstancePoolConnector, -}; +use crate::cache::error::CacheError; +use crate::cache::error::InternalError; +use crate::cache::job::SharedJobControlBlock; +use crate::db::InternalJobOrchestration; +use crate::ready_queue::ReadyQueueSender; +use crate::state::StorageServerError; +use crate::task_instance_pool::TaskInstancePoolConnector; /// An in-memory cache for job control blocks. /// @@ -157,25 +153,23 @@ type SharedJobMap; diff --git a/components/spider-storage/src/state/mod.rs b/components/spider-storage/src/state/mod.rs new file mode 100644 index 000000000..71f9b9cd5 --- /dev/null +++ b/components/spider-storage/src/state/mod.rs @@ -0,0 +1,18 @@ +pub mod error; +pub mod job_cache; +pub mod job_cache_gc; +pub mod runtime; +pub mod service; + +pub use error::StorageServerError; +pub use job_cache::JobCache; +pub use job_cache_gc::JobCacheGcConfig; +pub use job_cache_gc::JobCacheGcHandle; +pub use job_cache_gc::create_job_cache_gc; +pub use runtime::Runtime; +pub use runtime::create_runtime; +pub use service::ServiceState; +pub use service::ServiceStateParams; + +#[cfg(test)] +mod test_utils; diff --git a/components/spider-storage/src/state/runtime.rs b/components/spider-storage/src/state/runtime.rs index 23f425ba7..bada722d0 100644 --- a/components/spider-storage/src/state/runtime.rs +++ b/components/spider-storage/src/state/runtime.rs @@ -4,29 +4,27 @@ use serde::Deserialize; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; -use crate::{ - cache::{ - error::{CacheError, InternalError}, - job::SharedJobControlBlock, - }, - config::DatabaseConfig, - db::{DbStorage, MariaDbStorageConnector, SessionManagement}, - ready_queue::{ReadyQueueConfig, ReadyQueueSender, ReadyQueueSenderHandle, create_ready_queue}, - state::{ - JobCache, - JobCacheGcConfig, - ServiceState, - ServiceStateParams, - StorageServerError, - create_job_cache_gc, - }, - task_instance_pool::{ - TaskInstancePoolConfig, - TaskInstancePoolConnector, - TaskInstancePoolHandle, - create_task_instance_pool, - }, -}; +use crate::cache::error::CacheError; +use crate::cache::error::InternalError; +use crate::cache::job::SharedJobControlBlock; +use crate::config::DatabaseConfig; +use crate::db::DbStorage; +use crate::db::MariaDbStorageConnector; +use crate::db::SessionManagement; +use crate::ready_queue::ReadyQueueConfig; +use crate::ready_queue::ReadyQueueSender; +use crate::ready_queue::ReadyQueueSenderHandle; +use crate::ready_queue::create_ready_queue; +use crate::state::JobCache; +use crate::state::JobCacheGcConfig; +use crate::state::ServiceState; +use crate::state::ServiceStateParams; +use crate::state::StorageServerError; +use crate::state::create_job_cache_gc; +use crate::task_instance_pool::TaskInstancePoolConfig; +use crate::task_instance_pool::TaskInstancePoolConnector; +use crate::task_instance_pool::TaskInstancePoolHandle; +use crate::task_instance_pool::create_task_instance_pool; /// Runtime configuration for the storage service. #[derive(Clone, Debug, Deserialize)] @@ -258,18 +256,17 @@ mod tests { use tokio_util::sync::CancellationToken; use super::*; - use crate::{ - cache::error::InternalError, - db::SessionManagement, - ready_queue::{ReadyQueueConfig, ReadyQueueSenderHandle, create_ready_queue}, - state::{ - JobCache, - ServiceState, - ServiceStateParams, - StorageServerError, - test_utils::{MockDbConnector, MockTaskInstancePoolConnector}, - }, - }; + use crate::cache::error::InternalError; + use crate::db::SessionManagement; + use crate::ready_queue::ReadyQueueConfig; + use crate::ready_queue::ReadyQueueSenderHandle; + use crate::ready_queue::create_ready_queue; + use crate::state::JobCache; + use crate::state::ServiceState; + use crate::state::ServiceStateParams; + use crate::state::StorageServerError; + use crate::state::test_utils::MockDbConnector; + use crate::state::test_utils::MockTaskInstancePoolConnector; type TestServerRuntime = Runtime; diff --git a/components/spider-storage/src/state/service.rs b/components/spider-storage/src/state/service.rs index 4ea7e9a64..73908869d 100644 --- a/components/spider-storage/src/state/service.rs +++ b/components/spider-storage/src/state/service.rs @@ -1,42 +1,37 @@ -use std::{net::IpAddr, sync::Arc, time::Duration}; - -use spider_core::{ - job::JobState, - task::TaskIndex, - types::{ - id::{ - ExecutionManagerId, - JobId, - ResourceGroupId, - SchedulerId, - SessionId, - TaskId, - TaskInstanceId, - }, - io::{ExecutionContext, TaskOutput, TaskOutputsSerializer}, - scheduler::RegisteredScheduler, - }, -}; +use std::net::IpAddr; +use std::sync::Arc; +use std::time::Duration; + +use spider_core::job::JobState; +use spider_core::task::TaskIndex; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::ExecutionContext; +use spider_core::types::io::TaskOutput; +use spider_core::types::io::TaskOutputsSerializer; +use spider_core::types::scheduler::RegisteredScheduler; use spider_tdl::error::TdlError; use tokio_util::sync::CancellationToken; -use crate::{ - cache::{ - error::{CacheError, InternalError}, - job::SharedJobControlBlock, - }, - db::DbStorage, - job_submission::ValidatedJobSubmission, - ready_queue::{ - CleanupTaskMarker, - CommitTaskMarker, - ReadyQueueEntry, - ReadyQueueReceiverHandle, - ReadyQueueSender, - }, - state::{JobCache, JobCacheGcHandle, StorageServerError}, - task_instance_pool::TaskInstancePoolConnector, -}; +use crate::cache::error::CacheError; +use crate::cache::error::InternalError; +use crate::cache::job::SharedJobControlBlock; +use crate::db::DbStorage; +use crate::job_submission::ValidatedJobSubmission; +use crate::ready_queue::CleanupTaskMarker; +use crate::ready_queue::CommitTaskMarker; +use crate::ready_queue::ReadyQueueEntry; +use crate::ready_queue::ReadyQueueReceiverHandle; +use crate::ready_queue::ReadyQueueSender; +use crate::state::JobCache; +use crate::state::JobCacheGcHandle; +use crate::state::StorageServerError; +use crate::task_instance_pool::TaskInstancePoolConnector; /// Bundle of constructor parameters for [`ServiceState::new`]. /// @@ -811,35 +806,32 @@ struct ServiceStateInner< #[cfg(test)] mod tests { - use spider_core::{ - compression::encode_zstd_bytes, - job::JobState, - task::{ - DataTypeDescriptor, - ExecutionPolicy, - TaskDescriptor, - TaskGraph as SubmittedTaskGraph, - TdlContext, - ValueTypeDescriptor, - }, - types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId}, - io::{TaskInput, TaskOutput}, - }, - }; + use spider_core::compression::encode_zstd_bytes; + use spider_core::job::JobState; + use spider_core::task::DataTypeDescriptor; + use spider_core::task::ExecutionPolicy; + use spider_core::task::TaskDescriptor; + use spider_core::task::TaskGraph as SubmittedTaskGraph; + use spider_core::task::TdlContext; + use spider_core::task::ValueTypeDescriptor; + use spider_core::types::id::ExecutionManagerId; + use spider_core::types::id::JobId; + use spider_core::types::id::ResourceGroupId; + use spider_core::types::io::TaskInput; + use spider_core::types::io::TaskOutput; use super::*; - use crate::{ - cache::job::SharedJobControlBlock, - db::DbError, - job_submission::{compress_job_inputs, compress_task_graph, create_validated_submission}, - ready_queue::ReadyQueueSenderHandle, - state::{ - JobCacheGcHandle, - StorageServerError, - test_utils::{MockDbConnector, MockReadyQueueSender, MockTaskInstancePoolConnector}, - }, - }; + use crate::cache::job::SharedJobControlBlock; + use crate::db::DbError; + use crate::job_submission::compress_job_inputs; + use crate::job_submission::compress_task_graph; + use crate::job_submission::create_validated_submission; + use crate::ready_queue::ReadyQueueSenderHandle; + use crate::state::JobCacheGcHandle; + use crate::state::StorageServerError; + use crate::state::test_utils::MockDbConnector; + use crate::state::test_utils::MockReadyQueueSender; + use crate::state::test_utils::MockTaskInstancePoolConnector; type TestServiceState = ServiceState; @@ -874,7 +866,8 @@ mod tests { } fn create_ready_queue_receiver() -> ReadyQueueReceiverHandle { - use crate::ready_queue::{ReadyQueueConfig, create_ready_queue}; + use crate::ready_queue::ReadyQueueConfig; + use crate::ready_queue::create_ready_queue; let (_sender, receiver) = create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); receiver @@ -888,7 +881,8 @@ mod tests { fn create_test_service_with_ready_queue( db: MockDbConnector, ) -> (TestServiceStateWithReadyQueue, ReadyQueueSenderHandle) { - use crate::ready_queue::{ReadyQueueConfig, create_ready_queue}; + use crate::ready_queue::ReadyQueueConfig; + use crate::ready_queue::create_ready_queue; let (sender, receiver) = create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); let service = TestServiceStateWithReadyQueue::new(ServiceStateParams { diff --git a/components/spider-storage/src/state/test_utils.rs b/components/spider-storage/src/state/test_utils.rs index 7cd4d9072..6229326aa 100644 --- a/components/spider-storage/src/state/test_utils.rs +++ b/components/spider-storage/src/state/test_utils.rs @@ -1,41 +1,35 @@ -use std::{ - net::IpAddr, - sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }, -}; +use std::net::IpAddr; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use dashmap::DashMap; -use spider_core::{ - job::JobState, - types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, SchedulerId, SessionId, TaskInstanceId}, - io::TaskOutput, - scheduler::RegisteredScheduler, - }, -}; - -use crate::{ - cache::{ - error::InternalError, - task::{SharedTaskControlBlock, SharedTerminationTaskControlBlock}, - }, - db::{ - DbError, - DbStorage, - ExecutionManagerLivenessManagement, - ExternalJobOrchestration, - InternalJobOrchestration, - RecoverableJobContext, - ResourceGroupManagement, - SchedulerRegistrationManagement, - SessionManagement, - }, - job_submission::ValidatedJobSubmission, - ready_queue::ReadyQueueSender, - task_instance_pool::{TaskInstanceMetadata, TaskInstancePoolConnector}, -}; +use spider_core::job::JobState; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::TaskOutput; +use spider_core::types::scheduler::RegisteredScheduler; + +use crate::cache::error::InternalError; +use crate::cache::task::SharedTaskControlBlock; +use crate::cache::task::SharedTerminationTaskControlBlock; +use crate::db::DbError; +use crate::db::DbStorage; +use crate::db::ExecutionManagerLivenessManagement; +use crate::db::ExternalJobOrchestration; +use crate::db::InternalJobOrchestration; +use crate::db::RecoverableJobContext; +use crate::db::ResourceGroupManagement; +use crate::db::SchedulerRegistrationManagement; +use crate::db::SessionManagement; +use crate::job_submission::ValidatedJobSubmission; +use crate::ready_queue::ReadyQueueSender; +use crate::task_instance_pool::TaskInstanceMetadata; +use crate::task_instance_pool::TaskInstancePoolConnector; /// A mock ready queue sender for testing. #[derive(Clone, Default)] diff --git a/components/spider-storage/src/task_instance_pool.rs b/components/spider-storage/src/task_instance_pool.rs index c29298a81..91ef50308 100644 --- a/components/spider-storage/src/task_instance_pool.rs +++ b/components/spider-storage/src/task_instance_pool.rs @@ -13,29 +13,29 @@ //! Internally, the pool runs as a single-owner coroutine: a tokio task owns the mutable state //! directly (no mutex), processing registration messages and GC timers via `tokio::select!`. -use std::{ - collections::HashSet, - sync::{ - Arc, - atomic::{AtomicU64, Ordering}, - }, - time::{Duration, SystemTime}, -}; +use std::collections::HashSet; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::SystemTime; use async_trait::async_trait; use serde::Deserialize; -use spider_core::types::id::{ExecutionManagerId, JobId, ResourceGroupId, TaskId, TaskInstanceId}; -use tokio::{sync::mpsc, task::JoinHandle}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; -use crate::{ - cache::{ - error::InternalError, - task::{SharedTaskControlBlock, SharedTerminationTaskControlBlock}, - }, - db::ExecutionManagerLivenessManagement, - ready_queue::ReadyQueueSender, -}; +use crate::cache::error::InternalError; +use crate::cache::task::SharedTaskControlBlock; +use crate::cache::task::SharedTerminationTaskControlBlock; +use crate::db::ExecutionManagerLivenessManagement; +use crate::ready_queue::ReadyQueueSender; /// Configuration for a task instance pool actor. /// @@ -539,30 +539,26 @@ pub fn create_task_instance_pool< #[cfg(test)] mod tests { - use std::{ - net::IpAddr, - time::{Duration, SystemTime}, - }; + use std::net::IpAddr; + use std::time::Duration; + use std::time::SystemTime; use async_trait::async_trait; - use spider_core::{ - task::{ - DataTypeDescriptor, - ExecutionPolicy, - TaskDescriptor, - TaskGraph as SubmittedTaskGraph, - TdlContext, - ValueTypeDescriptor, - }, - types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId}, - io::TaskInput, - }, - }; + use spider_core::task::DataTypeDescriptor; + use spider_core::task::ExecutionPolicy; + use spider_core::task::TaskDescriptor; + use spider_core::task::TaskGraph as SubmittedTaskGraph; + use spider_core::task::TdlContext; + use spider_core::task::ValueTypeDescriptor; + use spider_core::types::id::ExecutionManagerId; + use spider_core::types::id::JobId; + use spider_core::types::id::ResourceGroupId; + use spider_core::types::io::TaskInput; use tokio::sync::Mutex; use super::*; - use crate::{db::DbError, job_submission::create_validated_submission}; + use crate::db::DbError; + use crate::job_submission::create_validated_submission; const DEFAULT_CHANNEL_SIZE: usize = 128; diff --git a/components/spider-storage/tests/jcb_test.rs b/components/spider-storage/tests/jcb_test.rs index 2dd58d91e..5a23c89bd 100644 --- a/components/spider-storage/tests/jcb_test.rs +++ b/components/spider-storage/tests/jcb_test.rs @@ -1,25 +1,20 @@ use spider_core::job::JobState; -use spider_storage::db::{ExternalJobOrchestration, InternalJobOrchestration}; - -use super::{ - scheduling_infra::{ - CancelPolicy, - DbConnectorFactory, - InstrumentSender, - WorkloadResult, - default_output_handler, - mariadb_db_connector_factory, - noop_db_connector_factory, - run_workload, - try_create_instrument_channel, - write_instrument_results, - }, - task_graph_builder::{ - build_flat_task_graph, - build_neural_net_task_graph, - create_validated_submission, - }, -}; +use spider_storage::db::ExternalJobOrchestration; +use spider_storage::db::InternalJobOrchestration; + +use super::scheduling_infra::CancelPolicy; +use super::scheduling_infra::DbConnectorFactory; +use super::scheduling_infra::InstrumentSender; +use super::scheduling_infra::WorkloadResult; +use super::scheduling_infra::default_output_handler; +use super::scheduling_infra::mariadb_db_connector_factory; +use super::scheduling_infra::noop_db_connector_factory; +use super::scheduling_infra::run_workload; +use super::scheduling_infra::try_create_instrument_channel; +use super::scheduling_infra::write_instrument_results; +use super::task_graph_builder::build_flat_task_graph; +use super::task_graph_builder::build_neural_net_task_graph; +use super::task_graph_builder::create_validated_submission; /// Evaluates to the fully-qualified name of the enclosing function, stripping internal suffixes /// like `::_f` and `::{{closure}}` that result from the macro expansion and `#[tokio::test]`. @@ -322,7 +317,8 @@ async fn test_concurrent_success_and_cancel_without_db() { test_concurrent_success_and_cancel(noop_db_connector_factory()).await; } -use super::mariadb_infra::{create_mariadb_connector, create_test_resource_group}; +use super::mariadb_infra::create_mariadb_connector; +use super::mariadb_infra::create_test_resource_group; #[tokio::test(flavor = "multi_thread")] #[ignore = "requires MariaDB"] diff --git a/components/spider-storage/tests/mariadb_infra.rs b/components/spider-storage/tests/mariadb_infra.rs index 806ef4bd7..8358f199a 100644 --- a/components/spider-storage/tests/mariadb_infra.rs +++ b/components/spider-storage/tests/mariadb_infra.rs @@ -1,9 +1,8 @@ use secrecy::SecretString; use spider_core::types::id::ResourceGroupId; -use spider_storage::{ - DatabaseConfig, - db::{MariaDbStorageConnector, ResourceGroupManagement}, -}; +use spider_storage::DatabaseConfig; +use spider_storage::db::MariaDbStorageConnector; +use spider_storage::db::ResourceGroupManagement; /// Creates a [`MariaDbStorageConnector`] from environment variables. /// diff --git a/components/spider-storage/tests/mariadb_test.rs b/components/spider-storage/tests/mariadb_test.rs index ff7eff582..c9a375afc 100644 --- a/components/spider-storage/tests/mariadb_test.rs +++ b/components/spider-storage/tests/mariadb_test.rs @@ -1,31 +1,29 @@ -use std::{ - net::{IpAddr, Ipv4Addr}, - time::Duration, -}; - -use spider_core::{ - job::JobState, - types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, SchedulerId}, - io::TaskInput, - }, -}; -use spider_storage::db::{ - DbError, - ExecutionManagerLivenessManagement, - ExternalJobOrchestration, - InternalJobOrchestration, - MariaDbStorageConnector, - ResourceGroupManagement, - SchedulerRegistrationManagement, - SessionManagement, -}; +use std::net::IpAddr; +use std::net::Ipv4Addr; +use std::time::Duration; + +use spider_core::job::JobState; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SchedulerId; +use spider_core::types::io::TaskInput; +use spider_storage::db::DbError; +use spider_storage::db::ExecutionManagerLivenessManagement; +use spider_storage::db::ExternalJobOrchestration; +use spider_storage::db::InternalJobOrchestration; +use spider_storage::db::MariaDbStorageConnector; +use spider_storage::db::ResourceGroupManagement; +use spider_storage::db::SchedulerRegistrationManagement; +use spider_storage::db::SessionManagement; use tokio::task::JoinSet; -use super::{ - mariadb_infra::{create_mariadb_config, create_mariadb_connector, create_test_resource_group}, - task_graph_builder::{SubmittedTaskGraph, build_flat_task_graph, create_validated_submission}, -}; +use super::mariadb_infra::create_mariadb_config; +use super::mariadb_infra::create_mariadb_connector; +use super::mariadb_infra::create_test_resource_group; +use super::task_graph_builder::SubmittedTaskGraph; +use super::task_graph_builder::build_flat_task_graph; +use super::task_graph_builder::create_validated_submission; /// Input payload size in bytes for the single-task graph used by DB-layer tests. const TEST_INPUT_PAYLOAD_SIZE: usize = 128; diff --git a/components/spider-storage/tests/runtime_recovery_test.rs b/components/spider-storage/tests/runtime_recovery_test.rs index 4971b537e..012f0cee5 100644 --- a/components/spider-storage/tests/runtime_recovery_test.rs +++ b/components/spider-storage/tests/runtime_recovery_test.rs @@ -1,32 +1,33 @@ -use std::{net::IpAddr, time::Duration}; - -use spider_core::{ - job::JobState, - task::TaskIndex, - types::{ - id::{ExecutionManagerId, JobId, TaskInstanceId}, - io::{TaskOutput, TaskOutputsSerializer}, - }, -}; -use spider_storage::{ - cache::error::{CacheError, StaleStateError}, - db::ExternalJobOrchestration, - ready_queue::{CleanupTaskMarker, CommitTaskMarker, ReadyQueueConfig, ReadyQueueEntry}, - state::{ - JobCacheGcConfig, - Runtime, - ServiceState, - StorageServerError, - create_runtime, - runtime::RuntimeConfig, - }, - task_instance_pool::TaskInstancePoolConfig, -}; - -use crate::{ - mariadb_infra::{create_mariadb_config, create_mariadb_connector}, - task_graph_builder::{build_flat_task_graph, compress_job_inputs, compress_task_graph}, -}; +use std::net::IpAddr; +use std::time::Duration; + +use spider_core::job::JobState; +use spider_core::task::TaskIndex; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::TaskOutput; +use spider_core::types::io::TaskOutputsSerializer; +use spider_storage::cache::error::CacheError; +use spider_storage::cache::error::StaleStateError; +use spider_storage::db::ExternalJobOrchestration; +use spider_storage::ready_queue::CleanupTaskMarker; +use spider_storage::ready_queue::CommitTaskMarker; +use spider_storage::ready_queue::ReadyQueueConfig; +use spider_storage::ready_queue::ReadyQueueEntry; +use spider_storage::state::JobCacheGcConfig; +use spider_storage::state::Runtime; +use spider_storage::state::ServiceState; +use spider_storage::state::StorageServerError; +use spider_storage::state::create_runtime; +use spider_storage::state::runtime::RuntimeConfig; +use spider_storage::task_instance_pool::TaskInstancePoolConfig; + +use crate::mariadb_infra::create_mariadb_config; +use crate::mariadb_infra::create_mariadb_connector; +use crate::task_graph_builder::build_flat_task_graph; +use crate::task_graph_builder::compress_job_inputs; +use crate::task_graph_builder::compress_task_graph; #[tokio::test] #[ignore = "requires MariaDB"] diff --git a/components/spider-storage/tests/scheduling_infra.rs b/components/spider-storage/tests/scheduling_infra.rs index f1580cca1..3e762d5da 100644 --- a/components/spider-storage/tests/scheduling_infra.rs +++ b/components/spider-storage/tests/scheduling_infra.rs @@ -71,45 +71,45 @@ //! `AsyncFnOnce() -> DbConnectorType`, so the same tests can be rerun with a real //! DB connector later. -use std::{ - sync::{ - Arc, - atomic::{AtomicU64, AtomicUsize, Ordering}, - }, - time::{Duration, Instant}, -}; +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::Instant; use anyhow::bail; use async_trait::async_trait; use dashmap::DashMap; -use rand::{Rng, SeedableRng}; -use spider_core::{ - job::JobState, - task::TaskIndex, - types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, TaskId, TaskInstanceId}, - io::{ExecutionContext, TaskOutput}, - }, -}; -use spider_storage::{ - cache::{ - error::{CacheError, InternalError}, - job::SharedJobControlBlock, - task::{SharedTaskControlBlock, SharedTerminationTaskControlBlock}, - }, - db::{ - DbError, - ExternalJobOrchestration, - InternalJobOrchestration, - MariaDbStorageConnector, - RecoverableJobContext, - }, - job_submission::ValidatedJobSubmission, - ready_queue::ReadyQueueSender, - task_instance_pool::{TaskInstanceMetadata, TaskInstancePoolConnector}, -}; -use tabled::{Table, Tabled}; -use tokio::sync::{mpsc, watch}; +use rand::Rng; +use rand::SeedableRng; +use spider_core::job::JobState; +use spider_core::task::TaskIndex; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::ExecutionContext; +use spider_core::types::io::TaskOutput; +use spider_storage::cache::error::CacheError; +use spider_storage::cache::error::InternalError; +use spider_storage::cache::job::SharedJobControlBlock; +use spider_storage::cache::task::SharedTaskControlBlock; +use spider_storage::cache::task::SharedTerminationTaskControlBlock; +use spider_storage::db::DbError; +use spider_storage::db::ExternalJobOrchestration; +use spider_storage::db::InternalJobOrchestration; +use spider_storage::db::MariaDbStorageConnector; +use spider_storage::db::RecoverableJobContext; +use spider_storage::job_submission::ValidatedJobSubmission; +use spider_storage::ready_queue::ReadyQueueSender; +use spider_storage::task_instance_pool::TaskInstanceMetadata; +use spider_storage::task_instance_pool::TaskInstancePoolConnector; +use tabled::Table; +use tabled::Tabled; +use tokio::sync::mpsc; +use tokio::sync::watch; /// A handler that generates mock task outputs from an [`ExecutionContext`], which simulates the /// execution of a TDL task. diff --git a/components/spider-storage/tests/task_graph_builder.rs b/components/spider-storage/tests/task_graph_builder.rs index 6935de4d3..cd450bb0f 100644 --- a/components/spider-storage/tests/task_graph_builder.rs +++ b/components/spider-storage/tests/task_graph_builder.rs @@ -1,19 +1,17 @@ use anyhow::Result; -use rand::{Rng, SeedableRng}; -use spider_core::{ - compression::encode_zstd_bytes, - task::{ - DataTypeDescriptor, - ExecutionPolicy, - TaskDescriptor, - TaskIndex, - TaskInputOutputIndex, - TdlContext, - TerminationTaskDescriptor, - ValueTypeDescriptor, - }, - types::io::{TaskInput, TaskInputsSerializer}, -}; +use rand::Rng; +use rand::SeedableRng; +use spider_core::compression::encode_zstd_bytes; +use spider_core::task::DataTypeDescriptor; +use spider_core::task::ExecutionPolicy; +use spider_core::task::TaskDescriptor; +use spider_core::task::TaskIndex; +use spider_core::task::TaskInputOutputIndex; +use spider_core::task::TdlContext; +use spider_core::task::TerminationTaskDescriptor; +use spider_core::task::ValueTypeDescriptor; +use spider_core::types::io::TaskInput; +use spider_core::types::io::TaskInputsSerializer; use spider_storage::job_submission::ValidatedJobSubmission; /// The submitted task graph type from spider-core. diff --git a/components/spider-task-executor/src/bin/spider_task_executor.rs b/components/spider-task-executor/src/bin/spider_task_executor.rs index 5be95bf05..c694e4e33 100644 --- a/components/spider-task-executor/src/bin/spider_task_executor.rs +++ b/components/spider-task-executor/src/bin/spider_task_executor.rs @@ -13,21 +13,25 @@ //! the executor itself has no concurrency requirements, and exactly one task runs for the //! lifetime of the process. -use std::{ - path::{Path, PathBuf}, - time::Instant, -}; +use std::path::Path; +use std::path::PathBuf; +use std::time::Instant; -use anyhow::{Result, anyhow}; +use anyhow::Result; +use anyhow::anyhow; use bytes::Bytes; -use futures_util::{SinkExt, StreamExt}; -use spider_task_executor::{ - ExecutorError, - TdlPackageManager, - protocol::{ExecutorOutcome, Request, Response}, -}; -use tokio::io::{stdin, stdout}; -use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec}; +use futures_util::SinkExt; +use futures_util::StreamExt; +use spider_task_executor::ExecutorError; +use spider_task_executor::TdlPackageManager; +use spider_task_executor::protocol::ExecutorOutcome; +use spider_task_executor::protocol::Request; +use spider_task_executor::protocol::Response; +use tokio::io::stdin; +use tokio::io::stdout; +use tokio_util::codec::FramedRead; +use tokio_util::codec::FramedWrite; +use tokio_util::codec::LengthDelimitedCodec; /// Env var that points to the directory where compiled TDL packages live. const SPIDER_TDL_PACKAGE_DIR: &str = "SPIDER_TDL_PACKAGE_DIR"; diff --git a/components/spider-task-executor/src/error.rs b/components/spider-task-executor/src/error.rs index c8da04ef2..f25b4769e 100644 --- a/components/spider-task-executor/src/error.rs +++ b/components/spider-task-executor/src/error.rs @@ -1,6 +1,7 @@ //! Errors produced while loading TDL packages or executing tasks across the FFI boundary. -use spider_tdl::{TdlError, Version}; +use spider_tdl::TdlError; +use spider_tdl::Version; /// All possible errors produced by the task executor. /// diff --git a/components/spider-task-executor/src/lib.rs b/components/spider-task-executor/src/lib.rs index 3afb04846..1502eba87 100644 --- a/components/spider-task-executor/src/lib.rs +++ b/components/spider-task-executor/src/lib.rs @@ -5,4 +5,5 @@ pub mod manager; pub mod protocol; pub use error::ExecutorError; -pub use manager::{TdlPackage, TdlPackageManager}; +pub use manager::TdlPackage; +pub use manager::TdlPackageManager; diff --git a/components/spider-task-executor/src/manager.rs b/components/spider-task-executor/src/manager.rs index d7cfe8c0d..80194085b 100644 --- a/components/spider-task-executor/src/manager.rs +++ b/components/spider-task-executor/src/manager.rs @@ -3,14 +3,16 @@ //! See [`TdlPackage`] for the per-library wrapper and [`TdlPackageManager`] for the top-level //! collection that enforces unique package names. -use std::{collections::HashMap, path::Path}; +use std::collections::HashMap; +use std::path::Path; -use libloading::{Library, Symbol}; -use spider_tdl::{ - TdlError, - Version, - ffi::{CByteArray, CCharArray, TaskExecutionResult}, -}; +use libloading::Library; +use libloading::Symbol; +use spider_tdl::TdlError; +use spider_tdl::Version; +use spider_tdl::ffi::CByteArray; +use spider_tdl::ffi::CCharArray; +use spider_tdl::ffi::TaskExecutionResult; use crate::error::ExecutorError; diff --git a/components/spider-task-executor/src/protocol.rs b/components/spider-task-executor/src/protocol.rs index 935d60d7b..b2ac25119 100644 --- a/components/spider-task-executor/src/protocol.rs +++ b/components/spider-task-executor/src/protocol.rs @@ -8,7 +8,8 @@ //! stderr; how those bytes are disposed of (inherited, piped, redirected to a log file) is a choice //! made by whoever spawned the process. -use serde::{Deserialize, Serialize}; +use serde::Deserialize; +use serde::Serialize; use spider_core::task::TdlContext; /// Request from the parent process (execution manager) to the executor. diff --git a/components/spider-tdl-derive/src/lib.rs b/components/spider-tdl-derive/src/lib.rs index de5dfbc47..7f3852e9f 100644 --- a/components/spider-tdl-derive/src/lib.rs +++ b/components/spider-tdl-derive/src/lib.rs @@ -7,7 +7,8 @@ mod task_macro; use proc_macro::TokenStream; -use syn::{ItemFn, parse_macro_input}; +use syn::ItemFn; +use syn::parse_macro_input; use task_macro::TaskAttr; /// Attribute macro that transforms a function into a Spider TDL–compatible task. diff --git a/components/spider-tdl-derive/src/task_macro.rs b/components/spider-tdl-derive/src/task_macro.rs index f932aeb8f..6329d94f2 100644 --- a/components/spider-tdl-derive/src/task_macro.rs +++ b/components/spider-tdl-derive/src/task_macro.rs @@ -14,19 +14,19 @@ //! See the crate-level docs of [`spider_tdl_derive`] for usage examples. use proc_macro2::TokenStream; -use quote::{format_ident, quote}; -use syn::{ - FnArg, - GenericArgument, - ItemFn, - LitStr, - Pat, - PathArguments, - ReturnType, - Token, - Type, - parse::{Parse, ParseStream}, -}; +use quote::format_ident; +use quote::quote; +use syn::FnArg; +use syn::GenericArgument; +use syn::ItemFn; +use syn::LitStr; +use syn::Pat; +use syn::PathArguments; +use syn::ReturnType; +use syn::Token; +use syn::Type; +use syn::parse::Parse; +use syn::parse::ParseStream; /// Parsed representation of the `#[task(...)]` attribute arguments. pub struct TaskAttr { diff --git a/components/spider-tdl/src/ffi.rs b/components/spider-tdl/src/ffi.rs index aba8bd6d8..544cd128d 100644 --- a/components/spider-tdl/src/ffi.rs +++ b/components/spider-tdl/src/ffi.rs @@ -4,7 +4,12 @@ //! so buffers allocated on one side can be reclaimed on the other via [`Box::into_raw`] / //! [`Box::from_raw`]. -use std::{ffi::c_char, fmt, marker::PhantomData, mem::ManuallyDrop, ops::Deref, str::Utf8Error}; +use std::ffi::c_char; +use std::fmt; +use std::marker::PhantomData; +use std::mem::ManuallyDrop; +use std::ops::Deref; +use std::str::Utf8Error; /// Borrowed, C-ABI-compatible view of a contiguous slice `&'borrow_lifetime [ElementType]`. /// diff --git a/components/spider-tdl/src/lib.rs b/components/spider-tdl/src/lib.rs index 61ac579c3..7f1cdd8cc 100644 --- a/components/spider-tdl/src/lib.rs +++ b/components/spider-tdl/src/lib.rs @@ -9,6 +9,9 @@ pub mod version; pub use error::TdlError; #[cfg(feature = "derive")] pub use spider_tdl_derive::task; -pub use task::{ExecutionResult, Task, TaskHandler, TaskHandlerImpl}; +pub use task::ExecutionResult; +pub use task::Task; +pub use task::TaskHandler; +pub use task::TaskHandlerImpl; pub use task_context::TaskContext; pub use version::Version; diff --git a/components/spider-tdl/src/task.rs b/components/spider-tdl/src/task.rs index a44206dca..d235dfeeb 100644 --- a/components/spider-tdl/src/task.rs +++ b/components/spider-tdl/src/task.rs @@ -6,9 +6,11 @@ use std::marker::PhantomData; -use spider_core::types::io::{TaskInputsSerializer, TaskOutputsSerializer}; +use spider_core::types::io::TaskInputsSerializer; +use spider_core::types::io::TaskOutputsSerializer; -use crate::{error::TdlError, task_context::TaskContext}; +use crate::error::TdlError; +use crate::task_context::TaskContext; /// The result of executing a task through the [`TaskHandler`] interface. /// @@ -168,14 +170,19 @@ fn serialize_error(err: &TdlError) -> Vec { #[cfg(test)] mod tests { - use serde::{Deserialize, Serialize}; - use spider_core::types::{ - id::{JobId, ResourceGroupId, TaskId}, - io::{TaskInput, TaskInputsSerializer, TaskOutputsSerializer}, - }; + use serde::Deserialize; + use serde::Serialize; + use spider_core::types::id::JobId; + use spider_core::types::id::ResourceGroupId; + use spider_core::types::id::TaskId; + use spider_core::types::io::TaskInput; + use spider_core::types::io::TaskInputsSerializer; + use spider_core::types::io::TaskOutputsSerializer; use super::*; - use crate::{error::TdlError, r#std::int32, task_context::TaskContext}; + use crate::error::TdlError; + use crate::r#std::int32; + use crate::task_context::TaskContext; const INTENTIONAL_FAILURE: &str = "intentional failure"; diff --git a/components/spider-tdl/src/task_context.rs b/components/spider-tdl/src/task_context.rs index d79dea6d1..6e0d648d7 100644 --- a/components/spider-tdl/src/task_context.rs +++ b/components/spider-tdl/src/task_context.rs @@ -4,7 +4,10 @@ //! through the task executor, and deserialized inside the TDL package before being handed to the //! user's task function as the first parameter. -use spider_core::types::id::{JobId, ResourceGroupId, TaskId, TaskInstanceId}; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; /// Runtime metadata about the current task execution. /// @@ -23,7 +26,9 @@ pub struct TaskContext { #[cfg(test)] mod tests { - use spider_core::types::id::{JobId, ResourceGroupId, TaskId}; + use spider_core::types::id::JobId; + use spider_core::types::id::ResourceGroupId; + use spider_core::types::id::TaskId; use super::TaskContext; diff --git a/components/spider-tdl/tests/test_task_macro.rs b/components/spider-tdl/tests/test_task_macro.rs index 4ad854bd2..3f772267a 100644 --- a/components/spider-tdl/tests/test_task_macro.rs +++ b/components/spider-tdl/tests/test_task_macro.rs @@ -4,12 +4,21 @@ //! downstream TDL package would consume it: through the `spider_tdl::task` re-export, with the //! generated code resolving `::spider_tdl` and `::serde` as external crates. -use serde::{Deserialize, Serialize}; -use spider_core::types::{ - id::{JobId, ResourceGroupId, TaskId}, - io::{TaskInput, TaskInputsSerializer, TaskOutputsSerializer}, -}; -use spider_tdl::{Task, TaskContext, TaskHandler, TaskHandlerImpl, TdlError, r#std::int32, task}; +use serde::Deserialize; +use serde::Serialize; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::io::TaskInput; +use spider_core::types::io::TaskInputsSerializer; +use spider_core::types::io::TaskOutputsSerializer; +use spider_tdl::Task; +use spider_tdl::TaskContext; +use spider_tdl::TaskHandler; +use spider_tdl::TaskHandlerImpl; +use spider_tdl::TdlError; +use spider_tdl::r#std::int32; +use spider_tdl::task; type AliasedContext = TaskContext; type _AliasedTdlError = TdlError; diff --git a/components/spider-utils/src/config.rs b/components/spider-utils/src/config.rs index 41eb44d39..3a7342a3a 100644 --- a/components/spider-utils/src/config.rs +++ b/components/spider-utils/src/config.rs @@ -1,11 +1,11 @@ -use std::{ - fs::File, - io, - net::{IpAddr, SocketAddr}, - path::Path, -}; +use std::fs::File; +use std::io; +use std::net::IpAddr; +use std::net::SocketAddr; +use std::path::Path; -use serde::{Deserialize, de::DeserializeOwned}; +use serde::Deserialize; +use serde::de::DeserializeOwned; use thiserror::Error; use tonic::transport::Endpoint; diff --git a/components/spider-utils/src/grpc/client.rs b/components/spider-utils/src/grpc/client.rs index 71c126625..fa06ff256 100644 --- a/components/spider-utils/src/grpc/client.rs +++ b/components/spider-utils/src/grpc/client.rs @@ -1,11 +1,11 @@ //! A round-robin pool of gRPC service-client connections. -use std::{ - num::NonZeroUsize, - sync::{Arc, atomic::AtomicUsize}, -}; +use std::num::NonZeroUsize; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; -use tonic::transport::{Channel, Endpoint}; +use tonic::transport::Channel; +use tonic::transport::Endpoint; use super::Error; diff --git a/components/spider-utils/src/logging.rs b/components/spider-utils/src/logging.rs index 4d8617c16..0d16b1739 100644 --- a/components/spider-utils/src/logging.rs +++ b/components/spider-utils/src/logging.rs @@ -1,6 +1,7 @@ //! Logging setup for Spider binaries. -use tracing_appender::non_blocking::{NonBlockingBuilder, WorkerGuard}; +use tracing_appender::non_blocking::NonBlockingBuilder; +use tracing_appender::non_blocking::WorkerGuard; use tracing_subscriber::EnvFilter; /// Initializes the global tracing subscriber with JSON-formatted, non-blocking output to `stderr`. diff --git a/components/spider-utils/src/wire.rs b/components/spider-utils/src/wire.rs index fe0767cf7..25c5d0d97 100644 --- a/components/spider-utils/src/wire.rs +++ b/components/spider-utils/src/wire.rs @@ -16,10 +16,11 @@ use std::fmt; -use serde::{ - de::{self, DeserializeSeed, SeqAccess, Visitor}, - ser, -}; +use serde::de::DeserializeSeed; +use serde::de::SeqAccess; +use serde::de::Visitor; +use serde::de::{self}; +use serde::ser; /// Errors produced while framing or unframing a wire buffer. #[derive(Debug, thiserror::Error)] diff --git a/examples/huntsman/complex/tasks/src/lib.rs b/examples/huntsman/complex/tasks/src/lib.rs index 4f6bc9452..2c080b074 100644 --- a/examples/huntsman/complex/tasks/src/lib.rs +++ b/examples/huntsman/complex/tasks/src/lib.rs @@ -1,8 +1,11 @@ //! Reference TDL package: complex-vector arithmetic. mod task_decl { - use huntsman_complex_types::{Complex, ComplexVec}; - use spider_tdl::{TaskContext, TdlError, task}; + use huntsman_complex_types::Complex; + use huntsman_complex_types::ComplexVec; + use spider_tdl::TaskContext; + use spider_tdl::TdlError; + use spider_tdl::task; #[task(name = "complex::add")] pub fn add(_ctx: TaskContext, a: ComplexVec, b: ComplexVec) -> Result { diff --git a/examples/huntsman/complex/types/src/lib.rs b/examples/huntsman/complex/types/src/lib.rs index b9853700f..3b1ffdc62 100644 --- a/examples/huntsman/complex/types/src/lib.rs +++ b/examples/huntsman/complex/types/src/lib.rs @@ -4,8 +4,10 @@ //! Splitting the types out of the cdylib lets the integration test crate depend on this rlib and //! reuse the canonical struct definitions instead of declaring a parallel mirror. -use serde::{Deserialize, Serialize}; -use spider_tdl::r#std::{List, double}; +use serde::Deserialize; +use serde::Serialize; +use spider_tdl::r#std::List; +use spider_tdl::r#std::double; /// A complex number with [`double`] components. #[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] diff --git a/tests/huntsman/em-runtime/tests/test_runtime.rs b/tests/huntsman/em-runtime/tests/test_runtime.rs index d06216930..c4d5d8251 100644 --- a/tests/huntsman/em-runtime/tests/test_runtime.rs +++ b/tests/huntsman/em-runtime/tests/test_runtime.rs @@ -7,39 +7,39 @@ //! //! All tests are `#[ignore]` so the workspace's plain `cargo test` doesn't run them. -use std::{path::PathBuf, sync::Arc, time::Duration}; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; use anyhow::Context; -use spider_core::{ - task::{TdlContext, TimeoutPolicy}, - types::{ - id::{ - ExecutionManagerId, - JobId, - ResourceGroupId, - SchedulerId, - SessionId, - TaskAssignmentId, - TaskId, - }, - io::{ExecutionContext, TaskInput, TaskInputsSerializer}, - scheduler::{TaskAssignment, TaskAssignmentRecord}, - }, -}; -use spider_execution_manager::{ - client::{SchedulerError, SchedulerResponse, StorageResponseError}, - runtime::{Runtime, RuntimeConfig, RuntimeError}, -}; -use test_utils::{ - MockLiveness, - MockScheduler, - MockStorage, - PACKAGE_NAME, - decode_single_output, - single_input, - task_executor_bin, - tdl_package_dir, -}; +use spider_core::task::TdlContext; +use spider_core::task::TimeoutPolicy; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::SchedulerId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskAssignmentId; +use spider_core::types::id::TaskId; +use spider_core::types::io::ExecutionContext; +use spider_core::types::io::TaskInput; +use spider_core::types::io::TaskInputsSerializer; +use spider_core::types::scheduler::TaskAssignment; +use spider_core::types::scheduler::TaskAssignmentRecord; +use spider_execution_manager::client::SchedulerError; +use spider_execution_manager::client::SchedulerResponse; +use spider_execution_manager::client::StorageResponseError; +use spider_execution_manager::runtime::Runtime; +use spider_execution_manager::runtime::RuntimeConfig; +use spider_execution_manager::runtime::RuntimeError; +use test_utils::MockLiveness; +use test_utils::MockScheduler; +use test_utils::MockStorage; +use test_utils::PACKAGE_NAME; +use test_utils::decode_single_output; +use test_utils::single_input; +use test_utils::task_executor_bin; +use test_utils::tdl_package_dir; const HEARTBEAT_INTERVAL: Duration = Duration::from_millis(100); const SLOW_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5); diff --git a/tests/huntsman/integration-test-tasks/src/lib.rs b/tests/huntsman/integration-test-tasks/src/lib.rs index 1f6bc7312..c432bcb83 100644 --- a/tests/huntsman/integration-test-tasks/src/lib.rs +++ b/tests/huntsman/integration-test-tasks/src/lib.rs @@ -17,9 +17,12 @@ pub const INSTRUMENT_SLEEP_US: u64 = 50; mod task_decl { - use std::{thread::sleep, time::Duration}; + use std::thread::sleep; + use std::time::Duration; - use spider_tdl::{TaskContext, TdlError, task}; + use spider_tdl::TaskContext; + use spider_tdl::TdlError; + use spider_tdl::task; use crate::INSTRUMENT_SLEEP_US; diff --git a/tests/huntsman/task-executor/tests/overhead_instrument.rs b/tests/huntsman/task-executor/tests/overhead_instrument.rs index 64bba93e7..ce40b3806 100644 --- a/tests/huntsman/task-executor/tests/overhead_instrument.rs +++ b/tests/huntsman/task-executor/tests/overhead_instrument.rs @@ -18,17 +18,21 @@ //! Aggregates (avg, p50, p95, p99) for each metric land in a markdown table at //! `${SPIDER_TEST_INSTRUMENT_OUTPUT_DIR}/task_executor_overhead.md`. -use std::{ - fs::File, - io::Write, - path::PathBuf, - time::{Duration, Instant}, -}; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use std::time::Duration; +use std::time::Instant; use integration_test_tasks::INSTRUMENT_SLEEP_US; -use spider_task_executor::protocol::{ExecutorOutcome, Response}; -use tabled::{Table, Tabled}; -use test_utils::{ExecutorHandle, decode_single_output, encode_single_input, execute_request}; +use spider_task_executor::protocol::ExecutorOutcome; +use spider_task_executor::protocol::Response; +use tabled::Table; +use tabled::Tabled; +use test_utils::ExecutorHandle; +use test_utils::decode_single_output; +use test_utils::encode_single_input; +use test_utils::execute_request; const PAYLOAD_LEN: usize = 100; const ITERATIONS: usize = 10; diff --git a/tests/huntsman/task-executor/tests/test_executor.rs b/tests/huntsman/task-executor/tests/test_executor.rs index cd91c1d6a..acc0b3903 100644 --- a/tests/huntsman/task-executor/tests/test_executor.rs +++ b/tests/huntsman/task-executor/tests/test_executor.rs @@ -3,18 +3,15 @@ //! Each test spawns a fresh executor subprocess via [`ExecutorHandle::spawn`], exchanges one framed //! bincode request/response over the binary's stdin/stdout, and asserts on the result. -use spider_task_executor::{ - ExecutorError, - protocol::{ExecutorOutcome, Response}, -}; +use spider_task_executor::ExecutorError; +use spider_task_executor::protocol::ExecutorOutcome; +use spider_task_executor::protocol::Response; use spider_tdl::TdlError; -use test_utils::{ - ExecutorHandle, - decode_single_output, - encode_no_inputs, - encode_single_input, - execute_request, -}; +use test_utils::ExecutorHandle; +use test_utils::decode_single_output; +use test_utils::encode_no_inputs; +use test_utils::encode_single_input; +use test_utils::execute_request; #[tokio::test] #[ignore = "requires `integration-test-tasks` cdylib and `spider-task-executor` binary"] diff --git a/tests/huntsman/task-executor/tests/test_process_pool.rs b/tests/huntsman/task-executor/tests/test_process_pool.rs index 1542c6a51..26a6c3ad7 100644 --- a/tests/huntsman/task-executor/tests/test_process_pool.rs +++ b/tests/huntsman/task-executor/tests/test_process_pool.rs @@ -13,28 +13,26 @@ use std::time::Duration; -use spider_core::{ - task::{TdlContext, TimeoutPolicy}, - types::{ - id::{ExecutionManagerId, JobId, ResourceGroupId, TaskId}, - io::{ExecutionContext, TaskInput, TaskInputsSerializer}, - }, -}; -use spider_execution_manager::process_pool::{ - ExecuteRequest, - Outcome, - ProcessPool, - ProcessPoolConfig, -}; +use spider_core::task::TdlContext; +use spider_core::task::TimeoutPolicy; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::io::ExecutionContext; +use spider_core::types::io::TaskInput; +use spider_core::types::io::TaskInputsSerializer; +use spider_execution_manager::process_pool::ExecuteRequest; +use spider_execution_manager::process_pool::Outcome; +use spider_execution_manager::process_pool::ProcessPool; +use spider_execution_manager::process_pool::ProcessPoolConfig; use spider_task_executor::ExecutorError; use spider_tdl::TdlError; -use test_utils::{ - PACKAGE_NAME, - decode_single_output, - single_input, - task_executor_bin, - tdl_package_dir, -}; +use test_utils::PACKAGE_NAME; +use test_utils::decode_single_output; +use test_utils::single_input; +use test_utils::task_executor_bin; +use test_utils::tdl_package_dir; /// Generous timeout for tasks expected to finish quickly. const NORMAL_TIMEOUT: Duration = Duration::from_secs(5); diff --git a/tests/huntsman/tdl-integration/tests/complex.rs b/tests/huntsman/tdl-integration/tests/complex.rs index 024fd5d63..2c5a69aef 100644 --- a/tests/huntsman/tdl-integration/tests/complex.rs +++ b/tests/huntsman/tdl-integration/tests/complex.rs @@ -1,12 +1,18 @@ //! End-to-end tests for the TDL package executor against the `huntsman-complex` example crate. -use huntsman_complex_types::{Complex, ComplexVec}; -use spider_core::types::{ - id::{JobId, ResourceGroupId, TaskId}, - io::{TaskInput, TaskInputsSerializer, TaskOutputsSerializer}, -}; -use spider_task_executor::{ExecutorError, TdlPackageManager}; -use spider_tdl::{TaskContext, TdlError, Version}; +use huntsman_complex_types::Complex; +use huntsman_complex_types::ComplexVec; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::io::TaskInput; +use spider_core::types::io::TaskInputsSerializer; +use spider_core::types::io::TaskOutputsSerializer; +use spider_task_executor::ExecutorError; +use spider_task_executor::TdlPackageManager; +use spider_tdl::TaskContext; +use spider_tdl::TdlError; +use spider_tdl::Version; const PACKAGE_NAME: &str = "complex"; diff --git a/tests/huntsman/test-utils/src/executor.rs b/tests/huntsman/test-utils/src/executor.rs index 8dd912d2e..0b93f8351 100644 --- a/tests/huntsman/test-utils/src/executor.rs +++ b/tests/huntsman/test-utils/src/executor.rs @@ -14,21 +14,29 @@ //! * `SPIDER_TDL_PACKAGE_DIR` — directory the binary searches for TDL packages; gets forwarded to //! the child verbatim. -use std::{path::PathBuf, process::Stdio}; +use std::path::PathBuf; +use std::process::Stdio; use bytes::Bytes; -use futures_util::{SinkExt, StreamExt}; -use spider_core::{ - task::TdlContext, - types::{ - id::{JobId, ResourceGroupId, TaskId}, - io::{TaskInput, TaskInputsSerializer, TaskOutputsSerializer}, - }, -}; -use spider_task_executor::protocol::{Request, Response}; +use futures_util::SinkExt; +use futures_util::StreamExt; +use spider_core::task::TdlContext; +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::io::TaskInput; +use spider_core::types::io::TaskInputsSerializer; +use spider_core::types::io::TaskOutputsSerializer; +use spider_task_executor::protocol::Request; +use spider_task_executor::protocol::Response; use spider_tdl::TaskContext; -use tokio::process::{Child, ChildStdin, ChildStdout, Command}; -use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec}; +use tokio::process::Child; +use tokio::process::ChildStdin; +use tokio::process::ChildStdout; +use tokio::process::Command; +use tokio_util::codec::FramedRead; +use tokio_util::codec::FramedWrite; +use tokio_util::codec::LengthDelimitedCodec; /// The TDL package name registered by `integration-test-tasks`. pub const PACKAGE_NAME: &str = "integration_test_tasks"; diff --git a/tests/huntsman/test-utils/src/mock.rs b/tests/huntsman/test-utils/src/mock.rs index 7d4b5d028..2a4f96e12 100644 --- a/tests/huntsman/test-utils/src/mock.rs +++ b/tests/huntsman/test-utils/src/mock.rs @@ -4,36 +4,34 @@ //! while the runtime owns a clone. Response queues let the test drive deterministic call sequences; //! inboxes record every call so assertions can be made. -use std::{ - collections::VecDeque, - net::IpAddr, - sync::{ - Arc, - Mutex, - MutexGuard, - PoisonError, - atomic::{AtomicBool, AtomicU64, Ordering}, - }, - time::Duration, -}; +use std::collections::VecDeque; +use std::net::IpAddr; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; use async_trait::async_trait; use dashmap::DashSet; -use spider_core::types::{ - id::{ExecutionManagerId, JobId, SessionId, TaskId, TaskInstanceId}, - io::ExecutionContext, - scheduler::TaskAssignmentRecord, -}; -use spider_execution_manager::client::{ - LivenessClient, - LivenessResponseError, - RegistrationResponse, - SchedulerClient, - SchedulerError, - SchedulerResponse, - StorageClient, - StorageResponseError, -}; +use spider_core::types::id::ExecutionManagerId; +use spider_core::types::id::JobId; +use spider_core::types::id::SessionId; +use spider_core::types::id::TaskId; +use spider_core::types::id::TaskInstanceId; +use spider_core::types::io::ExecutionContext; +use spider_core::types::scheduler::TaskAssignmentRecord; +use spider_execution_manager::client::LivenessClient; +use spider_execution_manager::client::LivenessResponseError; +use spider_execution_manager::client::RegistrationResponse; +use spider_execution_manager::client::SchedulerClient; +use spider_execution_manager::client::SchedulerError; +use spider_execution_manager::client::SchedulerResponse; +use spider_execution_manager::client::StorageClient; +use spider_execution_manager::client::StorageResponseError; use tokio::sync::Notify; /// Mock [`SchedulerClient`]. diff --git a/tools/yscope-dev-utils b/tools/yscope-dev-utils index 38bf51eff..70f63f266 160000 --- a/tools/yscope-dev-utils +++ b/tools/yscope-dev-utils @@ -1 +1 @@ -Subproject commit 38bf51effc500a528c37052345c141a65ea88447 +Subproject commit 70f63f266c35657362b2c26df8403621d9400b4f