rsketch_server/
lib.rs

1// Copyright 2025 Crrow
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub mod error;
16pub mod grpc;
17pub mod http;
18
19use futures::future::join_all;
20use rsketch_error::Result;
21use tokio::{sync::oneshot::Receiver, task::JoinHandle};
22use tokio_util::sync::CancellationToken;
23
24/// Handle for managing a running service: grpc or http.
25///
26/// This handle provides control over a running service, allowing you to:
27/// - Wait for the service to start accepting connections
28/// - Signal graceful shutdown
29/// - Wait for the service to fully stop
30/// - Check if the service task has completed
31///
32/// The handle uses a cancellation token for graceful shutdown and provides
33/// async methods for coordinating server lifecycle events.
34pub struct ServiceHandler {
35    /// Join handle for the server task
36    join_handle:        JoinHandle<()>,
37    /// Token for signalling shutdown
38    cancellation_token: CancellationToken,
39    /// Receiver for server start notification
40    started_rx:         Option<Receiver<()>>,
41    /// Join handles for readiness reporting tasks
42    reporter_handles:   Vec<JoinHandle<()>>,
43}
44
45impl ServiceHandler {
46    /// Waits for the server to start accepting connections.
47    ///
48    /// This method blocks until the server has successfully bound to its
49    /// configured address and is ready to accept gRPC requests.
50    ///
51    /// # Panics
52    /// Panics if called more than once, as the start signal is consumed.
53    pub async fn wait_for_start(&mut self) -> Result<()> {
54        self.started_rx
55            .take()
56            .expect("Server start signal already consumed")
57            .await
58            .expect("Failed to receive server start signal");
59        Ok(())
60    }
61
62    /// Waits for the server to completely stop.
63    ///
64    /// This method consumes the handle and blocks until the server task
65    /// has finished executing. Use this after calling `shutdown()` to
66    /// ensure clean termination.
67    ///
68    /// # Panics
69    /// Panics if the server task panicked during execution.
70    pub async fn wait_for_stop(self) -> Result<()> {
71        let handles = self
72            .reporter_handles
73            .into_iter()
74            .chain(std::iter::once(self.join_handle));
75        join_all(handles).await;
76        Ok(())
77    }
78
79    /// Signals the server to begin graceful shutdown.
80    ///
81    /// This method triggers the shutdown process but does not wait for
82    /// completion. Use `wait_for_stop()` to wait for the server to fully stop.
83    pub fn shutdown(&self) { self.cancellation_token.cancel(); }
84
85    /// Checks if the server task has completed.
86    ///
87    /// Returns `true` if the server has finished running, either due to
88    /// shutdown or an error condition.
89    #[must_use]
90    pub fn is_finished(&self) -> bool { self.join_handle.is_finished() }
91}
rsketch_server/lib.rs

rsketch_server/
lib.rs
