LCOV - differential code coverage report
Current view: top level - pageserver/src - task_mgr.rs (source / functions) Coverage Total Hit UBC CBC
Current: cd44433dd675caa99df17a61b18949c8387e2242.info Lines: 88.0 % 249 219 30 219
Current Date: 2024-01-09 02:06:09 Functions: 67.3 % 162 109 53 109
Baseline: 66c52a629a0f4a503e193045e0df4c77139e344b.info
Baseline Date: 2024-01-08 15:34:46

           TLA  Line data    Source code
       1                 : //!
       2                 : //! This module provides centralized handling of tokio tasks in the Page Server.
       3                 : //!
       4                 : //! We provide a few basic facilities:
       5                 : //! - A global registry of tasks that lists what kind of tasks they are, and
       6                 : //!   which tenant or timeline they are working on
       7                 : //!
       8                 : //! - The ability to request a task to shut down.
       9                 : //!
      10                 : //!
      11                 : //! # How it works?
      12                 : //!
      13                 : //! There is a global hashmap of all the tasks (`TASKS`). Whenever a new
      14                 : //! task is spawned, a PageServerTask entry is added there, and when a
      15                 : //! task dies, it removes itself from the hashmap. If you want to kill a
      16                 : //! task, you can scan the hashmap to find it.
      17                 : //!
      18                 : //! # Task shutdown
      19                 : //!
      20                 : //! To kill a task, we rely on co-operation from the victim. Each task is
      21                 : //! expected to periodically call the `is_shutdown_requested()` function, and
      22                 : //! if it returns true, exit gracefully. In addition to that, when waiting for
      23                 : //! the network or other long-running operation, you can use
      24                 : //! `shutdown_watcher()` function to get a Future that will become ready if
      25                 : //! the current task has been requested to shut down. You can use that with
      26                 : //! Tokio select!().
      27                 : //!
      28                 : //! TODO: This would be a good place to also handle panics in a somewhat sane way.
      29                 : //! Depending on what task panics, we might want to kill the whole server, or
      30                 : //! only a single tenant or timeline.
      31                 : //!
      32                 : 
      33                 : // Clippy 1.60 incorrectly complains about the tokio::task_local!() macro.
      34                 : // Silence it. See https://github.com/rust-lang/rust-clippy/issues/9224.
      35                 : #![allow(clippy::declare_interior_mutable_const)]
      36                 : 
      37                 : use std::collections::HashMap;
      38                 : use std::fmt;
      39                 : use std::future::Future;
      40                 : use std::panic::AssertUnwindSafe;
      41                 : use std::sync::atomic::{AtomicU64, Ordering};
      42                 : use std::sync::{Arc, Mutex};
      43                 : 
      44                 : use futures::FutureExt;
      45                 : use pageserver_api::shard::TenantShardId;
      46                 : use tokio::runtime::Runtime;
      47                 : use tokio::task::JoinHandle;
      48                 : use tokio::task_local;
      49                 : use tokio_util::sync::CancellationToken;
      50                 : 
      51                 : use tracing::{debug, error, info, warn};
      52                 : 
      53                 : use once_cell::sync::Lazy;
      54                 : 
      55                 : use utils::id::TimelineId;
      56                 : 
      57                 : use crate::shutdown_pageserver;
      58                 : 
      59                 : //
      60                 : // There are four runtimes:
      61                 : //
      62                 : // Compute request runtime
      63                 : //  - used to handle connections from compute nodes. Any tasks related to satisfying
      64                 : //    GetPage requests, base backups, import, and other such compute node operations
      65                 : //    are handled by the Compute request runtime
      66                 : //  - page_service.rs
      67                 : //  - this includes layer downloads from remote storage, if a layer is needed to
      68                 : //    satisfy a GetPage request
      69                 : //
      70                 : // Management request runtime
      71                 : //  - used to handle HTTP API requests
      72                 : //
      73                 : // WAL receiver runtime:
      74                 : //  - used to handle WAL receiver connections.
      75                 : //  - and to receiver updates from storage_broker
      76                 : //
      77                 : // Background runtime
      78                 : //  - layer flushing
      79                 : //  - garbage collection
      80                 : //  - compaction
      81                 : //  - remote storage uploads
      82                 : //  - initial tenant loading
      83                 : //
      84                 : // Everything runs in a tokio task. If you spawn new tasks, spawn it using the correct
      85                 : // runtime.
      86                 : //
      87                 : // There might be situations when one task needs to wait for a task running in another
      88                 : // Runtime to finish. For example, if a background operation needs a layer from remote
      89                 : // storage, it will start to download it. If a background operation needs a remote layer,
      90                 : // and the download was already initiated by a GetPage request, the background task
      91                 : // will wait for the download - running in the Page server runtime - to finish.
      92                 : // Another example: the initial tenant loading tasks are launched in the background ops
      93                 : // runtime. If a GetPage request comes in before the load of a tenant has finished, the
      94                 : // GetPage request will wait for the tenant load to finish.
      95                 : //
      96                 : // The core Timeline code is synchronous, and uses a bunch of std Mutexes and RWLocks to
      97                 : // protect data structures. Let's keep it that way. Synchronous code is easier to debug
      98                 : // and analyze, and there's a lot of hairy, low-level, performance critical code there.
      99                 : //
     100                 : // It's nice to have different runtimes, so that you can quickly eyeball how much CPU
     101                 : // time each class of operations is taking, with 'top -H' or similar.
     102                 : //
     103                 : // It's also good to avoid hogging all threads that would be needed to process
     104                 : // other operations, if the upload tasks e.g. get blocked on locks. It shouldn't
     105                 : // happen, but still.
     106                 : //
     107 CBC         557 : pub static COMPUTE_REQUEST_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
     108             557 :     tokio::runtime::Builder::new_multi_thread()
     109             557 :         .thread_name("compute request worker")
     110             557 :         .enable_all()
     111             557 :         .build()
     112             557 :         .expect("Failed to create compute request runtime")
     113             557 : });
     114                 : 
     115             557 : pub static MGMT_REQUEST_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
     116             557 :     tokio::runtime::Builder::new_multi_thread()
     117             557 :         .thread_name("mgmt request worker")
     118             557 :         .enable_all()
     119             557 :         .build()
     120             557 :         .expect("Failed to create mgmt request runtime")
     121             557 : });
     122                 : 
     123             562 : pub static WALRECEIVER_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
     124             562 :     tokio::runtime::Builder::new_multi_thread()
     125             562 :         .thread_name("walreceiver worker")
     126             562 :         .enable_all()
     127             562 :         .build()
     128             562 :         .expect("Failed to create walreceiver runtime")
     129             562 : });
     130                 : 
     131             593 : pub static BACKGROUND_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
     132             593 :     tokio::runtime::Builder::new_multi_thread()
     133             593 :         .thread_name("background op worker")
     134             593 :         // if you change the number of worker threads please change the constant below
     135             593 :         .enable_all()
     136             593 :         .build()
     137             593 :         .expect("Failed to create background op runtime")
     138             593 : });
     139                 : 
     140             447 : pub(crate) static BACKGROUND_RUNTIME_WORKER_THREADS: Lazy<usize> = Lazy::new(|| {
     141             447 :     // force init and thus panics
     142             447 :     let _ = BACKGROUND_RUNTIME.handle();
     143             447 :     // replicates tokio-1.28.1::loom::sys::num_cpus which is not available publicly
     144             447 :     // tokio would had already panicked for parsing errors or NotUnicode
     145             447 :     //
     146             447 :     // this will be wrong if any of the runtimes gets their worker threads configured to something
     147             447 :     // else, but that has not been needed in a long time.
     148             447 :     std::env::var("TOKIO_WORKER_THREADS")
     149             447 :         .map(|s| s.parse::<usize>().unwrap())
     150             447 :         .unwrap_or_else(|_e| usize::max(2, num_cpus::get()))
     151             447 : });
     152                 : 
     153 UBC           0 : #[derive(Debug, Clone, Copy)]
     154                 : pub struct PageserverTaskId(u64);
     155                 : 
     156                 : impl fmt::Display for PageserverTaskId {
     157 CBC          18 :     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
     158              18 :         self.0.fmt(f)
     159              18 :     }
     160                 : }
     161                 : 
     162                 : /// Each task that we track is associated with a "task ID". It's just an
     163                 : /// increasing number that we assign. Note that it is different from tokio::task::Id.
     164                 : static NEXT_TASK_ID: AtomicU64 = AtomicU64::new(1);
     165                 : 
     166                 : /// Global registry of tasks
     167                 : static TASKS: Lazy<Mutex<HashMap<u64, Arc<PageServerTask>>>> =
     168             595 :     Lazy::new(|| Mutex::new(HashMap::new()));
     169                 : 
     170              88 : task_local! {
     171              88 :     // This is a cancellation token which will be cancelled when a task needs to shut down. The
     172              88 :     // root token is kept in the global registry, so that anyone can send the signal to request
     173              88 :     // task shutdown.
     174              88 :     static SHUTDOWN_TOKEN: CancellationToken;
     175              88 : 
     176              88 :     // Each task holds reference to its own PageServerTask here.
     177              88 :     static CURRENT_TASK: Arc<PageServerTask>;
     178              88 : }
     179                 : 
     180                 : ///
     181                 : /// There are many kinds of tasks in the system. Some are associated with a particular
     182                 : /// tenant or timeline, while others are global.
     183                 : ///
     184                 : /// Note that we don't try to limit how many task of a certain kind can be running
     185                 : /// at the same time.
     186                 : ///
     187                 : #[derive(
     188               7 :     Debug,
     189                 :     // NB: enumset::EnumSetType derives PartialEq, Eq, Clone, Copy
     190            1192 :     enumset::EnumSetType,
     191       313699139 :     enum_map::Enum,
     192 UBC           0 :     serde::Serialize,
     193               0 :     serde::Deserialize,
     194 CBC       25669 :     strum_macros::IntoStaticStr,
     195                 : )]
     196                 : pub enum TaskKind {
     197                 :     // Pageserver startup, i.e., `main`
     198                 :     Startup,
     199                 : 
     200                 :     // libpq listener task. It just accepts connection and spawns a
     201                 :     // PageRequestHandler task for each connection.
     202                 :     LibpqEndpointListener,
     203                 : 
     204                 :     // HTTP endpoint listener.
     205                 :     HttpEndpointListener,
     206                 : 
     207                 :     // Task that handles a single connection. A PageRequestHandler task
     208                 :     // starts detached from any particular tenant or timeline, but it can be
     209                 :     // associated with one later, after receiving a command from the client.
     210                 :     PageRequestHandler,
     211                 : 
     212                 :     /// Manages the WAL receiver connection for one timeline.
     213                 :     /// It subscribes to events from storage_broker and decides which safekeeper to connect to.
     214                 :     /// Once the decision has been made, it establishes the connection using the `tokio-postgres` library.
     215                 :     /// There is at most one connection at any given time.
     216                 :     ///
     217                 :     /// That `tokio-postgres` library represents a connection as two objects: a `Client` and a `Connection`.
     218                 :     /// The `Client` object is what library users use to make requests & get responses.
     219                 :     /// Internally, `Client` hands over requests to the `Connection` object.
     220                 :     /// The `Connection` object is responsible for speaking the wire protocol.
     221                 :     ///
     222                 :     /// Walreceiver uses its own abstraction called `TaskHandle` to represent the activity of establishing and handling a connection.
     223                 :     /// That abstraction doesn't use `task_mgr`.
     224                 :     /// The `WalReceiverManager` task ensures that this `TaskHandle` task does not outlive the `WalReceiverManager` task.
     225                 :     /// For the `RequestContext` that we hand to the TaskHandle, we use the [`WalReceiverConnectionHandler`] task kind.
     226                 :     ///
     227                 :     /// Once the connection is established, the `TaskHandle` task creates a
     228                 :     /// [`WalReceiverConnectionPoller`] task_mgr task that is responsible for polling
     229                 :     /// the `Connection` object.
     230                 :     /// A `CancellationToken` created by the `TaskHandle` task ensures
     231                 :     /// that the [`WalReceiverConnectionPoller`] task will cancel soon after as the `TaskHandle` is dropped.
     232                 :     ///
     233                 :     /// [`WalReceiverConnectionHandler`]: Self::WalReceiverConnectionHandler
     234                 :     /// [`WalReceiverConnectionPoller`]: Self::WalReceiverConnectionPoller
     235                 :     WalReceiverManager,
     236                 : 
     237                 :     /// The `TaskHandle` task that executes `handle_walreceiver_connection`.
     238                 :     /// Not a `task_mgr` task, but we use this `TaskKind` for its `RequestContext`.
     239                 :     /// See the comment on [`WalReceiverManager`].
     240                 :     ///
     241                 :     /// [`WalReceiverManager`]: Self::WalReceiverManager
     242                 :     WalReceiverConnectionHandler,
     243                 : 
     244                 :     /// The task that polls the `tokio-postgres::Connection` object.
     245                 :     /// Spawned by task [`WalReceiverConnectionHandler`](Self::WalReceiverConnectionHandler).
     246                 :     /// See the comment on [`WalReceiverManager`](Self::WalReceiverManager).
     247                 :     WalReceiverConnectionPoller,
     248                 : 
     249                 :     // Garbage collection worker. One per tenant
     250                 :     GarbageCollector,
     251                 : 
     252                 :     // Compaction. One per tenant.
     253                 :     Compaction,
     254                 : 
     255                 :     // Eviction. One per timeline.
     256                 :     Eviction,
     257                 : 
     258                 :     /// See [`crate::disk_usage_eviction_task`].
     259                 :     DiskUsageEviction,
     260                 : 
     261                 :     /// See [`crate::tenant::secondary`].
     262                 :     SecondaryDownloads,
     263                 : 
     264                 :     /// See [`crate::tenant::secondary`].
     265                 :     SecondaryUploads,
     266                 : 
     267                 :     // Initial logical size calculation
     268                 :     InitialLogicalSizeCalculation,
     269                 : 
     270                 :     OndemandLogicalSizeCalculation,
     271                 : 
     272                 :     // Task that flushes frozen in-memory layers to disk
     273                 :     LayerFlushTask,
     274                 : 
     275                 :     // Task that uploads a file to remote storage
     276                 :     RemoteUploadTask,
     277                 : 
     278                 :     // Task that downloads a file from remote storage
     279                 :     RemoteDownloadTask,
     280                 : 
     281                 :     // task that handles the initial downloading of all tenants
     282                 :     InitialLoad,
     283                 : 
     284                 :     // task that handles attaching a tenant
     285                 :     Attach,
     286                 : 
     287                 :     // Used mostly for background deletion from s3
     288                 :     TimelineDeletionWorker,
     289                 : 
     290                 :     // task that handhes metrics collection
     291                 :     MetricsCollection,
     292                 : 
     293                 :     // task that drives downloading layers
     294                 :     DownloadAllRemoteLayers,
     295                 :     // Task that calculates synthetis size for all active tenants
     296                 :     CalculateSyntheticSize,
     297                 : 
     298                 :     // A request that comes in via the pageserver HTTP API.
     299                 :     MgmtRequest,
     300                 : 
     301                 :     DebugTool,
     302                 : 
     303                 :     #[cfg(test)]
     304                 :     UnitTest,
     305                 : }
     306                 : 
     307 UBC           0 : #[derive(Default)]
     308                 : struct MutableTaskState {
     309                 :     /// Handle for waiting for the task to exit. It can be None, if the
     310                 :     /// the task has already exited.
     311                 :     join_handle: Option<JoinHandle<()>>,
     312                 : }
     313                 : 
     314                 : struct PageServerTask {
     315                 :     #[allow(dead_code)] // unused currently
     316                 :     task_id: PageserverTaskId,
     317                 : 
     318                 :     kind: TaskKind,
     319                 : 
     320                 :     name: String,
     321                 : 
     322                 :     // To request task shutdown, just cancel this token.
     323                 :     cancel: CancellationToken,
     324                 : 
     325                 :     /// Tasks may optionally be launched for a particular tenant/timeline, enabling
     326                 :     /// later cancelling tasks for that tenant/timeline in [`shutdown_tasks`]
     327                 :     tenant_shard_id: Option<TenantShardId>,
     328                 :     timeline_id: Option<TimelineId>,
     329                 : 
     330                 :     mutable: Mutex<MutableTaskState>,
     331                 : }
     332                 : 
     333                 : /// Launch a new task
     334                 : /// Note: if shutdown_process_on_error is set to true failure
     335                 : ///   of the task will lead to shutdown of entire process
     336 CBC       55517 : pub fn spawn<F>(
     337           55517 :     runtime: &tokio::runtime::Handle,
     338           55517 :     kind: TaskKind,
     339           55517 :     tenant_shard_id: Option<TenantShardId>,
     340           55517 :     timeline_id: Option<TimelineId>,
     341           55517 :     name: &str,
     342           55517 :     shutdown_process_on_error: bool,
     343           55517 :     future: F,
     344           55517 : ) -> PageserverTaskId
     345           55517 : where
     346           55517 :     F: Future<Output = anyhow::Result<()>> + Send + 'static,
     347           55517 : {
     348           55517 :     let cancel = CancellationToken::new();
     349           55517 :     let task_id = NEXT_TASK_ID.fetch_add(1, Ordering::Relaxed);
     350           55517 :     let task = Arc::new(PageServerTask {
     351           55517 :         task_id: PageserverTaskId(task_id),
     352           55517 :         kind,
     353           55517 :         name: name.to_string(),
     354           55517 :         cancel: cancel.clone(),
     355           55517 :         tenant_shard_id,
     356           55517 :         timeline_id,
     357           55517 :         mutable: Mutex::new(MutableTaskState { join_handle: None }),
     358           55517 :     });
     359           55517 : 
     360           55517 :     TASKS.lock().unwrap().insert(task_id, Arc::clone(&task));
     361           55517 : 
     362           55517 :     let mut task_mut = task.mutable.lock().unwrap();
     363           55517 : 
     364           55517 :     let task_name = name.to_string();
     365           55517 :     let task_cloned = Arc::clone(&task);
     366           55517 :     let join_handle = runtime.spawn(task_wrapper(
     367           55517 :         task_name,
     368           55517 :         task_id,
     369           55517 :         task_cloned,
     370           55517 :         cancel,
     371           55517 :         shutdown_process_on_error,
     372           55517 :         future,
     373           55517 :     ));
     374           55517 :     task_mut.join_handle = Some(join_handle);
     375           55517 :     drop(task_mut);
     376           55517 : 
     377           55517 :     // The task is now running. Nothing more to do here
     378           55517 :     PageserverTaskId(task_id)
     379           55517 : }
     380                 : 
     381                 : /// This wrapper function runs in a newly-spawned task. It initializes the
     382                 : /// task-local variables and calls the payload function.
     383           55517 : async fn task_wrapper<F>(
     384           55517 :     task_name: String,
     385           55517 :     task_id: u64,
     386           55517 :     task: Arc<PageServerTask>,
     387           55517 :     shutdown_token: CancellationToken,
     388           55517 :     shutdown_process_on_error: bool,
     389           55517 :     future: F,
     390           55517 : ) where
     391           55517 :     F: Future<Output = anyhow::Result<()>> + Send + 'static,
     392           55517 : {
     393 UBC           0 :     debug!("Starting task '{}'", task_name);
     394                 : 
     395 CBC       55508 :     let result = SHUTDOWN_TOKEN
     396           55508 :         .scope(
     397           55508 :             shutdown_token,
     398           55508 :             CURRENT_TASK.scope(task, {
     399           55508 :                 // We use AssertUnwindSafe here so that the payload function
     400           55508 :                 // doesn't need to be UnwindSafe. We don't do anything after the
     401           55508 :                 // unwinding that would expose us to unwind-unsafe behavior.
     402           55508 :                 AssertUnwindSafe(future).catch_unwind()
     403           55508 :             }),
     404           55508 :         )
     405         8193618 :         .await;
     406           50937 :     task_finish(result, task_name, task_id, shutdown_process_on_error).await;
     407           50937 : }
     408                 : 
     409           50937 : async fn task_finish(
     410           50937 :     result: std::result::Result<
     411           50937 :         anyhow::Result<()>,
     412           50937 :         std::boxed::Box<dyn std::any::Any + std::marker::Send>,
     413           50937 :     >,
     414           50937 :     task_name: String,
     415           50937 :     task_id: u64,
     416           50937 :     shutdown_process_on_error: bool,
     417           50937 : ) {
     418           50937 :     // Remove our entry from the global hashmap.
     419           50937 :     let task = TASKS
     420           50937 :         .lock()
     421           50937 :         .unwrap()
     422           50937 :         .remove(&task_id)
     423           50937 :         .expect("no task in registry");
     424           50937 : 
     425           50937 :     let mut shutdown_process = false;
     426                 :     {
     427           50937 :         match result {
     428                 :             Ok(Ok(())) => {
     429 UBC           0 :                 debug!("Task '{}' exited normally", task_name);
     430                 :             }
     431 CBC           8 :             Ok(Err(err)) => {
     432               8 :                 if shutdown_process_on_error {
     433 UBC           0 :                     error!(
     434               0 :                         "Shutting down: task '{}' tenant_shard_id: {:?}, timeline_id: {:?} exited with error: {:?}",
     435               0 :                         task_name, task.tenant_shard_id, task.timeline_id, err
     436               0 :                     );
     437               0 :                     shutdown_process = true;
     438                 :                 } else {
     439 CBC           8 :                     error!(
     440               8 :                         "Task '{}' tenant_shard_id: {:?}, timeline_id: {:?} exited with error: {:?}",
     441               8 :                         task_name, task.tenant_shard_id, task.timeline_id, err
     442               8 :                     );
     443                 :                 }
     444                 :             }
     445 UBC           0 :             Err(err) => {
     446               0 :                 if shutdown_process_on_error {
     447               0 :                     error!(
     448               0 :                         "Shutting down: task '{}' tenant_shard_id: {:?}, timeline_id: {:?} panicked: {:?}",
     449               0 :                         task_name, task.tenant_shard_id, task.timeline_id, err
     450               0 :                     );
     451               0 :                     shutdown_process = true;
     452                 :                 } else {
     453               0 :                     error!(
     454               0 :                         "Task '{}' tenant_shard_id: {:?}, timeline_id: {:?} panicked: {:?}",
     455               0 :                         task_name, task.tenant_shard_id, task.timeline_id, err
     456               0 :                     );
     457                 :                 }
     458                 :             }
     459                 :         }
     460                 :     }
     461                 : 
     462 CBC       50937 :     if shutdown_process {
     463 UBC           0 :         shutdown_pageserver(None, 1).await;
     464 CBC       50937 :     }
     465           50937 : }
     466                 : 
     467                 : /// Signal and wait for tasks to shut down.
     468                 : ///
     469                 : ///
     470                 : /// The arguments are used to select the tasks to kill. Any None arguments are
     471                 : /// ignored. For example, to shut down all WalReceiver tasks:
     472                 : ///
     473                 : ///   shutdown_tasks(Some(TaskKind::WalReceiver), None, None)
     474                 : ///
     475                 : /// Or to shut down all tasks for given timeline:
     476                 : ///
     477                 : ///   shutdown_tasks(None, Some(tenant_shard_id), Some(timeline_id))
     478                 : ///
     479            2747 : pub async fn shutdown_tasks(
     480            2747 :     kind: Option<TaskKind>,
     481            2747 :     tenant_shard_id: Option<TenantShardId>,
     482            2747 :     timeline_id: Option<TimelineId>,
     483            2747 : ) {
     484            2747 :     let mut victim_tasks = Vec::new();
     485            2747 : 
     486            2747 :     {
     487            2747 :         let tasks = TASKS.lock().unwrap();
     488           32346 :         for task in tasks.values() {
     489           32346 :             if (kind.is_none() || Some(task.kind) == kind)
     490           14766 :                 && (tenant_shard_id.is_none() || task.tenant_shard_id == tenant_shard_id)
     491            9363 :                 && (timeline_id.is_none() || task.timeline_id == timeline_id)
     492            5123 :             {
     493            5123 :                 task.cancel.cancel();
     494            5123 :                 victim_tasks.push((
     495            5123 :                     Arc::clone(task),
     496            5123 :                     task.kind,
     497            5123 :                     task.tenant_shard_id,
     498            5123 :                     task.timeline_id,
     499            5123 :                 ));
     500           27223 :             }
     501                 :         }
     502                 :     }
     503                 : 
     504            2747 :     let log_all = kind.is_none() && tenant_shard_id.is_none() && timeline_id.is_none();
     505                 : 
     506            7870 :     for (task, task_kind, tenant_shard_id, timeline_id) in victim_tasks {
     507            5123 :         let join_handle = {
     508            5123 :             let mut task_mut = task.mutable.lock().unwrap();
     509            5123 :             task_mut.join_handle.take()
     510                 :         };
     511            5123 :         if let Some(mut join_handle) = join_handle {
     512            5123 :             if log_all {
     513               6 :                 if tenant_shard_id.is_none() {
     514                 :                     // there are quite few of these
     515               6 :                     info!(name = task.name, kind = ?task_kind, "stopping global task");
     516                 :                 } else {
     517                 :                     // warn to catch these in tests; there shouldn't be any
     518 UBC           0 :                     warn!(name = task.name, tenant_shard_id = ?tenant_shard_id, timeline_id = ?timeline_id, kind = ?task_kind, "stopping left-over");
     519                 :                 }
     520 CBC        5117 :             }
     521            5123 :             if tokio::time::timeout(std::time::Duration::from_secs(1), &mut join_handle)
     522            2014 :                 .await
     523            5123 :                 .is_err()
     524                 :             {
     525                 :                 // allow some time to elapse before logging to cut down the number of log
     526                 :                 // lines.
     527               8 :                 info!("waiting for task {} to shut down", task.name);
     528                 :                 // we never handled this return value, but:
     529                 :                 // - we don't deschedule which would lead to is_cancelled
     530                 :                 // - panics are already logged (is_panicked)
     531                 :                 // - task errors are already logged in the wrapper
     532               8 :                 let _ = join_handle.await;
     533               8 :                 info!("task {} completed", task.name);
     534            5115 :             }
     535 UBC           0 :         } else {
     536               0 :             // Possibly one of:
     537               0 :             //  * The task had not even fully started yet.
     538               0 :             //  * It was shut down concurrently and already exited
     539               0 :         }
     540                 :     }
     541 CBC        2747 : }
     542                 : 
     543         4258380 : pub fn current_task_kind() -> Option<TaskKind> {
     544         4258380 :     CURRENT_TASK.try_with(|ct| ct.kind).ok()
     545         4258380 : }
     546                 : 
     547              16 : pub fn current_task_id() -> Option<PageserverTaskId> {
     548              16 :     CURRENT_TASK.try_with(|ct| ct.task_id).ok()
     549              16 : }
     550                 : 
     551                 : /// A Future that can be used to check if the current task has been requested to
     552                 : /// shut down.
     553           56895 : pub async fn shutdown_watcher() {
     554           52770 :     let token = SHUTDOWN_TOKEN
     555           52770 :         .try_with(|t| t.clone())
     556           52770 :         .expect("shutdown_watcher() called in an unexpected task or thread");
     557           52770 : 
     558          615687 :     token.cancelled().await;
     559             361 : }
     560                 : 
     561                 : /// Clone the current task's cancellation token, which can be moved across tasks.
     562                 : ///
     563                 : /// When the task which is currently executing is shutdown, the cancellation token will be
     564                 : /// cancelled. It can however be moved to other tasks, such as `tokio::task::spawn_blocking` or
     565                 : /// `tokio::task::JoinSet::spawn`.
     566           26464 : pub fn shutdown_token() -> CancellationToken {
     567           26464 :     let res = SHUTDOWN_TOKEN.try_with(|t| t.clone());
     568           26464 : 
     569           26464 :     if cfg!(test) {
     570                 :         // in tests this method is called from non-taskmgr spawned tasks, and that is all ok.
     571             154 :         res.unwrap_or_default()
     572                 :     } else {
     573           26310 :         res.expect("shutdown_token() called in an unexpected task or thread")
     574                 :     }
     575           26464 : }
     576                 : 
     577                 : /// Has the current task been requested to shut down?
     578           31261 : pub fn is_shutdown_requested() -> bool {
     579           31261 :     if let Ok(cancel) = SHUTDOWN_TOKEN.try_with(|t| t.clone()) {
     580           31257 :         cancel.is_cancelled()
     581                 :     } else {
     582               4 :         if !cfg!(test) {
     583 UBC           0 :             warn!("is_shutdown_requested() called in an unexpected task or thread");
     584 CBC           4 :         }
     585               4 :         false
     586                 :     }
     587           31261 : }
        

Generated by: LCOV version 2.1-beta