TLA  Line data    Source code

       1                 : use crate::{background_process, local_env::LocalEnv};
       2                 : use anyhow::anyhow;
       3                 : use camino::Utf8PathBuf;
       4                 : use serde::{Deserialize, Serialize};
       5                 : use serde_with::{serde_as, DisplayFromStr};
       6                 : use std::{path::PathBuf, process::Child};
       7                 : use utils::id::{NodeId, TenantId};
       8                 : 
       9                 : pub struct AttachmentService {
      10                 :     env: LocalEnv,
      11                 :     listen: String,
      12                 :     path: PathBuf,
      13                 : }
      14                 : 
      15                 : const COMMAND: &str = "attachment_service";
      16                 : 
      17                 : #[serde_as]
      18 CBC          12 : #[derive(Serialize, Deserialize)]
      19                 : pub struct AttachHookRequest {
      20                 :     #[serde_as(as = "DisplayFromStr")]
      21                 :     pub tenant_id: TenantId,
      22                 :     pub pageserver_id: Option<NodeId>,
      23                 : }
      24                 : 
      25              36 : #[derive(Serialize, Deserialize)]
      26                 : pub struct AttachHookResponse {
      27                 :     pub gen: Option<u32>,
      28                 : }
      29                 : 
      30                 : impl AttachmentService {
      31              38 :     pub fn from_env(env: &LocalEnv) -> Self {
      32              38 :         let path = env.base_data_dir.join("attachments.json");
      33              38 : 
      34              38 :         // Makes no sense to construct this if pageservers aren't going to use it: assume
      35              38 :         // pageservers have control plane API set
      36              38 :         let listen_url = env.control_plane_api.clone().unwrap();
      37              38 : 
      38              38 :         let listen = format!(
      39              38 :             "{}:{}",
      40              38 :             listen_url.host_str().unwrap(),
      41              38 :             listen_url.port().unwrap()
      42              38 :         );
      43              38 : 
      44              38 :         Self {
      45              38 :             env: env.clone(),
      46              38 :             path,
      47              38 :             listen,
      48              38 :         }
      49              38 :     }
      50                 : 
      51              26 :     fn pid_file(&self) -> Utf8PathBuf {
      52              26 :         Utf8PathBuf::from_path_buf(self.env.base_data_dir.join("attachment_service.pid"))
      53              26 :             .expect("non-Unicode path")
      54              26 :     }
      55                 : 
      56              13 :     pub fn start(&self) -> anyhow::Result<Child> {
      57              13 :         let path_str = self.path.to_string_lossy();
      58              13 : 
      59              13 :         background_process::start_process(
      60              13 :             COMMAND,
      61              13 :             &self.env.base_data_dir,
      62              13 :             &self.env.attachment_service_bin(),
      63              13 :             ["-l", &self.listen, "-p", &path_str],
      64              13 :             [],
      65              13 :             background_process::InitialPidFile::Create(&self.pid_file()),
      66              13 :             // TODO: a real status check
      67              13 :             || Ok(true),
      68              13 :         )
      69              13 :     }
      70                 : 
      71              13 :     pub fn stop(&self, immediate: bool) -> anyhow::Result<()> {
      72              13 :         background_process::stop_process(immediate, COMMAND, &self.pid_file())
      73              13 :     }
      74                 : 
      75                 :     /// Call into the attach_hook API, for use before handing out attachments to pageservers
      76              12 :     pub fn attach_hook(
      77              12 :         &self,
      78              12 :         tenant_id: TenantId,
      79              12 :         pageserver_id: NodeId,
      80              12 :     ) -> anyhow::Result<Option<u32>> {
      81              12 :         use hyper::StatusCode;
      82              12 : 
      83              12 :         let url = self
      84              12 :             .env
      85              12 :             .control_plane_api
      86              12 :             .clone()
      87              12 :             .unwrap()
      88              12 :             .join("attach_hook")
      89              12 :             .unwrap();
      90              12 :         let client = reqwest::blocking::ClientBuilder::new()
      91              12 :             .build()
      92              12 :             .expect("Failed to construct http client");
      93              12 : 
      94              12 :         let request = AttachHookRequest {
      95              12 :             tenant_id,
      96              12 :             pageserver_id: Some(pageserver_id),
      97              12 :         };
      98                 : 
      99              12 :         let response = client.post(url).json(&request).send()?;
     100              12 :         if response.status() != StatusCode::OK {
     101 UBC           0 :             return Err(anyhow!("Unexpected status {}", response.status()));
     102 CBC          12 :         }
     103                 : 
     104              12 :         let response = response.json::<AttachHookResponse>()?;
     105              12 :         Ok(response.gen)
     106              12 :     }
     107                 : }