Line data Source code
1 : pub(crate) mod split_state;
2 : use std::collections::HashMap;
3 : use std::io::Write;
4 : use std::str::FromStr;
5 : use std::sync::Arc;
6 : use std::time::{Duration, Instant};
7 :
8 : use diesel::deserialize::{FromSql, FromSqlRow};
9 : use diesel::expression::AsExpression;
10 : use diesel::pg::Pg;
11 : use diesel::prelude::*;
12 : use diesel::serialize::{IsNull, ToSql};
13 : use diesel_async::async_connection_wrapper::AsyncConnectionWrapper;
14 : use diesel_async::pooled_connection::bb8::Pool;
15 : use diesel_async::pooled_connection::{AsyncDieselConnectionManager, ManagerConfig};
16 : use diesel_async::{AsyncPgConnection, RunQueryDsl};
17 : use diesel_migrations::{EmbeddedMigrations, embed_migrations};
18 : use futures::FutureExt;
19 : use futures::future::BoxFuture;
20 : use itertools::Itertools;
21 : use pageserver_api::controller_api::{
22 : AvailabilityZone, MetadataHealthRecord, NodeLifecycle, NodeSchedulingPolicy, PlacementPolicy,
23 : SafekeeperDescribeResponse, ShardSchedulingPolicy, SkSchedulingPolicy,
24 : };
25 : use pageserver_api::models::{ShardImportStatus, TenantConfig};
26 : use pageserver_api::shard::{
27 : ShardConfigError, ShardCount, ShardIdentity, ShardNumber, ShardStripeSize, TenantShardId,
28 : };
29 : use rustls::client::WebPkiServerVerifier;
30 : use rustls::client::danger::{ServerCertVerified, ServerCertVerifier};
31 : use rustls::crypto::ring;
32 : use scoped_futures::ScopedBoxFuture;
33 : use serde::{Deserialize, Serialize};
34 : use utils::generation::Generation;
35 : use utils::id::{NodeId, TenantId, TimelineId};
36 : use utils::lsn::Lsn;
37 :
38 : use self::split_state::SplitState;
39 : use crate::metrics::{
40 : DatabaseQueryErrorLabelGroup, DatabaseQueryLatencyLabelGroup, METRICS_REGISTRY,
41 : };
42 : use crate::node::Node;
43 : use crate::timeline_import::{
44 : TimelineImport, TimelineImportUpdateError, TimelineImportUpdateFollowUp,
45 : };
46 : const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations");
47 :
48 : /// ## What do we store?
49 : ///
50 : /// The storage controller service does not store most of its state durably.
51 : ///
52 : /// The essential things to store durably are:
53 : /// - generation numbers, as these must always advance monotonically to ensure data safety.
54 : /// - Tenant's PlacementPolicy and TenantConfig, as the source of truth for these is something external.
55 : /// - Node's scheduling policies, as the source of truth for these is something external.
56 : ///
57 : /// Other things we store durably as an implementation detail:
58 : /// - Node's host/port: this could be avoided it we made nodes emit a self-registering heartbeat,
59 : /// but it is operationally simpler to make this service the authority for which nodes
60 : /// it talks to.
61 : ///
62 : /// ## Performance/efficiency
63 : ///
64 : /// The storage controller service does not go via the database for most things: there are
65 : /// a couple of places where we must, and where efficiency matters:
66 : /// - Incrementing generation numbers: the Reconciler has to wait for this to complete
67 : /// before it can attach a tenant, so this acts as a bound on how fast things like
68 : /// failover can happen.
69 : /// - Pageserver re-attach: we will increment many shards' generations when this happens,
70 : /// so it is important to avoid e.g. issuing O(N) queries.
71 : ///
72 : /// Database calls relating to nodes have low performance requirements, as they are very rarely
73 : /// updated, and reads of nodes are always from memory, not the database. We only require that
74 : /// we can UPDATE a node's scheduling mode reasonably quickly to mark a bad node offline.
75 : pub struct Persistence {
76 : connection_pool: Pool<AsyncPgConnection>,
77 : }
78 :
79 : /// Legacy format, for use in JSON compat objects in test environment
80 0 : #[derive(Serialize, Deserialize)]
81 : struct JsonPersistence {
82 : tenants: HashMap<TenantShardId, TenantShardPersistence>,
83 : }
84 :
85 : #[derive(thiserror::Error, Debug)]
86 : pub(crate) enum DatabaseError {
87 : #[error(transparent)]
88 : Query(#[from] diesel::result::Error),
89 : #[error(transparent)]
90 : Connection(#[from] diesel::result::ConnectionError),
91 : #[error(transparent)]
92 : ConnectionPool(#[from] diesel_async::pooled_connection::bb8::RunError),
93 : #[error("Logical error: {0}")]
94 : Logical(String),
95 : #[error("Migration error: {0}")]
96 : Migration(String),
97 : }
98 :
99 : #[derive(measured::FixedCardinalityLabel, Copy, Clone)]
100 : pub(crate) enum DatabaseOperation {
101 : InsertNode,
102 : UpdateNode,
103 : DeleteNode,
104 : ListNodes,
105 : ListTombstones,
106 : BeginShardSplit,
107 : CompleteShardSplit,
108 : AbortShardSplit,
109 : Detach,
110 : ReAttach,
111 : IncrementGeneration,
112 : TenantGenerations,
113 : ShardGenerations,
114 : ListTenantShards,
115 : LoadTenant,
116 : InsertTenantShards,
117 : UpdateTenantShard,
118 : DeleteTenant,
119 : UpdateTenantConfig,
120 : UpdateMetadataHealth,
121 : ListMetadataHealth,
122 : ListMetadataHealthUnhealthy,
123 : ListMetadataHealthOutdated,
124 : ListSafekeepers,
125 : GetLeader,
126 : UpdateLeader,
127 : SetPreferredAzs,
128 : InsertTimeline,
129 : GetTimeline,
130 : InsertTimelineReconcile,
131 : RemoveTimelineReconcile,
132 : ListTimelineReconcile,
133 : ListTimelineReconcileStartup,
134 : InsertTimelineImport,
135 : UpdateTimelineImport,
136 : DeleteTimelineImport,
137 : ListTimelineImports,
138 : IsTenantImportingTimeline,
139 : }
140 :
141 : #[must_use]
142 : pub(crate) enum AbortShardSplitStatus {
143 : /// We aborted the split in the database by reverting to the parent shards
144 : Aborted,
145 : /// The split had already been persisted.
146 : Complete,
147 : }
148 :
149 : pub(crate) type DatabaseResult<T> = Result<T, DatabaseError>;
150 :
151 : /// Some methods can operate on either a whole tenant or a single shard
152 : #[derive(Clone)]
153 : pub(crate) enum TenantFilter {
154 : Tenant(TenantId),
155 : Shard(TenantShardId),
156 : }
157 :
158 : /// Represents the results of looking up generation+pageserver for the shards of a tenant
159 : pub(crate) struct ShardGenerationState {
160 : pub(crate) tenant_shard_id: TenantShardId,
161 : pub(crate) generation: Option<Generation>,
162 : pub(crate) generation_pageserver: Option<NodeId>,
163 : }
164 :
165 : // A generous allowance for how many times we may retry serializable transactions
166 : // before giving up. This is not expected to be hit: it is a defensive measure in case we
167 : // somehow engineer a situation where duelling transactions might otherwise live-lock.
168 : const MAX_RETRIES: usize = 128;
169 :
170 : impl Persistence {
171 : // The default postgres connection limit is 100. We use up to 99, to leave one free for a human admin under
172 : // normal circumstances. This assumes we have exclusive use of the database cluster to which we connect.
173 : pub const MAX_CONNECTIONS: u32 = 99;
174 :
175 : // We don't want to keep a lot of connections alive: close them down promptly if they aren't being used.
176 : const IDLE_CONNECTION_TIMEOUT: Duration = Duration::from_secs(10);
177 : const MAX_CONNECTION_LIFETIME: Duration = Duration::from_secs(60);
178 :
179 0 : pub async fn new(database_url: String) -> Self {
180 0 : let mut mgr_config = ManagerConfig::default();
181 0 : mgr_config.custom_setup = Box::new(establish_connection_rustls);
182 0 :
183 0 : let manager = AsyncDieselConnectionManager::<AsyncPgConnection>::new_with_config(
184 0 : database_url,
185 0 : mgr_config,
186 0 : );
187 :
188 : // We will use a connection pool: this is primarily to _limit_ our connection count, rather than to optimize time
189 : // to execute queries (database queries are not generally on latency-sensitive paths).
190 0 : let connection_pool = Pool::builder()
191 0 : .max_size(Self::MAX_CONNECTIONS)
192 0 : .max_lifetime(Some(Self::MAX_CONNECTION_LIFETIME))
193 0 : .idle_timeout(Some(Self::IDLE_CONNECTION_TIMEOUT))
194 0 : // Always keep at least one connection ready to go
195 0 : .min_idle(Some(1))
196 0 : .test_on_check_out(true)
197 0 : .build(manager)
198 0 : .await
199 0 : .expect("Could not build connection pool");
200 0 :
201 0 : Self { connection_pool }
202 0 : }
203 :
204 : /// A helper for use during startup, where we would like to tolerate concurrent restarts of the
205 : /// database and the storage controller, therefore the database might not be available right away
206 0 : pub async fn await_connection(
207 0 : database_url: &str,
208 0 : timeout: Duration,
209 0 : ) -> Result<(), diesel::ConnectionError> {
210 0 : let started_at = Instant::now();
211 0 : log_postgres_connstr_info(database_url)
212 0 : .map_err(|e| diesel::ConnectionError::InvalidConnectionUrl(e.to_string()))?;
213 : loop {
214 0 : match establish_connection_rustls(database_url).await {
215 : Ok(_) => {
216 0 : tracing::info!("Connected to database.");
217 0 : return Ok(());
218 : }
219 0 : Err(e) => {
220 0 : if started_at.elapsed() > timeout {
221 0 : return Err(e);
222 : } else {
223 0 : tracing::info!("Database not yet available, waiting... ({e})");
224 0 : tokio::time::sleep(Duration::from_millis(100)).await;
225 : }
226 : }
227 : }
228 : }
229 0 : }
230 :
231 : /// Execute the diesel migrations that are built into this binary
232 0 : pub(crate) async fn migration_run(&self) -> DatabaseResult<()> {
233 : use diesel_migrations::{HarnessWithOutput, MigrationHarness};
234 :
235 : // Can't use self.with_conn here as we do spawn_blocking which requires static.
236 0 : let conn = self
237 0 : .connection_pool
238 0 : .dedicated_connection()
239 0 : .await
240 0 : .map_err(|e| DatabaseError::Migration(e.to_string()))?;
241 0 : let mut async_wrapper: AsyncConnectionWrapper<AsyncPgConnection> =
242 0 : AsyncConnectionWrapper::from(conn);
243 0 : tokio::task::spawn_blocking(move || {
244 0 : let mut retry_count = 0;
245 0 : loop {
246 0 : let result = HarnessWithOutput::write_to_stdout(&mut async_wrapper)
247 0 : .run_pending_migrations(MIGRATIONS)
248 0 : .map(|_| ())
249 0 : .map_err(|e| DatabaseError::Migration(e.to_string()));
250 0 : match result {
251 0 : Ok(r) => break Ok(r),
252 : Err(
253 0 : err @ DatabaseError::Query(diesel::result::Error::DatabaseError(
254 0 : diesel::result::DatabaseErrorKind::SerializationFailure,
255 0 : _,
256 0 : )),
257 0 : ) => {
258 0 : retry_count += 1;
259 0 : if retry_count > MAX_RETRIES {
260 0 : tracing::error!(
261 0 : "Exceeded max retries on SerializationFailure errors: {err:?}"
262 : );
263 0 : break Err(err);
264 : } else {
265 : // Retry on serialization errors: these are expected, because even though our
266 : // transactions don't fight for the same rows, they will occasionally collide
267 : // on index pages (e.g. increment_generation for unrelated shards can collide)
268 0 : tracing::debug!(
269 0 : "Retrying transaction on serialization failure {err:?}"
270 : );
271 0 : continue;
272 : }
273 : }
274 0 : Err(e) => break Err(e),
275 : }
276 : }
277 0 : })
278 0 : .await
279 0 : .map_err(|e| DatabaseError::Migration(e.to_string()))??;
280 0 : Ok(())
281 0 : }
282 :
283 : /// Wraps `with_conn` in order to collect latency and error metrics
284 0 : async fn with_measured_conn<'a, 'b, F, R>(
285 0 : &self,
286 0 : op: DatabaseOperation,
287 0 : func: F,
288 0 : ) -> DatabaseResult<R>
289 0 : where
290 0 : F: for<'r> Fn(&'r mut AsyncPgConnection) -> ScopedBoxFuture<'b, 'r, DatabaseResult<R>>
291 0 : + Send
292 0 : + std::marker::Sync
293 0 : + 'a,
294 0 : R: Send + 'b,
295 0 : {
296 0 : let latency = &METRICS_REGISTRY
297 0 : .metrics_group
298 0 : .storage_controller_database_query_latency;
299 0 : let _timer = latency.start_timer(DatabaseQueryLatencyLabelGroup { operation: op });
300 :
301 0 : let res = self.with_conn(func).await;
302 :
303 0 : if let Err(err) = &res {
304 0 : let error_counter = &METRICS_REGISTRY
305 0 : .metrics_group
306 0 : .storage_controller_database_query_error;
307 0 : error_counter.inc(DatabaseQueryErrorLabelGroup {
308 0 : error_type: err.error_label(),
309 0 : operation: op,
310 0 : })
311 0 : }
312 :
313 0 : res
314 0 : }
315 :
316 : /// Call the provided function with a Diesel database connection in a retry loop
317 0 : async fn with_conn<'a, 'b, F, R>(&self, func: F) -> DatabaseResult<R>
318 0 : where
319 0 : F: for<'r> Fn(&'r mut AsyncPgConnection) -> ScopedBoxFuture<'b, 'r, DatabaseResult<R>>
320 0 : + Send
321 0 : + std::marker::Sync
322 0 : + 'a,
323 0 : R: Send + 'b,
324 0 : {
325 0 : let mut retry_count = 0;
326 : loop {
327 0 : let mut conn = self.connection_pool.get().await?;
328 0 : match conn
329 0 : .build_transaction()
330 0 : .serializable()
331 0 : .run(|c| func(c))
332 0 : .await
333 : {
334 0 : Ok(r) => break Ok(r),
335 : Err(
336 0 : err @ DatabaseError::Query(diesel::result::Error::DatabaseError(
337 0 : diesel::result::DatabaseErrorKind::SerializationFailure,
338 0 : _,
339 0 : )),
340 0 : ) => {
341 0 : retry_count += 1;
342 0 : if retry_count > MAX_RETRIES {
343 0 : tracing::error!(
344 0 : "Exceeded max retries on SerializationFailure errors: {err:?}"
345 : );
346 0 : break Err(err);
347 : } else {
348 : // Retry on serialization errors: these are expected, because even though our
349 : // transactions don't fight for the same rows, they will occasionally collide
350 : // on index pages (e.g. increment_generation for unrelated shards can collide)
351 0 : tracing::debug!("Retrying transaction on serialization failure {err:?}");
352 0 : continue;
353 : }
354 : }
355 0 : Err(e) => break Err(e),
356 : }
357 : }
358 0 : }
359 :
360 : /// When a node is first registered, persist it before using it for anything
361 : /// If the provided node_id already exists, it will be error.
362 : /// The common case is when a node marked for deletion wants to register.
363 0 : pub(crate) async fn insert_node(&self, node: &Node) -> DatabaseResult<()> {
364 0 : let np = &node.to_persistent();
365 0 : self.with_measured_conn(DatabaseOperation::InsertNode, move |conn| {
366 0 : Box::pin(async move {
367 0 : diesel::insert_into(crate::schema::nodes::table)
368 0 : .values(np)
369 0 : .execute(conn)
370 0 : .await?;
371 0 : Ok(())
372 0 : })
373 0 : })
374 0 : .await
375 0 : }
376 :
377 : /// At startup, populate the list of nodes which our shards may be placed on
378 0 : pub(crate) async fn list_nodes(&self) -> DatabaseResult<Vec<NodePersistence>> {
379 : use crate::schema::nodes::dsl::*;
380 :
381 0 : let result: Vec<NodePersistence> = self
382 0 : .with_measured_conn(DatabaseOperation::ListNodes, move |conn| {
383 0 : Box::pin(async move {
384 0 : Ok(crate::schema::nodes::table
385 0 : .filter(lifecycle.ne(String::from(NodeLifecycle::Deleted)))
386 0 : .load::<NodePersistence>(conn)
387 0 : .await?)
388 0 : })
389 0 : })
390 0 : .await?;
391 :
392 0 : tracing::info!("list_nodes: loaded {} nodes", result.len());
393 :
394 0 : Ok(result)
395 0 : }
396 :
397 0 : pub(crate) async fn list_tombstones(&self) -> DatabaseResult<Vec<NodePersistence>> {
398 : use crate::schema::nodes::dsl::*;
399 :
400 0 : let result: Vec<NodePersistence> = self
401 0 : .with_measured_conn(DatabaseOperation::ListTombstones, move |conn| {
402 0 : Box::pin(async move {
403 0 : Ok(crate::schema::nodes::table
404 0 : .filter(lifecycle.eq(String::from(NodeLifecycle::Deleted)))
405 0 : .load::<NodePersistence>(conn)
406 0 : .await?)
407 0 : })
408 0 : })
409 0 : .await?;
410 :
411 0 : tracing::info!("list_tombstones: loaded {} nodes", result.len());
412 :
413 0 : Ok(result)
414 0 : }
415 :
416 0 : pub(crate) async fn update_node<V>(
417 0 : &self,
418 0 : input_node_id: NodeId,
419 0 : values: V,
420 0 : ) -> DatabaseResult<()>
421 0 : where
422 0 : V: diesel::AsChangeset<Target = crate::schema::nodes::table> + Clone + Send + Sync,
423 0 : V::Changeset: diesel::query_builder::QueryFragment<diesel::pg::Pg> + Send, // valid Postgres SQL
424 0 : {
425 : use crate::schema::nodes::dsl::*;
426 0 : let updated = self
427 0 : .with_measured_conn(DatabaseOperation::UpdateNode, move |conn| {
428 0 : let values = values.clone();
429 0 : Box::pin(async move {
430 0 : let updated = diesel::update(nodes)
431 0 : .filter(node_id.eq(input_node_id.0 as i64))
432 0 : .filter(lifecycle.ne(String::from(NodeLifecycle::Deleted)))
433 0 : .set(values)
434 0 : .execute(conn)
435 0 : .await?;
436 0 : Ok(updated)
437 0 : })
438 0 : })
439 0 : .await?;
440 :
441 0 : if updated != 1 {
442 0 : Err(DatabaseError::Logical(format!(
443 0 : "Node {node_id:?} not found for update",
444 0 : )))
445 : } else {
446 0 : Ok(())
447 : }
448 0 : }
449 :
450 0 : pub(crate) async fn update_node_scheduling_policy(
451 0 : &self,
452 0 : input_node_id: NodeId,
453 0 : input_scheduling: NodeSchedulingPolicy,
454 0 : ) -> DatabaseResult<()> {
455 : use crate::schema::nodes::dsl::*;
456 0 : self.update_node(
457 0 : input_node_id,
458 0 : scheduling_policy.eq(String::from(input_scheduling)),
459 0 : )
460 0 : .await
461 0 : }
462 :
463 0 : pub(crate) async fn update_node_on_registration(
464 0 : &self,
465 0 : input_node_id: NodeId,
466 0 : input_https_port: Option<u16>,
467 0 : ) -> DatabaseResult<()> {
468 : use crate::schema::nodes::dsl::*;
469 0 : self.update_node(
470 0 : input_node_id,
471 0 : listen_https_port.eq(input_https_port.map(|x| x as i32)),
472 0 : )
473 0 : .await
474 0 : }
475 :
476 : /// Tombstone is a special state where the node is not deleted from the database,
477 : /// but it is not available for usage.
478 : /// The main reason for it is to prevent the flaky node to register.
479 0 : pub(crate) async fn set_tombstone(&self, del_node_id: NodeId) -> DatabaseResult<()> {
480 : use crate::schema::nodes::dsl::*;
481 0 : self.update_node(
482 0 : del_node_id,
483 0 : lifecycle.eq(String::from(NodeLifecycle::Deleted)),
484 0 : )
485 0 : .await
486 0 : }
487 :
488 0 : pub(crate) async fn delete_node(&self, del_node_id: NodeId) -> DatabaseResult<()> {
489 : use crate::schema::nodes::dsl::*;
490 0 : self.with_measured_conn(DatabaseOperation::DeleteNode, move |conn| {
491 0 : Box::pin(async move {
492 : // You can hard delete a node only if it has a tombstone.
493 : // So we need to check if the node has lifecycle set to deleted.
494 0 : let node_to_delete = nodes
495 0 : .filter(node_id.eq(del_node_id.0 as i64))
496 0 : .first::<NodePersistence>(conn)
497 0 : .await
498 0 : .optional()?;
499 :
500 0 : if let Some(np) = node_to_delete {
501 0 : let lc = NodeLifecycle::from_str(&np.lifecycle).map_err(|e| {
502 0 : DatabaseError::Logical(format!(
503 0 : "Node {} has invalid lifecycle: {}",
504 0 : del_node_id, e
505 0 : ))
506 0 : })?;
507 :
508 0 : if lc != NodeLifecycle::Deleted {
509 0 : return Err(DatabaseError::Logical(format!(
510 0 : "Node {} was not soft deleted before, cannot hard delete it",
511 0 : del_node_id
512 0 : )));
513 0 : }
514 0 :
515 0 : diesel::delete(nodes)
516 0 : .filter(node_id.eq(del_node_id.0 as i64))
517 0 : .execute(conn)
518 0 : .await?;
519 0 : }
520 :
521 0 : Ok(())
522 0 : })
523 0 : })
524 0 : .await
525 0 : }
526 :
527 : /// At startup, load the high level state for shards, such as their config + policy. This will
528 : /// be enriched at runtime with state discovered on pageservers.
529 : ///
530 : /// We exclude shards configured to be detached. During startup, if we see any attached locations
531 : /// for such shards, they will automatically be detached as 'orphans'.
532 0 : pub(crate) async fn load_active_tenant_shards(
533 0 : &self,
534 0 : ) -> DatabaseResult<Vec<TenantShardPersistence>> {
535 : use crate::schema::tenant_shards::dsl::*;
536 0 : self.with_measured_conn(DatabaseOperation::ListTenantShards, move |conn| {
537 0 : Box::pin(async move {
538 0 : let query = tenant_shards.filter(
539 0 : placement_policy.ne(serde_json::to_string(&PlacementPolicy::Detached).unwrap()),
540 0 : );
541 0 : let result = query.load::<TenantShardPersistence>(conn).await?;
542 :
543 0 : Ok(result)
544 0 : })
545 0 : })
546 0 : .await
547 0 : }
548 :
549 : /// When restoring a previously detached tenant into memory, load it from the database
550 0 : pub(crate) async fn load_tenant(
551 0 : &self,
552 0 : filter_tenant_id: TenantId,
553 0 : ) -> DatabaseResult<Vec<TenantShardPersistence>> {
554 : use crate::schema::tenant_shards::dsl::*;
555 0 : self.with_measured_conn(DatabaseOperation::LoadTenant, move |conn| {
556 0 : Box::pin(async move {
557 0 : let query = tenant_shards.filter(tenant_id.eq(filter_tenant_id.to_string()));
558 0 : let result = query.load::<TenantShardPersistence>(conn).await?;
559 :
560 0 : Ok(result)
561 0 : })
562 0 : })
563 0 : .await
564 0 : }
565 :
566 : /// Tenants must be persisted before we schedule them for the first time. This enables us
567 : /// to correctly retain generation monotonicity, and the externally provided placement policy & config.
568 0 : pub(crate) async fn insert_tenant_shards(
569 0 : &self,
570 0 : shards: Vec<TenantShardPersistence>,
571 0 : ) -> DatabaseResult<()> {
572 : use crate::schema::{metadata_health, tenant_shards};
573 :
574 0 : let now = chrono::Utc::now();
575 0 :
576 0 : let metadata_health_records = shards
577 0 : .iter()
578 0 : .map(|t| MetadataHealthPersistence {
579 0 : tenant_id: t.tenant_id.clone(),
580 0 : shard_number: t.shard_number,
581 0 : shard_count: t.shard_count,
582 0 : healthy: true,
583 0 : last_scrubbed_at: now,
584 0 : })
585 0 : .collect::<Vec<_>>();
586 0 :
587 0 : let shards = &shards;
588 0 : let metadata_health_records = &metadata_health_records;
589 0 : self.with_measured_conn(DatabaseOperation::InsertTenantShards, move |conn| {
590 0 : Box::pin(async move {
591 0 : diesel::insert_into(tenant_shards::table)
592 0 : .values(shards)
593 0 : .execute(conn)
594 0 : .await?;
595 :
596 0 : diesel::insert_into(metadata_health::table)
597 0 : .values(metadata_health_records)
598 0 : .execute(conn)
599 0 : .await?;
600 0 : Ok(())
601 0 : })
602 0 : })
603 0 : .await
604 0 : }
605 :
606 : /// Ordering: call this _after_ deleting the tenant on pageservers, but _before_ dropping state for
607 : /// the tenant from memory on this server.
608 0 : pub(crate) async fn delete_tenant(&self, del_tenant_id: TenantId) -> DatabaseResult<()> {
609 : use crate::schema::tenant_shards::dsl::*;
610 0 : self.with_measured_conn(DatabaseOperation::DeleteTenant, move |conn| {
611 0 : Box::pin(async move {
612 0 : // `metadata_health` status (if exists) is also deleted based on the cascade behavior.
613 0 : diesel::delete(tenant_shards)
614 0 : .filter(tenant_id.eq(del_tenant_id.to_string()))
615 0 : .execute(conn)
616 0 : .await?;
617 0 : Ok(())
618 0 : })
619 0 : })
620 0 : .await
621 0 : }
622 :
623 : /// When a tenant invokes the /re-attach API, this function is responsible for doing an efficient
624 : /// batched increment of the generations of all tenants whose generation_pageserver is equal to
625 : /// the node that called /re-attach.
626 : #[tracing::instrument(skip_all, fields(node_id))]
627 : pub(crate) async fn re_attach(
628 : &self,
629 : input_node_id: NodeId,
630 : ) -> DatabaseResult<HashMap<TenantShardId, Generation>> {
631 : use crate::schema::nodes::dsl::{scheduling_policy, *};
632 : use crate::schema::tenant_shards::dsl::*;
633 : let updated = self
634 0 : .with_measured_conn(DatabaseOperation::ReAttach, move |conn| {
635 0 : Box::pin(async move {
636 : // Check if the node is not marked as deleted
637 0 : let deleted_node: i64 = nodes
638 0 : .filter(node_id.eq(input_node_id.0 as i64))
639 0 : .filter(lifecycle.eq(String::from(NodeLifecycle::Deleted)))
640 0 : .count()
641 0 : .get_result(conn)
642 0 : .await?;
643 0 : if deleted_node > 0 {
644 0 : return Err(DatabaseError::Logical(format!(
645 0 : "Node {} is marked as deleted, re-attach is not allowed",
646 0 : input_node_id
647 0 : )));
648 0 : }
649 :
650 0 : let rows_updated = diesel::update(tenant_shards)
651 0 : .filter(generation_pageserver.eq(input_node_id.0 as i64))
652 0 : .set(generation.eq(generation + 1))
653 0 : .execute(conn)
654 0 : .await?;
655 :
656 0 : tracing::info!("Incremented {} tenants' generations", rows_updated);
657 :
658 : // TODO: UPDATE+SELECT in one query
659 :
660 0 : let updated = tenant_shards
661 0 : .filter(generation_pageserver.eq(input_node_id.0 as i64))
662 0 : .select(TenantShardPersistence::as_select())
663 0 : .load(conn)
664 0 : .await?;
665 :
666 : // If the node went through a drain and restart phase before re-attaching,
667 : // then reset it's node scheduling policy to active.
668 0 : diesel::update(nodes)
669 0 : .filter(node_id.eq(input_node_id.0 as i64))
670 0 : .filter(
671 0 : scheduling_policy
672 0 : .eq(String::from(NodeSchedulingPolicy::PauseForRestart))
673 0 : .or(scheduling_policy
674 0 : .eq(String::from(NodeSchedulingPolicy::Draining)))
675 0 : .or(scheduling_policy
676 0 : .eq(String::from(NodeSchedulingPolicy::Filling))),
677 0 : )
678 0 : .set(scheduling_policy.eq(String::from(NodeSchedulingPolicy::Active)))
679 0 : .execute(conn)
680 0 : .await?;
681 :
682 0 : Ok(updated)
683 0 : })
684 0 : })
685 : .await?;
686 :
687 : let mut result = HashMap::new();
688 : for tsp in updated {
689 : let tenant_shard_id = TenantShardId {
690 : tenant_id: TenantId::from_str(tsp.tenant_id.as_str())
691 0 : .map_err(|e| DatabaseError::Logical(format!("Malformed tenant id: {e}")))?,
692 : shard_number: ShardNumber(tsp.shard_number as u8),
693 : shard_count: ShardCount::new(tsp.shard_count as u8),
694 : };
695 :
696 : let Some(g) = tsp.generation else {
697 : // If the generation_pageserver column was non-NULL, then the generation column should also be non-NULL:
698 : // we only set generation_pageserver when setting generation.
699 : return Err(DatabaseError::Logical(
700 : "Generation should always be set after incrementing".to_string(),
701 : ));
702 : };
703 : result.insert(tenant_shard_id, Generation::new(g as u32));
704 : }
705 :
706 : Ok(result)
707 : }
708 :
709 : /// Reconciler calls this immediately before attaching to a new pageserver, to acquire a unique, monotonically
710 : /// advancing generation number. We also store the NodeId for which the generation was issued, so that in
711 : /// [`Self::re_attach`] we can do a bulk UPDATE on the generations for that node.
712 0 : pub(crate) async fn increment_generation(
713 0 : &self,
714 0 : tenant_shard_id: TenantShardId,
715 0 : node_id: NodeId,
716 0 : ) -> anyhow::Result<Generation> {
717 : use crate::schema::tenant_shards::dsl::*;
718 0 : let generation_value = self
719 0 : .with_measured_conn(DatabaseOperation::IncrementGeneration, move |conn| {
720 0 : Box::pin(async move {
721 0 : let generation_value: Option<i32> = diesel::update(tenant_shards)
722 0 : .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
723 0 : .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
724 0 : .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
725 0 : .set((
726 0 : generation.eq(generation + 1),
727 0 : generation_pageserver.eq(node_id.0 as i64),
728 0 : ))
729 0 : .returning(generation)
730 0 : .get_result(conn)
731 0 : .await?;
732 :
733 0 : Ok(generation_value)
734 0 : })
735 0 : })
736 0 : .await?;
737 :
738 : // Generation is always non-null in the rseult: if the generation column had been NULL, then we
739 : // should have experienced an SQL Confilict error while executing a query that tries to increment it.
740 0 : debug_assert!(generation_value.is_some());
741 0 : let Some(g) = generation_value else {
742 0 : return Err(DatabaseError::Logical(
743 0 : "Generation should always be set after incrementing".to_string(),
744 0 : )
745 0 : .into());
746 : };
747 :
748 0 : Ok(Generation::new(g as u32))
749 0 : }
750 :
751 : /// When we want to call out to the running shards for a tenant, e.g. during timeline CRUD operations,
752 : /// we need to know where the shard is attached, _and_ the generation, so that we can re-check the generation
753 : /// afterwards to confirm that our timeline CRUD operation is truly persistent (it must have happened in the
754 : /// latest generation)
755 : ///
756 : /// If the tenant doesn't exist, an empty vector is returned.
757 : ///
758 : /// Output is sorted by shard number
759 0 : pub(crate) async fn tenant_generations(
760 0 : &self,
761 0 : filter_tenant_id: TenantId,
762 0 : ) -> Result<Vec<ShardGenerationState>, DatabaseError> {
763 : use crate::schema::tenant_shards::dsl::*;
764 0 : let rows = self
765 0 : .with_measured_conn(DatabaseOperation::TenantGenerations, move |conn| {
766 0 : Box::pin(async move {
767 0 : let result = tenant_shards
768 0 : .filter(tenant_id.eq(filter_tenant_id.to_string()))
769 0 : .select(TenantShardPersistence::as_select())
770 0 : .order(shard_number)
771 0 : .load(conn)
772 0 : .await?;
773 0 : Ok(result)
774 0 : })
775 0 : })
776 0 : .await?;
777 :
778 0 : Ok(rows
779 0 : .into_iter()
780 0 : .map(|p| ShardGenerationState {
781 0 : tenant_shard_id: p
782 0 : .get_tenant_shard_id()
783 0 : .expect("Corrupt tenant shard id in database"),
784 0 : generation: p.generation.map(|g| Generation::new(g as u32)),
785 0 : generation_pageserver: p.generation_pageserver.map(|n| NodeId(n as u64)),
786 0 : })
787 0 : .collect())
788 0 : }
789 :
790 : /// Read the generation number of specific tenant shards
791 : ///
792 : /// Output is unsorted. Output may not include values for all inputs, if they are missing in the database.
793 0 : pub(crate) async fn shard_generations(
794 0 : &self,
795 0 : mut tenant_shard_ids: impl Iterator<Item = &TenantShardId>,
796 0 : ) -> Result<Vec<(TenantShardId, Option<Generation>)>, DatabaseError> {
797 0 : let mut rows = Vec::with_capacity(tenant_shard_ids.size_hint().0);
798 :
799 : // We will chunk our input to avoid composing arbitrarily long `IN` clauses. Typically we are
800 : // called with a single digit number of IDs, but in principle we could be called with tens
801 : // of thousands (all the shards on one pageserver) from the generation validation API.
802 0 : loop {
803 0 : // A modest hardcoded chunk size to handle typical cases in a single query but never generate particularly
804 0 : // large query strings.
805 0 : let chunk_ids = tenant_shard_ids.by_ref().take(32);
806 0 :
807 0 : // Compose a comma separated list of tuples for matching on (tenant_id, shard_number, shard_count)
808 0 : let in_clause = chunk_ids
809 0 : .map(|tsid| {
810 0 : format!(
811 0 : "('{}', {}, {})",
812 0 : tsid.tenant_id, tsid.shard_number.0, tsid.shard_count.0
813 0 : )
814 0 : })
815 0 : .join(",");
816 0 :
817 0 : // We are done when our iterator gives us nothing to filter on
818 0 : if in_clause.is_empty() {
819 0 : break;
820 0 : }
821 0 :
822 0 : let in_clause = &in_clause;
823 0 : let chunk_rows = self
824 0 : .with_measured_conn(DatabaseOperation::ShardGenerations, move |conn| {
825 0 : Box::pin(async move {
826 : // diesel doesn't support multi-column IN queries, so we compose raw SQL. No escaping is required because
827 : // the inputs are strongly typed and cannot carry any user-supplied raw string content.
828 0 : let result : Vec<TenantShardPersistence> = diesel::sql_query(
829 0 : format!("SELECT * from tenant_shards where (tenant_id, shard_number, shard_count) in ({in_clause});").as_str()
830 0 : ).load(conn).await?;
831 :
832 0 : Ok(result)
833 0 : })
834 0 : })
835 0 : .await?;
836 0 : rows.extend(chunk_rows.into_iter())
837 : }
838 :
839 0 : Ok(rows
840 0 : .into_iter()
841 0 : .map(|tsp| {
842 0 : (
843 0 : tsp.get_tenant_shard_id()
844 0 : .expect("Bad tenant ID in database"),
845 0 : tsp.generation.map(|g| Generation::new(g as u32)),
846 0 : )
847 0 : })
848 0 : .collect())
849 0 : }
850 :
851 : #[allow(non_local_definitions)]
852 : /// For use when updating a persistent property of a tenant, such as its config or placement_policy.
853 : ///
854 : /// Do not use this for settting generation, unless in the special onboarding code path (/location_config)
855 : /// API: use [`Self::increment_generation`] instead. Setting the generation via this route is a one-time thing
856 : /// that we only do the first time a tenant is set to an attached policy via /location_config.
857 0 : pub(crate) async fn update_tenant_shard(
858 0 : &self,
859 0 : tenant: TenantFilter,
860 0 : input_placement_policy: Option<PlacementPolicy>,
861 0 : input_config: Option<TenantConfig>,
862 0 : input_generation: Option<Generation>,
863 0 : input_scheduling_policy: Option<ShardSchedulingPolicy>,
864 0 : ) -> DatabaseResult<()> {
865 : use crate::schema::tenant_shards::dsl::*;
866 :
867 0 : let tenant = &tenant;
868 0 : let input_placement_policy = &input_placement_policy;
869 0 : let input_config = &input_config;
870 0 : let input_generation = &input_generation;
871 0 : let input_scheduling_policy = &input_scheduling_policy;
872 0 : self.with_measured_conn(DatabaseOperation::UpdateTenantShard, move |conn| {
873 0 : Box::pin(async move {
874 0 : let query = match tenant {
875 0 : TenantFilter::Shard(tenant_shard_id) => diesel::update(tenant_shards)
876 0 : .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
877 0 : .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
878 0 : .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
879 0 : .into_boxed(),
880 0 : TenantFilter::Tenant(input_tenant_id) => diesel::update(tenant_shards)
881 0 : .filter(tenant_id.eq(input_tenant_id.to_string()))
882 0 : .into_boxed(),
883 : };
884 :
885 : // Clear generation_pageserver if we are moving into a state where we won't have
886 : // any attached pageservers.
887 0 : let input_generation_pageserver = match input_placement_policy {
888 0 : None | Some(PlacementPolicy::Attached(_)) => None,
889 0 : Some(PlacementPolicy::Detached | PlacementPolicy::Secondary) => Some(None),
890 : };
891 :
892 0 : #[derive(AsChangeset)]
893 : #[diesel(table_name = crate::schema::tenant_shards)]
894 : struct ShardUpdate {
895 : generation: Option<i32>,
896 : placement_policy: Option<String>,
897 : config: Option<String>,
898 : scheduling_policy: Option<String>,
899 : generation_pageserver: Option<Option<i64>>,
900 : }
901 :
902 0 : let update = ShardUpdate {
903 0 : generation: input_generation.map(|g| g.into().unwrap() as i32),
904 0 : placement_policy: input_placement_policy
905 0 : .as_ref()
906 0 : .map(|p| serde_json::to_string(&p).unwrap()),
907 0 : config: input_config
908 0 : .as_ref()
909 0 : .map(|c| serde_json::to_string(&c).unwrap()),
910 0 : scheduling_policy: input_scheduling_policy
911 0 : .map(|p| serde_json::to_string(&p).unwrap()),
912 0 : generation_pageserver: input_generation_pageserver,
913 0 : };
914 0 :
915 0 : query.set(update).execute(conn).await?;
916 :
917 0 : Ok(())
918 0 : })
919 0 : })
920 0 : .await?;
921 :
922 0 : Ok(())
923 0 : }
924 :
925 : /// Note that passing None for a shard clears the preferred AZ (rather than leaving it unmodified)
926 0 : pub(crate) async fn set_tenant_shard_preferred_azs(
927 0 : &self,
928 0 : preferred_azs: Vec<(TenantShardId, Option<AvailabilityZone>)>,
929 0 : ) -> DatabaseResult<Vec<(TenantShardId, Option<AvailabilityZone>)>> {
930 : use crate::schema::tenant_shards::dsl::*;
931 :
932 0 : let preferred_azs = preferred_azs.as_slice();
933 0 : self.with_measured_conn(DatabaseOperation::SetPreferredAzs, move |conn| {
934 0 : Box::pin(async move {
935 0 : let mut shards_updated = Vec::default();
936 :
937 0 : for (tenant_shard_id, preferred_az) in preferred_azs.iter() {
938 0 : let updated = diesel::update(tenant_shards)
939 0 : .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
940 0 : .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
941 0 : .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
942 0 : .set(preferred_az_id.eq(preferred_az.as_ref().map(|az| az.0.clone())))
943 0 : .execute(conn)
944 0 : .await?;
945 :
946 0 : if updated == 1 {
947 0 : shards_updated.push((*tenant_shard_id, preferred_az.clone()));
948 0 : }
949 : }
950 :
951 0 : Ok(shards_updated)
952 0 : })
953 0 : })
954 0 : .await
955 0 : }
956 :
957 0 : pub(crate) async fn detach(&self, tenant_shard_id: TenantShardId) -> anyhow::Result<()> {
958 : use crate::schema::tenant_shards::dsl::*;
959 0 : self.with_measured_conn(DatabaseOperation::Detach, move |conn| {
960 0 : Box::pin(async move {
961 0 : let updated = diesel::update(tenant_shards)
962 0 : .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
963 0 : .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
964 0 : .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
965 0 : .set((
966 0 : generation_pageserver.eq(Option::<i64>::None),
967 0 : placement_policy
968 0 : .eq(serde_json::to_string(&PlacementPolicy::Detached).unwrap()),
969 0 : ))
970 0 : .execute(conn)
971 0 : .await?;
972 :
973 0 : Ok(updated)
974 0 : })
975 0 : })
976 0 : .await?;
977 :
978 0 : Ok(())
979 0 : }
980 :
981 : // When we start shard splitting, we must durably mark the tenant so that
982 : // on restart, we know that we must go through recovery.
983 : //
984 : // We create the child shards here, so that they will be available for increment_generation calls
985 : // if some pageserver holding a child shard needs to restart before the overall tenant split is complete.
986 0 : pub(crate) async fn begin_shard_split(
987 0 : &self,
988 0 : old_shard_count: ShardCount,
989 0 : split_tenant_id: TenantId,
990 0 : parent_to_children: Vec<(TenantShardId, Vec<TenantShardPersistence>)>,
991 0 : ) -> DatabaseResult<()> {
992 : use crate::schema::tenant_shards::dsl::*;
993 0 : let parent_to_children = parent_to_children.as_slice();
994 0 : self.with_measured_conn(DatabaseOperation::BeginShardSplit, move |conn| {
995 0 : Box::pin(async move {
996 : // Mark parent shards as splitting
997 :
998 0 : let updated = diesel::update(tenant_shards)
999 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
1000 0 : .filter(shard_count.eq(old_shard_count.literal() as i32))
1001 0 : .set((splitting.eq(1),))
1002 0 : .execute(conn).await?;
1003 0 : if u8::try_from(updated)
1004 0 : .map_err(|_| DatabaseError::Logical(
1005 0 : format!("Overflow existing shard count {} while splitting", updated))
1006 0 : )? != old_shard_count.count() {
1007 : // Perhaps a deletion or another split raced with this attempt to split, mutating
1008 : // the parent shards that we intend to split. In this case the split request should fail.
1009 0 : return Err(DatabaseError::Logical(
1010 0 : format!("Unexpected existing shard count {updated} when preparing tenant for split (expected {})", old_shard_count.count())
1011 0 : ));
1012 0 : }
1013 0 :
1014 0 : // FIXME: spurious clone to sidestep closure move rules
1015 0 : let parent_to_children = parent_to_children.to_vec();
1016 :
1017 : // Insert child shards
1018 0 : for (parent_shard_id, children) in parent_to_children {
1019 0 : let mut parent = crate::schema::tenant_shards::table
1020 0 : .filter(tenant_id.eq(parent_shard_id.tenant_id.to_string()))
1021 0 : .filter(shard_number.eq(parent_shard_id.shard_number.0 as i32))
1022 0 : .filter(shard_count.eq(parent_shard_id.shard_count.literal() as i32))
1023 0 : .load::<TenantShardPersistence>(conn).await?;
1024 0 : let parent = if parent.len() != 1 {
1025 0 : return Err(DatabaseError::Logical(format!(
1026 0 : "Parent shard {parent_shard_id} not found"
1027 0 : )));
1028 : } else {
1029 0 : parent.pop().unwrap()
1030 : };
1031 0 : for mut shard in children {
1032 : // Carry the parent's generation into the child
1033 0 : shard.generation = parent.generation;
1034 0 :
1035 0 : debug_assert!(shard.splitting == SplitState::Splitting);
1036 0 : diesel::insert_into(tenant_shards)
1037 0 : .values(shard)
1038 0 : .execute(conn).await?;
1039 : }
1040 : }
1041 :
1042 0 : Ok(())
1043 0 : })
1044 0 : })
1045 0 : .await
1046 0 : }
1047 :
1048 : // When we finish shard splitting, we must atomically clean up the old shards
1049 : // and insert the new shards, and clear the splitting marker.
1050 0 : pub(crate) async fn complete_shard_split(
1051 0 : &self,
1052 0 : split_tenant_id: TenantId,
1053 0 : old_shard_count: ShardCount,
1054 0 : new_shard_count: ShardCount,
1055 0 : ) -> DatabaseResult<()> {
1056 : use crate::schema::tenant_shards::dsl::*;
1057 0 : self.with_measured_conn(DatabaseOperation::CompleteShardSplit, move |conn| {
1058 0 : Box::pin(async move {
1059 0 : // Sanity: child shards must still exist, as we're deleting parent shards
1060 0 : let child_shards_query = tenant_shards
1061 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
1062 0 : .filter(shard_count.eq(new_shard_count.literal() as i32));
1063 0 : let child_shards = child_shards_query
1064 0 : .load::<TenantShardPersistence>(conn)
1065 0 : .await?;
1066 0 : if child_shards.len() != new_shard_count.count() as usize {
1067 0 : return Err(DatabaseError::Logical(format!(
1068 0 : "Unexpected child shard count {} while completing split to \
1069 0 : count {new_shard_count:?} on tenant {split_tenant_id}",
1070 0 : child_shards.len()
1071 0 : )));
1072 0 : }
1073 0 :
1074 0 : // Drop parent shards
1075 0 : diesel::delete(tenant_shards)
1076 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
1077 0 : .filter(shard_count.eq(old_shard_count.literal() as i32))
1078 0 : .execute(conn)
1079 0 : .await?;
1080 :
1081 : // Clear sharding flag
1082 0 : let updated = diesel::update(tenant_shards)
1083 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
1084 0 : .filter(shard_count.eq(new_shard_count.literal() as i32))
1085 0 : .set((splitting.eq(0),))
1086 0 : .execute(conn)
1087 0 : .await?;
1088 0 : assert!(updated == new_shard_count.count() as usize);
1089 :
1090 0 : Ok(())
1091 0 : })
1092 0 : })
1093 0 : .await
1094 0 : }
1095 :
1096 : /// Used when the remote part of a shard split failed: we will revert the database state to have only
1097 : /// the parent shards, with SplitState::Idle.
1098 0 : pub(crate) async fn abort_shard_split(
1099 0 : &self,
1100 0 : split_tenant_id: TenantId,
1101 0 : new_shard_count: ShardCount,
1102 0 : ) -> DatabaseResult<AbortShardSplitStatus> {
1103 : use crate::schema::tenant_shards::dsl::*;
1104 0 : self.with_measured_conn(DatabaseOperation::AbortShardSplit, move |conn| {
1105 0 : Box::pin(async move {
1106 : // Clear the splitting state on parent shards
1107 0 : let updated = diesel::update(tenant_shards)
1108 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
1109 0 : .filter(shard_count.ne(new_shard_count.literal() as i32))
1110 0 : .set((splitting.eq(0),))
1111 0 : .execute(conn)
1112 0 : .await?;
1113 :
1114 : // Parent shards are already gone: we cannot abort.
1115 0 : if updated == 0 {
1116 0 : return Ok(AbortShardSplitStatus::Complete);
1117 0 : }
1118 0 :
1119 0 : // Sanity check: if parent shards were present, their cardinality should
1120 0 : // be less than the number of child shards.
1121 0 : if updated >= new_shard_count.count() as usize {
1122 0 : return Err(DatabaseError::Logical(format!(
1123 0 : "Unexpected parent shard count {updated} while aborting split to \
1124 0 : count {new_shard_count:?} on tenant {split_tenant_id}"
1125 0 : )));
1126 0 : }
1127 0 :
1128 0 : // Erase child shards
1129 0 : diesel::delete(tenant_shards)
1130 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
1131 0 : .filter(shard_count.eq(new_shard_count.literal() as i32))
1132 0 : .execute(conn)
1133 0 : .await?;
1134 :
1135 0 : Ok(AbortShardSplitStatus::Aborted)
1136 0 : })
1137 0 : })
1138 0 : .await
1139 0 : }
1140 :
1141 : /// Stores all the latest metadata health updates durably. Updates existing entry on conflict.
1142 : ///
1143 : /// **Correctness:** `metadata_health_updates` should all belong the tenant shards managed by the storage controller.
1144 : #[allow(dead_code)]
1145 0 : pub(crate) async fn update_metadata_health_records(
1146 0 : &self,
1147 0 : healthy_records: Vec<MetadataHealthPersistence>,
1148 0 : unhealthy_records: Vec<MetadataHealthPersistence>,
1149 0 : now: chrono::DateTime<chrono::Utc>,
1150 0 : ) -> DatabaseResult<()> {
1151 : use crate::schema::metadata_health::dsl::*;
1152 :
1153 0 : let healthy_records = healthy_records.as_slice();
1154 0 : let unhealthy_records = unhealthy_records.as_slice();
1155 0 : self.with_measured_conn(DatabaseOperation::UpdateMetadataHealth, move |conn| {
1156 0 : Box::pin(async move {
1157 0 : diesel::insert_into(metadata_health)
1158 0 : .values(healthy_records)
1159 0 : .on_conflict((tenant_id, shard_number, shard_count))
1160 0 : .do_update()
1161 0 : .set((healthy.eq(true), last_scrubbed_at.eq(now)))
1162 0 : .execute(conn)
1163 0 : .await?;
1164 :
1165 0 : diesel::insert_into(metadata_health)
1166 0 : .values(unhealthy_records)
1167 0 : .on_conflict((tenant_id, shard_number, shard_count))
1168 0 : .do_update()
1169 0 : .set((healthy.eq(false), last_scrubbed_at.eq(now)))
1170 0 : .execute(conn)
1171 0 : .await?;
1172 0 : Ok(())
1173 0 : })
1174 0 : })
1175 0 : .await
1176 0 : }
1177 :
1178 : /// Lists all the metadata health records.
1179 : #[allow(dead_code)]
1180 0 : pub(crate) async fn list_metadata_health_records(
1181 0 : &self,
1182 0 : ) -> DatabaseResult<Vec<MetadataHealthPersistence>> {
1183 0 : self.with_measured_conn(DatabaseOperation::ListMetadataHealth, move |conn| {
1184 0 : Box::pin(async {
1185 0 : Ok(crate::schema::metadata_health::table
1186 0 : .load::<MetadataHealthPersistence>(conn)
1187 0 : .await?)
1188 0 : })
1189 0 : })
1190 0 : .await
1191 0 : }
1192 :
1193 : /// Lists all the metadata health records that is unhealthy.
1194 : #[allow(dead_code)]
1195 0 : pub(crate) async fn list_unhealthy_metadata_health_records(
1196 0 : &self,
1197 0 : ) -> DatabaseResult<Vec<MetadataHealthPersistence>> {
1198 : use crate::schema::metadata_health::dsl::*;
1199 0 : self.with_measured_conn(
1200 0 : DatabaseOperation::ListMetadataHealthUnhealthy,
1201 0 : move |conn| {
1202 0 : Box::pin(async {
1203 0 : DatabaseResult::Ok(
1204 0 : crate::schema::metadata_health::table
1205 0 : .filter(healthy.eq(false))
1206 0 : .load::<MetadataHealthPersistence>(conn)
1207 0 : .await?,
1208 : )
1209 0 : })
1210 0 : },
1211 0 : )
1212 0 : .await
1213 0 : }
1214 :
1215 : /// Lists all the metadata health records that have not been updated since an `earlier` time.
1216 : #[allow(dead_code)]
1217 0 : pub(crate) async fn list_outdated_metadata_health_records(
1218 0 : &self,
1219 0 : earlier: chrono::DateTime<chrono::Utc>,
1220 0 : ) -> DatabaseResult<Vec<MetadataHealthPersistence>> {
1221 : use crate::schema::metadata_health::dsl::*;
1222 :
1223 0 : self.with_measured_conn(DatabaseOperation::ListMetadataHealthOutdated, move |conn| {
1224 0 : Box::pin(async move {
1225 0 : let query = metadata_health.filter(last_scrubbed_at.lt(earlier));
1226 0 : let res = query.load::<MetadataHealthPersistence>(conn).await?;
1227 :
1228 0 : Ok(res)
1229 0 : })
1230 0 : })
1231 0 : .await
1232 0 : }
1233 :
1234 : /// Get the current entry from the `leader` table if one exists.
1235 : /// It is an error for the table to contain more than one entry.
1236 0 : pub(crate) async fn get_leader(&self) -> DatabaseResult<Option<ControllerPersistence>> {
1237 0 : let mut leader: Vec<ControllerPersistence> = self
1238 0 : .with_measured_conn(DatabaseOperation::GetLeader, move |conn| {
1239 0 : Box::pin(async move {
1240 0 : Ok(crate::schema::controllers::table
1241 0 : .load::<ControllerPersistence>(conn)
1242 0 : .await?)
1243 0 : })
1244 0 : })
1245 0 : .await?;
1246 :
1247 0 : if leader.len() > 1 {
1248 0 : return Err(DatabaseError::Logical(format!(
1249 0 : "More than one entry present in the leader table: {leader:?}"
1250 0 : )));
1251 0 : }
1252 0 :
1253 0 : Ok(leader.pop())
1254 0 : }
1255 :
1256 : /// Update the new leader with compare-exchange semantics. If `prev` does not
1257 : /// match the current leader entry, then the update is treated as a failure.
1258 : /// When `prev` is not specified, the update is forced.
1259 0 : pub(crate) async fn update_leader(
1260 0 : &self,
1261 0 : prev: Option<ControllerPersistence>,
1262 0 : new: ControllerPersistence,
1263 0 : ) -> DatabaseResult<()> {
1264 : use crate::schema::controllers::dsl::*;
1265 :
1266 0 : let updated = self
1267 0 : .with_measured_conn(DatabaseOperation::UpdateLeader, move |conn| {
1268 0 : let prev = prev.clone();
1269 0 : let new = new.clone();
1270 0 : Box::pin(async move {
1271 0 : let updated = match &prev {
1272 0 : Some(prev) => {
1273 0 : diesel::update(controllers)
1274 0 : .filter(address.eq(prev.address.clone()))
1275 0 : .filter(started_at.eq(prev.started_at))
1276 0 : .set((
1277 0 : address.eq(new.address.clone()),
1278 0 : started_at.eq(new.started_at),
1279 0 : ))
1280 0 : .execute(conn)
1281 0 : .await?
1282 : }
1283 : None => {
1284 0 : diesel::insert_into(controllers)
1285 0 : .values(new.clone())
1286 0 : .execute(conn)
1287 0 : .await?
1288 : }
1289 : };
1290 :
1291 0 : Ok(updated)
1292 0 : })
1293 0 : })
1294 0 : .await?;
1295 :
1296 0 : if updated == 0 {
1297 0 : return Err(DatabaseError::Logical(
1298 0 : "Leader table update failed".to_string(),
1299 0 : ));
1300 0 : }
1301 0 :
1302 0 : Ok(())
1303 0 : }
1304 :
1305 : /// At startup, populate the list of nodes which our shards may be placed on
1306 0 : pub(crate) async fn list_safekeepers(&self) -> DatabaseResult<Vec<SafekeeperPersistence>> {
1307 0 : let safekeepers: Vec<SafekeeperPersistence> = self
1308 0 : .with_measured_conn(DatabaseOperation::ListNodes, move |conn| {
1309 0 : Box::pin(async move {
1310 0 : Ok(crate::schema::safekeepers::table
1311 0 : .load::<SafekeeperPersistence>(conn)
1312 0 : .await?)
1313 0 : })
1314 0 : })
1315 0 : .await?;
1316 :
1317 0 : tracing::info!("list_safekeepers: loaded {} nodes", safekeepers.len());
1318 :
1319 0 : Ok(safekeepers)
1320 0 : }
1321 :
1322 0 : pub(crate) async fn safekeeper_upsert(
1323 0 : &self,
1324 0 : record: SafekeeperUpsert,
1325 0 : ) -> Result<(), DatabaseError> {
1326 : use crate::schema::safekeepers::dsl::*;
1327 :
1328 0 : self.with_conn(move |conn| {
1329 0 : let record = record.clone();
1330 0 : Box::pin(async move {
1331 0 : let bind = record
1332 0 : .as_insert_or_update()
1333 0 : .map_err(|e| DatabaseError::Logical(format!("{e}")))?;
1334 :
1335 0 : let inserted_updated = diesel::insert_into(safekeepers)
1336 0 : .values(&bind)
1337 0 : .on_conflict(id)
1338 0 : .do_update()
1339 0 : .set(&bind)
1340 0 : .execute(conn)
1341 0 : .await?;
1342 :
1343 0 : if inserted_updated != 1 {
1344 0 : return Err(DatabaseError::Logical(format!(
1345 0 : "unexpected number of rows ({})",
1346 0 : inserted_updated
1347 0 : )));
1348 0 : }
1349 0 :
1350 0 : Ok(())
1351 0 : })
1352 0 : })
1353 0 : .await
1354 0 : }
1355 :
1356 0 : pub(crate) async fn set_safekeeper_scheduling_policy(
1357 0 : &self,
1358 0 : id_: i64,
1359 0 : scheduling_policy_: SkSchedulingPolicy,
1360 0 : ) -> Result<(), DatabaseError> {
1361 : use crate::schema::safekeepers::dsl::*;
1362 :
1363 0 : self.with_conn(move |conn| {
1364 0 : Box::pin(async move {
1365 0 : #[derive(Insertable, AsChangeset)]
1366 : #[diesel(table_name = crate::schema::safekeepers)]
1367 : struct UpdateSkSchedulingPolicy<'a> {
1368 : id: i64,
1369 : scheduling_policy: &'a str,
1370 : }
1371 0 : let scheduling_policy_ = String::from(scheduling_policy_);
1372 :
1373 0 : let rows_affected = diesel::update(safekeepers.filter(id.eq(id_)))
1374 0 : .set(scheduling_policy.eq(scheduling_policy_))
1375 0 : .execute(conn)
1376 0 : .await?;
1377 :
1378 0 : if rows_affected != 1 {
1379 0 : return Err(DatabaseError::Logical(format!(
1380 0 : "unexpected number of rows ({rows_affected})",
1381 0 : )));
1382 0 : }
1383 0 :
1384 0 : Ok(())
1385 0 : })
1386 0 : })
1387 0 : .await
1388 0 : }
1389 :
1390 : /// Persist timeline. Returns if the timeline was newly inserted. If it wasn't, we haven't done any writes.
1391 0 : pub(crate) async fn insert_timeline(&self, entry: TimelinePersistence) -> DatabaseResult<bool> {
1392 : use crate::schema::timelines;
1393 :
1394 0 : let entry = &entry;
1395 0 : self.with_measured_conn(DatabaseOperation::InsertTimeline, move |conn| {
1396 0 : Box::pin(async move {
1397 0 : let inserted_updated = diesel::insert_into(timelines::table)
1398 0 : .values(entry)
1399 0 : .on_conflict((timelines::tenant_id, timelines::timeline_id))
1400 0 : .do_nothing()
1401 0 : .execute(conn)
1402 0 : .await?;
1403 :
1404 0 : match inserted_updated {
1405 0 : 0 => Ok(false),
1406 0 : 1 => Ok(true),
1407 0 : _ => Err(DatabaseError::Logical(format!(
1408 0 : "unexpected number of rows ({})",
1409 0 : inserted_updated
1410 0 : ))),
1411 : }
1412 0 : })
1413 0 : })
1414 0 : .await
1415 0 : }
1416 :
1417 : /// Load timeline from db. Returns `None` if not present.
1418 0 : pub(crate) async fn get_timeline(
1419 0 : &self,
1420 0 : tenant_id: TenantId,
1421 0 : timeline_id: TimelineId,
1422 0 : ) -> DatabaseResult<Option<TimelinePersistence>> {
1423 : use crate::schema::timelines::dsl;
1424 :
1425 0 : let tenant_id = &tenant_id;
1426 0 : let timeline_id = &timeline_id;
1427 0 : let timeline_from_db = self
1428 0 : .with_measured_conn(DatabaseOperation::GetTimeline, move |conn| {
1429 0 : Box::pin(async move {
1430 0 : let mut from_db: Vec<TimelineFromDb> = dsl::timelines
1431 0 : .filter(
1432 0 : dsl::tenant_id
1433 0 : .eq(&tenant_id.to_string())
1434 0 : .and(dsl::timeline_id.eq(&timeline_id.to_string())),
1435 0 : )
1436 0 : .load(conn)
1437 0 : .await?;
1438 0 : if from_db.is_empty() {
1439 0 : return Ok(None);
1440 0 : }
1441 0 : if from_db.len() != 1 {
1442 0 : return Err(DatabaseError::Logical(format!(
1443 0 : "unexpected number of rows ({})",
1444 0 : from_db.len()
1445 0 : )));
1446 0 : }
1447 0 :
1448 0 : Ok(Some(from_db.pop().unwrap().into_persistence()))
1449 0 : })
1450 0 : })
1451 0 : .await?;
1452 :
1453 0 : Ok(timeline_from_db)
1454 0 : }
1455 :
1456 : /// Set `delete_at` for the given timeline
1457 0 : pub(crate) async fn timeline_set_deleted_at(
1458 0 : &self,
1459 0 : tenant_id: TenantId,
1460 0 : timeline_id: TimelineId,
1461 0 : ) -> DatabaseResult<()> {
1462 : use crate::schema::timelines;
1463 :
1464 0 : let deletion_time = chrono::Local::now().to_utc();
1465 0 : self.with_measured_conn(DatabaseOperation::InsertTimeline, move |conn| {
1466 0 : Box::pin(async move {
1467 0 : let updated = diesel::update(timelines::table)
1468 0 : .filter(timelines::tenant_id.eq(tenant_id.to_string()))
1469 0 : .filter(timelines::timeline_id.eq(timeline_id.to_string()))
1470 0 : .set(timelines::deleted_at.eq(Some(deletion_time)))
1471 0 : .execute(conn)
1472 0 : .await?;
1473 :
1474 0 : match updated {
1475 0 : 0 => Ok(()),
1476 0 : 1 => Ok(()),
1477 0 : _ => Err(DatabaseError::Logical(format!(
1478 0 : "unexpected number of rows ({})",
1479 0 : updated
1480 0 : ))),
1481 : }
1482 0 : })
1483 0 : })
1484 0 : .await
1485 0 : }
1486 :
1487 : /// Load timeline from db. Returns `None` if not present.
1488 : ///
1489 : /// Only works if `deleted_at` is set, so you should call [`Self::timeline_set_deleted_at`] before.
1490 0 : pub(crate) async fn delete_timeline(
1491 0 : &self,
1492 0 : tenant_id: TenantId,
1493 0 : timeline_id: TimelineId,
1494 0 : ) -> DatabaseResult<()> {
1495 : use crate::schema::timelines::dsl;
1496 :
1497 0 : let tenant_id = &tenant_id;
1498 0 : let timeline_id = &timeline_id;
1499 0 : self.with_measured_conn(DatabaseOperation::GetTimeline, move |conn| {
1500 0 : Box::pin(async move {
1501 0 : diesel::delete(dsl::timelines)
1502 0 : .filter(dsl::tenant_id.eq(&tenant_id.to_string()))
1503 0 : .filter(dsl::timeline_id.eq(&timeline_id.to_string()))
1504 0 : .filter(dsl::deleted_at.is_not_null())
1505 0 : .execute(conn)
1506 0 : .await?;
1507 0 : Ok(())
1508 0 : })
1509 0 : })
1510 0 : .await?;
1511 :
1512 0 : Ok(())
1513 0 : }
1514 :
1515 : /// Loads a list of all timelines from database.
1516 0 : pub(crate) async fn list_timelines_for_tenant(
1517 0 : &self,
1518 0 : tenant_id: TenantId,
1519 0 : ) -> DatabaseResult<Vec<TimelinePersistence>> {
1520 : use crate::schema::timelines::dsl;
1521 :
1522 0 : let tenant_id = &tenant_id;
1523 0 : let timelines = self
1524 0 : .with_measured_conn(DatabaseOperation::GetTimeline, move |conn| {
1525 0 : Box::pin(async move {
1526 0 : let timelines: Vec<TimelineFromDb> = dsl::timelines
1527 0 : .filter(dsl::tenant_id.eq(&tenant_id.to_string()))
1528 0 : .load(conn)
1529 0 : .await?;
1530 0 : Ok(timelines)
1531 0 : })
1532 0 : })
1533 0 : .await?;
1534 :
1535 0 : let timelines = timelines
1536 0 : .into_iter()
1537 0 : .map(TimelineFromDb::into_persistence)
1538 0 : .collect();
1539 0 : Ok(timelines)
1540 0 : }
1541 :
1542 : /// Persist pending op. Returns if it was newly inserted. If it wasn't, we haven't done any writes.
1543 0 : pub(crate) async fn insert_pending_op(
1544 0 : &self,
1545 0 : entry: TimelinePendingOpPersistence,
1546 0 : ) -> DatabaseResult<bool> {
1547 : use crate::schema::safekeeper_timeline_pending_ops as skpo;
1548 : // This overrides the `filter` fn used in other functions, so contain the mayhem via a function-local use
1549 : use diesel::query_dsl::methods::FilterDsl;
1550 :
1551 0 : let entry = &entry;
1552 0 : self.with_measured_conn(DatabaseOperation::InsertTimelineReconcile, move |conn| {
1553 0 : Box::pin(async move {
1554 : // For simplicity it makes sense to keep only the last operation
1555 : // per (tenant, timeline, sk) tuple: if we migrated a timeline
1556 : // from node and adding it back it is not necessary to remove
1557 : // data on it. Hence, generation is not part of primary key and
1558 : // we override any rows with lower generations here.
1559 0 : let inserted_updated = diesel::insert_into(skpo::table)
1560 0 : .values(entry)
1561 0 : .on_conflict((skpo::tenant_id, skpo::timeline_id, skpo::sk_id))
1562 0 : .do_update()
1563 0 : .set(entry)
1564 0 : .filter(skpo::generation.lt(entry.generation))
1565 0 : .execute(conn)
1566 0 : .await?;
1567 :
1568 0 : match inserted_updated {
1569 0 : 0 => Ok(false),
1570 0 : 1 => Ok(true),
1571 0 : _ => Err(DatabaseError::Logical(format!(
1572 0 : "unexpected number of rows ({})",
1573 0 : inserted_updated
1574 0 : ))),
1575 : }
1576 0 : })
1577 0 : })
1578 0 : .await
1579 0 : }
1580 : /// Remove persisted pending op.
1581 0 : pub(crate) async fn remove_pending_op(
1582 0 : &self,
1583 0 : tenant_id: TenantId,
1584 0 : timeline_id: Option<TimelineId>,
1585 0 : sk_id: NodeId,
1586 0 : generation: u32,
1587 0 : ) -> DatabaseResult<()> {
1588 : use crate::schema::safekeeper_timeline_pending_ops::dsl;
1589 :
1590 0 : let tenant_id = &tenant_id;
1591 0 : let timeline_id = &timeline_id;
1592 0 : self.with_measured_conn(DatabaseOperation::RemoveTimelineReconcile, move |conn| {
1593 0 : let timeline_id_str = timeline_id.map(|tid| tid.to_string()).unwrap_or_default();
1594 0 : Box::pin(async move {
1595 0 : diesel::delete(dsl::safekeeper_timeline_pending_ops)
1596 0 : .filter(dsl::tenant_id.eq(tenant_id.to_string()))
1597 0 : .filter(dsl::timeline_id.eq(timeline_id_str))
1598 0 : .filter(dsl::sk_id.eq(sk_id.0 as i64))
1599 0 : .filter(dsl::generation.eq(generation as i32))
1600 0 : .execute(conn)
1601 0 : .await?;
1602 0 : Ok(())
1603 0 : })
1604 0 : })
1605 0 : .await
1606 0 : }
1607 :
1608 : /// Load pending operations from db, joined together with timeline data.
1609 0 : pub(crate) async fn list_pending_ops_with_timelines(
1610 0 : &self,
1611 0 : ) -> DatabaseResult<Vec<(TimelinePendingOpPersistence, Option<TimelinePersistence>)>> {
1612 : use crate::schema::safekeeper_timeline_pending_ops::dsl;
1613 : use crate::schema::timelines;
1614 :
1615 0 : let timeline_from_db = self
1616 0 : .with_measured_conn(
1617 0 : DatabaseOperation::ListTimelineReconcileStartup,
1618 0 : move |conn| {
1619 0 : Box::pin(async move {
1620 0 : let from_db: Vec<(TimelinePendingOpPersistence, Option<TimelineFromDb>)> =
1621 0 : dsl::safekeeper_timeline_pending_ops
1622 0 : .left_join(
1623 0 : timelines::table.on(timelines::tenant_id
1624 0 : .eq(dsl::tenant_id)
1625 0 : .and(timelines::timeline_id.eq(dsl::timeline_id))),
1626 0 : )
1627 0 : .select((
1628 0 : TimelinePendingOpPersistence::as_select(),
1629 0 : Option::<TimelineFromDb>::as_select(),
1630 0 : ))
1631 0 : .load(conn)
1632 0 : .await?;
1633 0 : Ok(from_db)
1634 0 : })
1635 0 : },
1636 0 : )
1637 0 : .await?;
1638 :
1639 0 : Ok(timeline_from_db
1640 0 : .into_iter()
1641 0 : .map(|(op, tl_opt)| (op, tl_opt.map(|tl_opt| tl_opt.into_persistence())))
1642 0 : .collect())
1643 0 : }
1644 : /// List pending operations for a given timeline (including tenant-global ones)
1645 0 : pub(crate) async fn list_pending_ops_for_timeline(
1646 0 : &self,
1647 0 : tenant_id: TenantId,
1648 0 : timeline_id: TimelineId,
1649 0 : ) -> DatabaseResult<Vec<TimelinePendingOpPersistence>> {
1650 : use crate::schema::safekeeper_timeline_pending_ops::dsl;
1651 :
1652 0 : let timelines_from_db = self
1653 0 : .with_measured_conn(DatabaseOperation::ListTimelineReconcile, move |conn| {
1654 0 : Box::pin(async move {
1655 0 : let from_db: Vec<TimelinePendingOpPersistence> =
1656 0 : dsl::safekeeper_timeline_pending_ops
1657 0 : .filter(dsl::tenant_id.eq(tenant_id.to_string()))
1658 0 : .filter(
1659 0 : dsl::timeline_id
1660 0 : .eq(timeline_id.to_string())
1661 0 : .or(dsl::timeline_id.eq("")),
1662 0 : )
1663 0 : .load(conn)
1664 0 : .await?;
1665 0 : Ok(from_db)
1666 0 : })
1667 0 : })
1668 0 : .await?;
1669 :
1670 0 : Ok(timelines_from_db)
1671 0 : }
1672 :
1673 : /// Delete all pending ops for the given timeline.
1674 : ///
1675 : /// Use this only at timeline deletion, otherwise use generation based APIs
1676 0 : pub(crate) async fn remove_pending_ops_for_timeline(
1677 0 : &self,
1678 0 : tenant_id: TenantId,
1679 0 : timeline_id: Option<TimelineId>,
1680 0 : ) -> DatabaseResult<()> {
1681 : use crate::schema::safekeeper_timeline_pending_ops::dsl;
1682 :
1683 0 : let tenant_id = &tenant_id;
1684 0 : let timeline_id = &timeline_id;
1685 0 : self.with_measured_conn(DatabaseOperation::RemoveTimelineReconcile, move |conn| {
1686 0 : let timeline_id_str = timeline_id.map(|tid| tid.to_string()).unwrap_or_default();
1687 0 : Box::pin(async move {
1688 0 : diesel::delete(dsl::safekeeper_timeline_pending_ops)
1689 0 : .filter(dsl::tenant_id.eq(tenant_id.to_string()))
1690 0 : .filter(dsl::timeline_id.eq(timeline_id_str))
1691 0 : .execute(conn)
1692 0 : .await?;
1693 0 : Ok(())
1694 0 : })
1695 0 : })
1696 0 : .await?;
1697 :
1698 0 : Ok(())
1699 0 : }
1700 :
1701 0 : pub(crate) async fn insert_timeline_import(
1702 0 : &self,
1703 0 : import: TimelineImportPersistence,
1704 0 : ) -> DatabaseResult<bool> {
1705 0 : self.with_measured_conn(DatabaseOperation::InsertTimelineImport, move |conn| {
1706 0 : Box::pin({
1707 0 : let import = import.clone();
1708 0 : async move {
1709 0 : let inserted = diesel::insert_into(crate::schema::timeline_imports::table)
1710 0 : .values(import)
1711 0 : .execute(conn)
1712 0 : .await?;
1713 0 : Ok(inserted == 1)
1714 0 : }
1715 0 : })
1716 0 : })
1717 0 : .await
1718 0 : }
1719 :
1720 0 : pub(crate) async fn list_timeline_imports(&self) -> DatabaseResult<Vec<TimelineImport>> {
1721 : use crate::schema::timeline_imports::dsl;
1722 0 : let persistent = self
1723 0 : .with_measured_conn(DatabaseOperation::ListTimelineImports, move |conn| {
1724 0 : Box::pin(async move {
1725 0 : let from_db: Vec<TimelineImportPersistence> =
1726 0 : dsl::timeline_imports.load(conn).await?;
1727 0 : Ok(from_db)
1728 0 : })
1729 0 : })
1730 0 : .await?;
1731 :
1732 0 : let imports: Result<Vec<TimelineImport>, _> = persistent
1733 0 : .into_iter()
1734 0 : .map(TimelineImport::from_persistent)
1735 0 : .collect();
1736 0 : match imports {
1737 0 : Ok(ok) => Ok(ok.into_iter().collect()),
1738 0 : Err(err) => Err(DatabaseError::Logical(format!(
1739 0 : "failed to deserialize import: {err}"
1740 0 : ))),
1741 : }
1742 0 : }
1743 :
1744 0 : pub(crate) async fn get_timeline_import(
1745 0 : &self,
1746 0 : tenant_id: TenantId,
1747 0 : timeline_id: TimelineId,
1748 0 : ) -> DatabaseResult<Option<TimelineImport>> {
1749 : use crate::schema::timeline_imports::dsl;
1750 0 : let persistent_import = self
1751 0 : .with_measured_conn(DatabaseOperation::ListTimelineImports, move |conn| {
1752 0 : Box::pin(async move {
1753 0 : let mut from_db: Vec<TimelineImportPersistence> = dsl::timeline_imports
1754 0 : .filter(dsl::tenant_id.eq(tenant_id.to_string()))
1755 0 : .filter(dsl::timeline_id.eq(timeline_id.to_string()))
1756 0 : .load(conn)
1757 0 : .await?;
1758 :
1759 0 : if from_db.len() > 1 {
1760 0 : return Err(DatabaseError::Logical(format!(
1761 0 : "unexpected number of rows ({})",
1762 0 : from_db.len()
1763 0 : )));
1764 0 : }
1765 0 :
1766 0 : Ok(from_db.pop())
1767 0 : })
1768 0 : })
1769 0 : .await?;
1770 :
1771 0 : persistent_import
1772 0 : .map(TimelineImport::from_persistent)
1773 0 : .transpose()
1774 0 : .map_err(|err| DatabaseError::Logical(format!("failed to deserialize import: {err}")))
1775 0 : }
1776 :
1777 0 : pub(crate) async fn delete_timeline_import(
1778 0 : &self,
1779 0 : tenant_id: TenantId,
1780 0 : timeline_id: TimelineId,
1781 0 : ) -> DatabaseResult<()> {
1782 : use crate::schema::timeline_imports::dsl;
1783 :
1784 0 : self.with_measured_conn(DatabaseOperation::DeleteTimelineImport, move |conn| {
1785 0 : Box::pin(async move {
1786 0 : diesel::delete(crate::schema::timeline_imports::table)
1787 0 : .filter(
1788 0 : dsl::tenant_id
1789 0 : .eq(tenant_id.to_string())
1790 0 : .and(dsl::timeline_id.eq(timeline_id.to_string())),
1791 0 : )
1792 0 : .execute(conn)
1793 0 : .await?;
1794 :
1795 0 : Ok(())
1796 0 : })
1797 0 : })
1798 0 : .await
1799 0 : }
1800 :
1801 : /// Idempotently update the status of one shard for an ongoing timeline import
1802 : ///
1803 : /// If the update was persisted to the database, then the current state of the
1804 : /// import is returned to the caller. In case of logical errors a bespoke
1805 : /// [`TimelineImportUpdateError`] instance is returned. Other database errors
1806 : /// are covered by the outer [`DatabaseError`].
1807 0 : pub(crate) async fn update_timeline_import(
1808 0 : &self,
1809 0 : tenant_shard_id: TenantShardId,
1810 0 : timeline_id: TimelineId,
1811 0 : shard_status: ShardImportStatus,
1812 0 : ) -> DatabaseResult<Result<Option<TimelineImport>, TimelineImportUpdateError>> {
1813 : use crate::schema::timeline_imports::dsl;
1814 :
1815 0 : self.with_measured_conn(DatabaseOperation::UpdateTimelineImport, move |conn| {
1816 0 : Box::pin({
1817 0 : let shard_status = shard_status.clone();
1818 0 : async move {
1819 : // Load the current state from the database
1820 0 : let mut from_db: Vec<TimelineImportPersistence> = dsl::timeline_imports
1821 0 : .filter(
1822 0 : dsl::tenant_id
1823 0 : .eq(tenant_shard_id.tenant_id.to_string())
1824 0 : .and(dsl::timeline_id.eq(timeline_id.to_string())),
1825 0 : )
1826 0 : .load(conn)
1827 0 : .await?;
1828 :
1829 0 : assert!(from_db.len() <= 1);
1830 :
1831 0 : let mut status = match from_db.pop() {
1832 0 : Some(some) => TimelineImport::from_persistent(some).unwrap(),
1833 : None => {
1834 0 : return Ok(Err(TimelineImportUpdateError::ImportNotFound {
1835 0 : tenant_id: tenant_shard_id.tenant_id,
1836 0 : timeline_id,
1837 0 : }));
1838 : }
1839 : };
1840 :
1841 : // Perform the update in-memory
1842 0 : let follow_up = match status.update(tenant_shard_id.to_index(), shard_status) {
1843 0 : Ok(ok) => ok,
1844 0 : Err(err) => {
1845 0 : return Ok(Err(err));
1846 : }
1847 : };
1848 :
1849 0 : let new_persistent = status.to_persistent();
1850 0 :
1851 0 : // Write back if required (in the same transaction)
1852 0 : match follow_up {
1853 : TimelineImportUpdateFollowUp::Persist => {
1854 0 : let updated = diesel::update(dsl::timeline_imports)
1855 0 : .filter(
1856 0 : dsl::tenant_id
1857 0 : .eq(tenant_shard_id.tenant_id.to_string())
1858 0 : .and(dsl::timeline_id.eq(timeline_id.to_string())),
1859 0 : )
1860 0 : .set(dsl::shard_statuses.eq(new_persistent.shard_statuses))
1861 0 : .execute(conn)
1862 0 : .await?;
1863 :
1864 0 : if updated != 1 {
1865 0 : return Ok(Err(TimelineImportUpdateError::ImportNotFound {
1866 0 : tenant_id: tenant_shard_id.tenant_id,
1867 0 : timeline_id,
1868 0 : }));
1869 0 : }
1870 0 :
1871 0 : Ok(Ok(Some(status)))
1872 : }
1873 0 : TimelineImportUpdateFollowUp::None => Ok(Ok(None)),
1874 : }
1875 0 : }
1876 0 : })
1877 0 : })
1878 0 : .await
1879 0 : }
1880 :
1881 0 : pub(crate) async fn is_tenant_importing_timeline(
1882 0 : &self,
1883 0 : tenant_id: TenantId,
1884 0 : ) -> DatabaseResult<bool> {
1885 : use crate::schema::timeline_imports::dsl;
1886 0 : self.with_measured_conn(DatabaseOperation::IsTenantImportingTimeline, move |conn| {
1887 0 : Box::pin(async move {
1888 0 : let imports: i64 = dsl::timeline_imports
1889 0 : .filter(dsl::tenant_id.eq(tenant_id.to_string()))
1890 0 : .count()
1891 0 : .get_result(conn)
1892 0 : .await?;
1893 :
1894 0 : Ok(imports > 0)
1895 0 : })
1896 0 : })
1897 0 : .await
1898 0 : }
1899 : }
1900 :
1901 0 : pub(crate) fn load_certs() -> anyhow::Result<Arc<rustls::RootCertStore>> {
1902 0 : let der_certs = rustls_native_certs::load_native_certs();
1903 0 :
1904 0 : if !der_certs.errors.is_empty() {
1905 0 : anyhow::bail!("could not parse certificates: {:?}", der_certs.errors);
1906 0 : }
1907 0 :
1908 0 : let mut store = rustls::RootCertStore::empty();
1909 0 : store.add_parsable_certificates(der_certs.certs);
1910 0 : Ok(Arc::new(store))
1911 0 : }
1912 :
1913 : #[derive(Debug)]
1914 : /// A verifier that accepts all certificates (but logs an error still)
1915 : struct AcceptAll(Arc<WebPkiServerVerifier>);
1916 : impl ServerCertVerifier for AcceptAll {
1917 0 : fn verify_server_cert(
1918 0 : &self,
1919 0 : end_entity: &rustls::pki_types::CertificateDer<'_>,
1920 0 : intermediates: &[rustls::pki_types::CertificateDer<'_>],
1921 0 : server_name: &rustls::pki_types::ServerName<'_>,
1922 0 : ocsp_response: &[u8],
1923 0 : now: rustls::pki_types::UnixTime,
1924 0 : ) -> Result<ServerCertVerified, rustls::Error> {
1925 0 : let r =
1926 0 : self.0
1927 0 : .verify_server_cert(end_entity, intermediates, server_name, ocsp_response, now);
1928 0 : if let Err(err) = r {
1929 0 : tracing::info!(
1930 : ?server_name,
1931 0 : "ignoring db connection TLS validation error: {err:?}"
1932 : );
1933 0 : return Ok(ServerCertVerified::assertion());
1934 0 : }
1935 0 : r
1936 0 : }
1937 0 : fn verify_tls12_signature(
1938 0 : &self,
1939 0 : message: &[u8],
1940 0 : cert: &rustls::pki_types::CertificateDer<'_>,
1941 0 : dss: &rustls::DigitallySignedStruct,
1942 0 : ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
1943 0 : self.0.verify_tls12_signature(message, cert, dss)
1944 0 : }
1945 0 : fn verify_tls13_signature(
1946 0 : &self,
1947 0 : message: &[u8],
1948 0 : cert: &rustls::pki_types::CertificateDer<'_>,
1949 0 : dss: &rustls::DigitallySignedStruct,
1950 0 : ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
1951 0 : self.0.verify_tls13_signature(message, cert, dss)
1952 0 : }
1953 0 : fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
1954 0 : self.0.supported_verify_schemes()
1955 0 : }
1956 : }
1957 :
1958 : /// Loads the root certificates and constructs a client config suitable for connecting.
1959 : /// This function is blocking.
1960 0 : fn client_config_with_root_certs() -> anyhow::Result<rustls::ClientConfig> {
1961 0 : let client_config =
1962 0 : rustls::ClientConfig::builder_with_provider(Arc::new(ring::default_provider()))
1963 0 : .with_safe_default_protocol_versions()
1964 0 : .expect("ring should support the default protocol versions");
1965 : static DO_CERT_CHECKS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1966 0 : let do_cert_checks =
1967 0 : DO_CERT_CHECKS.get_or_init(|| std::env::var("STORCON_DB_CERT_CHECKS").is_ok());
1968 0 : Ok(if *do_cert_checks {
1969 0 : client_config
1970 0 : .with_root_certificates(load_certs()?)
1971 0 : .with_no_client_auth()
1972 : } else {
1973 0 : let verifier = AcceptAll(
1974 : WebPkiServerVerifier::builder_with_provider(
1975 0 : load_certs()?,
1976 0 : Arc::new(ring::default_provider()),
1977 0 : )
1978 0 : .build()?,
1979 : );
1980 0 : client_config
1981 0 : .dangerous()
1982 0 : .with_custom_certificate_verifier(Arc::new(verifier))
1983 0 : .with_no_client_auth()
1984 : })
1985 0 : }
1986 :
1987 0 : fn establish_connection_rustls(config: &str) -> BoxFuture<ConnectionResult<AsyncPgConnection>> {
1988 0 : let fut = async {
1989 : // We first set up the way we want rustls to work.
1990 0 : let rustls_config = client_config_with_root_certs()
1991 0 : .map_err(|err| ConnectionError::BadConnection(format!("{err:?}")))?;
1992 0 : let tls = tokio_postgres_rustls::MakeRustlsConnect::new(rustls_config);
1993 0 : let (client, conn) = tokio_postgres::connect(config, tls)
1994 0 : .await
1995 0 : .map_err(|e| ConnectionError::BadConnection(e.to_string()))?;
1996 :
1997 0 : AsyncPgConnection::try_from_client_and_connection(client, conn).await
1998 0 : };
1999 0 : fut.boxed()
2000 0 : }
2001 :
2002 : #[cfg_attr(test, test)]
2003 1 : fn test_config_debug_censors_password() {
2004 1 : let has_pw =
2005 1 : "host=/var/lib/postgresql,localhost port=1234 user=specialuser password='NOT ALLOWED TAG'";
2006 1 : let has_pw_cfg = has_pw.parse::<tokio_postgres::Config>().unwrap();
2007 1 : assert!(format!("{has_pw_cfg:?}").contains("specialuser"));
2008 : // Ensure that the password is not leaked by the debug impl
2009 1 : assert!(!format!("{has_pw_cfg:?}").contains("NOT ALLOWED TAG"));
2010 1 : }
2011 :
2012 0 : fn log_postgres_connstr_info(config_str: &str) -> anyhow::Result<()> {
2013 0 : let config = config_str
2014 0 : .parse::<tokio_postgres::Config>()
2015 0 : .map_err(|_e| anyhow::anyhow!("Couldn't parse config str"))?;
2016 : // We use debug formatting here, and use a unit test to ensure that we don't leak the password.
2017 : // To make extra sure the test gets ran, run it every time the function is called
2018 : // (this is rather cold code, we can afford it).
2019 : #[cfg(not(test))]
2020 0 : test_config_debug_censors_password();
2021 0 : tracing::info!("database connection config: {config:?}");
2022 0 : Ok(())
2023 0 : }
2024 :
2025 : /// Parts of [`crate::tenant_shard::TenantShard`] that are stored durably
2026 : #[derive(
2027 0 : QueryableByName, Queryable, Selectable, Insertable, Serialize, Deserialize, Clone, Eq, PartialEq,
2028 : )]
2029 : #[diesel(table_name = crate::schema::tenant_shards)]
2030 : pub(crate) struct TenantShardPersistence {
2031 : #[serde(default)]
2032 : pub(crate) tenant_id: String,
2033 : #[serde(default)]
2034 : pub(crate) shard_number: i32,
2035 : #[serde(default)]
2036 : pub(crate) shard_count: i32,
2037 : #[serde(default)]
2038 : pub(crate) shard_stripe_size: i32,
2039 :
2040 : // Latest generation number: next time we attach, increment this
2041 : // and use the incremented number when attaching.
2042 : //
2043 : // Generation is only None when first onboarding a tenant, where it may
2044 : // be in PlacementPolicy::Secondary and therefore have no valid generation state.
2045 : pub(crate) generation: Option<i32>,
2046 :
2047 : // Currently attached pageserver
2048 : #[serde(rename = "pageserver")]
2049 : pub(crate) generation_pageserver: Option<i64>,
2050 :
2051 : #[serde(default)]
2052 : pub(crate) placement_policy: String,
2053 : #[serde(default)]
2054 : pub(crate) splitting: SplitState,
2055 : #[serde(default)]
2056 : pub(crate) config: String,
2057 : #[serde(default)]
2058 : pub(crate) scheduling_policy: String,
2059 :
2060 : // Hint that we should attempt to schedule this tenant shard the given
2061 : // availability zone in order to minimise the chances of cross-AZ communication
2062 : // with compute.
2063 : pub(crate) preferred_az_id: Option<String>,
2064 : }
2065 :
2066 : impl TenantShardPersistence {
2067 0 : fn get_shard_count(&self) -> Result<ShardCount, ShardConfigError> {
2068 0 : self.shard_count
2069 0 : .try_into()
2070 0 : .map(ShardCount)
2071 0 : .map_err(|_| ShardConfigError::InvalidCount)
2072 0 : }
2073 :
2074 0 : fn get_shard_number(&self) -> Result<ShardNumber, ShardConfigError> {
2075 0 : self.shard_number
2076 0 : .try_into()
2077 0 : .map(ShardNumber)
2078 0 : .map_err(|_| ShardConfigError::InvalidNumber)
2079 0 : }
2080 :
2081 0 : fn get_stripe_size(&self) -> Result<ShardStripeSize, ShardConfigError> {
2082 0 : self.shard_stripe_size
2083 0 : .try_into()
2084 0 : .map(ShardStripeSize)
2085 0 : .map_err(|_| ShardConfigError::InvalidStripeSize)
2086 0 : }
2087 :
2088 0 : pub(crate) fn get_shard_identity(&self) -> Result<ShardIdentity, ShardConfigError> {
2089 0 : if self.shard_count == 0 {
2090 : // NB: carry over the stripe size from the persisted record, to avoid consistency check
2091 : // failures if the persisted value differs from the default stripe size. The stripe size
2092 : // doesn't really matter for unsharded tenants anyway.
2093 : Ok(ShardIdentity::unsharded_with_stripe_size(
2094 0 : self.get_stripe_size()?,
2095 : ))
2096 : } else {
2097 : Ok(ShardIdentity::new(
2098 0 : self.get_shard_number()?,
2099 0 : self.get_shard_count()?,
2100 0 : self.get_stripe_size()?,
2101 0 : )?)
2102 : }
2103 0 : }
2104 :
2105 0 : pub(crate) fn get_tenant_shard_id(&self) -> anyhow::Result<TenantShardId> {
2106 0 : Ok(TenantShardId {
2107 0 : tenant_id: TenantId::from_str(self.tenant_id.as_str())?,
2108 0 : shard_number: self.get_shard_number()?,
2109 0 : shard_count: self.get_shard_count()?,
2110 : })
2111 0 : }
2112 : }
2113 :
2114 : /// Parts of [`crate::node::Node`] that are stored durably
2115 0 : #[derive(Serialize, Deserialize, Queryable, Selectable, Insertable, Eq, PartialEq)]
2116 : #[diesel(table_name = crate::schema::nodes)]
2117 : pub(crate) struct NodePersistence {
2118 : pub(crate) node_id: i64,
2119 : pub(crate) scheduling_policy: String,
2120 : pub(crate) listen_http_addr: String,
2121 : pub(crate) listen_http_port: i32,
2122 : pub(crate) listen_pg_addr: String,
2123 : pub(crate) listen_pg_port: i32,
2124 : pub(crate) availability_zone_id: String,
2125 : pub(crate) listen_https_port: Option<i32>,
2126 : pub(crate) lifecycle: String,
2127 : }
2128 :
2129 : /// Tenant metadata health status that are stored durably.
2130 0 : #[derive(Queryable, Selectable, Insertable, Serialize, Deserialize, Clone, Eq, PartialEq)]
2131 : #[diesel(table_name = crate::schema::metadata_health)]
2132 : pub(crate) struct MetadataHealthPersistence {
2133 : #[serde(default)]
2134 : pub(crate) tenant_id: String,
2135 : #[serde(default)]
2136 : pub(crate) shard_number: i32,
2137 : #[serde(default)]
2138 : pub(crate) shard_count: i32,
2139 :
2140 : pub(crate) healthy: bool,
2141 : pub(crate) last_scrubbed_at: chrono::DateTime<chrono::Utc>,
2142 : }
2143 :
2144 : impl MetadataHealthPersistence {
2145 0 : pub fn new(
2146 0 : tenant_shard_id: TenantShardId,
2147 0 : healthy: bool,
2148 0 : last_scrubbed_at: chrono::DateTime<chrono::Utc>,
2149 0 : ) -> Self {
2150 0 : let tenant_id = tenant_shard_id.tenant_id.to_string();
2151 0 : let shard_number = tenant_shard_id.shard_number.0 as i32;
2152 0 : let shard_count = tenant_shard_id.shard_count.literal() as i32;
2153 0 :
2154 0 : MetadataHealthPersistence {
2155 0 : tenant_id,
2156 0 : shard_number,
2157 0 : shard_count,
2158 0 : healthy,
2159 0 : last_scrubbed_at,
2160 0 : }
2161 0 : }
2162 :
2163 : #[allow(dead_code)]
2164 0 : pub(crate) fn get_tenant_shard_id(&self) -> Result<TenantShardId, hex::FromHexError> {
2165 0 : Ok(TenantShardId {
2166 0 : tenant_id: TenantId::from_str(self.tenant_id.as_str())?,
2167 0 : shard_number: ShardNumber(self.shard_number as u8),
2168 0 : shard_count: ShardCount::new(self.shard_count as u8),
2169 : })
2170 0 : }
2171 : }
2172 :
2173 : impl From<MetadataHealthPersistence> for MetadataHealthRecord {
2174 0 : fn from(value: MetadataHealthPersistence) -> Self {
2175 0 : MetadataHealthRecord {
2176 0 : tenant_shard_id: value
2177 0 : .get_tenant_shard_id()
2178 0 : .expect("stored tenant id should be valid"),
2179 0 : healthy: value.healthy,
2180 0 : last_scrubbed_at: value.last_scrubbed_at,
2181 0 : }
2182 0 : }
2183 : }
2184 :
2185 : #[derive(
2186 0 : Serialize, Deserialize, Queryable, Selectable, Insertable, Eq, PartialEq, Debug, Clone,
2187 : )]
2188 : #[diesel(table_name = crate::schema::controllers)]
2189 : pub(crate) struct ControllerPersistence {
2190 : pub(crate) address: String,
2191 : pub(crate) started_at: chrono::DateTime<chrono::Utc>,
2192 : }
2193 :
2194 : // What we store in the database
2195 0 : #[derive(Serialize, Deserialize, Queryable, Selectable, Eq, PartialEq, Debug, Clone)]
2196 : #[diesel(table_name = crate::schema::safekeepers)]
2197 : pub(crate) struct SafekeeperPersistence {
2198 : pub(crate) id: i64,
2199 : pub(crate) region_id: String,
2200 : /// 1 is special, it means just created (not currently posted to storcon).
2201 : /// Zero or negative is not really expected.
2202 : /// Otherwise the number from `release-$(number_of_commits_on_branch)` tag.
2203 : pub(crate) version: i64,
2204 : pub(crate) host: String,
2205 : pub(crate) port: i32,
2206 : pub(crate) http_port: i32,
2207 : pub(crate) availability_zone_id: String,
2208 : pub(crate) scheduling_policy: SkSchedulingPolicyFromSql,
2209 : pub(crate) https_port: Option<i32>,
2210 : }
2211 :
2212 : /// Wrapper struct around [`SkSchedulingPolicy`] because both it and [`FromSql`] are from foreign crates,
2213 : /// and we don't want to make [`safekeeper_api`] depend on [`diesel`].
2214 0 : #[derive(Serialize, Deserialize, FromSqlRow, Eq, PartialEq, Debug, Copy, Clone)]
2215 : pub(crate) struct SkSchedulingPolicyFromSql(pub(crate) SkSchedulingPolicy);
2216 :
2217 : impl From<SkSchedulingPolicy> for SkSchedulingPolicyFromSql {
2218 0 : fn from(value: SkSchedulingPolicy) -> Self {
2219 0 : SkSchedulingPolicyFromSql(value)
2220 0 : }
2221 : }
2222 :
2223 : impl FromSql<diesel::sql_types::VarChar, Pg> for SkSchedulingPolicyFromSql {
2224 0 : fn from_sql(
2225 0 : bytes: <Pg as diesel::backend::Backend>::RawValue<'_>,
2226 0 : ) -> diesel::deserialize::Result<Self> {
2227 0 : let bytes = bytes.as_bytes();
2228 0 : match core::str::from_utf8(bytes) {
2229 0 : Ok(s) => match SkSchedulingPolicy::from_str(s) {
2230 0 : Ok(policy) => Ok(SkSchedulingPolicyFromSql(policy)),
2231 0 : Err(e) => Err(format!("can't parse: {e}").into()),
2232 : },
2233 0 : Err(e) => Err(format!("invalid UTF-8 for scheduling policy: {e}").into()),
2234 : }
2235 0 : }
2236 : }
2237 :
2238 : impl SafekeeperPersistence {
2239 0 : pub(crate) fn from_upsert(
2240 0 : upsert: SafekeeperUpsert,
2241 0 : scheduling_policy: SkSchedulingPolicy,
2242 0 : ) -> Self {
2243 0 : crate::persistence::SafekeeperPersistence {
2244 0 : id: upsert.id,
2245 0 : region_id: upsert.region_id,
2246 0 : version: upsert.version,
2247 0 : host: upsert.host,
2248 0 : port: upsert.port,
2249 0 : http_port: upsert.http_port,
2250 0 : https_port: upsert.https_port,
2251 0 : availability_zone_id: upsert.availability_zone_id,
2252 0 : scheduling_policy: SkSchedulingPolicyFromSql(scheduling_policy),
2253 0 : }
2254 0 : }
2255 0 : pub(crate) fn as_describe_response(&self) -> Result<SafekeeperDescribeResponse, DatabaseError> {
2256 0 : Ok(SafekeeperDescribeResponse {
2257 0 : id: NodeId(self.id as u64),
2258 0 : region_id: self.region_id.clone(),
2259 0 : version: self.version,
2260 0 : host: self.host.clone(),
2261 0 : port: self.port,
2262 0 : http_port: self.http_port,
2263 0 : https_port: self.https_port,
2264 0 : availability_zone_id: self.availability_zone_id.clone(),
2265 0 : scheduling_policy: self.scheduling_policy.0,
2266 0 : })
2267 0 : }
2268 : }
2269 :
2270 : /// What we expect from the upsert http api
2271 0 : #[derive(Serialize, Deserialize, Eq, PartialEq, Debug, Clone)]
2272 : pub(crate) struct SafekeeperUpsert {
2273 : pub(crate) id: i64,
2274 : pub(crate) region_id: String,
2275 : /// 1 is special, it means just created (not currently posted to storcon).
2276 : /// Zero or negative is not really expected.
2277 : /// Otherwise the number from `release-$(number_of_commits_on_branch)` tag.
2278 : pub(crate) version: i64,
2279 : pub(crate) host: String,
2280 : pub(crate) port: i32,
2281 : /// The active flag will not be stored in the database and will be ignored.
2282 : pub(crate) active: Option<bool>,
2283 : pub(crate) http_port: i32,
2284 : pub(crate) https_port: Option<i32>,
2285 : pub(crate) availability_zone_id: String,
2286 : }
2287 :
2288 : impl SafekeeperUpsert {
2289 0 : fn as_insert_or_update(&self) -> anyhow::Result<InsertUpdateSafekeeper<'_>> {
2290 0 : if self.version < 0 {
2291 0 : anyhow::bail!("negative version: {}", self.version);
2292 0 : }
2293 0 : Ok(InsertUpdateSafekeeper {
2294 0 : id: self.id,
2295 0 : region_id: &self.region_id,
2296 0 : version: self.version,
2297 0 : host: &self.host,
2298 0 : port: self.port,
2299 0 : http_port: self.http_port,
2300 0 : https_port: self.https_port,
2301 0 : availability_zone_id: &self.availability_zone_id,
2302 0 : // None means a wish to not update this column. We expose abilities to update it via other means.
2303 0 : scheduling_policy: None,
2304 0 : })
2305 0 : }
2306 : }
2307 :
2308 0 : #[derive(Insertable, AsChangeset)]
2309 : #[diesel(table_name = crate::schema::safekeepers)]
2310 : struct InsertUpdateSafekeeper<'a> {
2311 : id: i64,
2312 : region_id: &'a str,
2313 : version: i64,
2314 : host: &'a str,
2315 : port: i32,
2316 : http_port: i32,
2317 : https_port: Option<i32>,
2318 : availability_zone_id: &'a str,
2319 : scheduling_policy: Option<&'a str>,
2320 : }
2321 :
2322 0 : #[derive(Serialize, Deserialize, FromSqlRow, AsExpression, Eq, PartialEq, Debug, Copy, Clone)]
2323 : #[diesel(sql_type = crate::schema::sql_types::PgLsn)]
2324 : pub(crate) struct LsnWrapper(pub(crate) Lsn);
2325 :
2326 : impl From<Lsn> for LsnWrapper {
2327 0 : fn from(value: Lsn) -> Self {
2328 0 : LsnWrapper(value)
2329 0 : }
2330 : }
2331 :
2332 : impl FromSql<crate::schema::sql_types::PgLsn, Pg> for LsnWrapper {
2333 0 : fn from_sql(
2334 0 : bytes: <Pg as diesel::backend::Backend>::RawValue<'_>,
2335 0 : ) -> diesel::deserialize::Result<Self> {
2336 0 : let byte_arr: diesel::deserialize::Result<[u8; 8]> = bytes
2337 0 : .as_bytes()
2338 0 : .try_into()
2339 0 : .map_err(|_| "Can't obtain lsn from sql".into());
2340 0 : Ok(LsnWrapper(Lsn(u64::from_be_bytes(byte_arr?))))
2341 0 : }
2342 : }
2343 :
2344 : impl ToSql<crate::schema::sql_types::PgLsn, Pg> for LsnWrapper {
2345 0 : fn to_sql<'b>(
2346 0 : &'b self,
2347 0 : out: &mut diesel::serialize::Output<'b, '_, Pg>,
2348 0 : ) -> diesel::serialize::Result {
2349 0 : out.write_all(&u64::to_be_bytes(self.0.0))
2350 0 : .map(|_| IsNull::No)
2351 0 : .map_err(Into::into)
2352 0 : }
2353 : }
2354 :
2355 0 : #[derive(Insertable, AsChangeset, Clone)]
2356 : #[diesel(table_name = crate::schema::timelines)]
2357 : pub(crate) struct TimelinePersistence {
2358 : pub(crate) tenant_id: String,
2359 : pub(crate) timeline_id: String,
2360 : pub(crate) start_lsn: LsnWrapper,
2361 : pub(crate) generation: i32,
2362 : pub(crate) sk_set: Vec<i64>,
2363 : pub(crate) new_sk_set: Option<Vec<i64>>,
2364 : pub(crate) cplane_notified_generation: i32,
2365 : pub(crate) deleted_at: Option<chrono::DateTime<chrono::Utc>>,
2366 : }
2367 :
2368 : /// This is separate from [TimelinePersistence] only because postgres allows NULLs
2369 : /// in arrays and there is no way to forbid that at schema level. Hence diesel
2370 : /// wants `sk_set` to be `Vec<Option<i64>>` instead of `Vec<i64>` for
2371 : /// Queryable/Selectable. It does however allow insertions without redundant
2372 : /// Option(s), so [TimelinePersistence] doesn't have them.
2373 0 : #[derive(Queryable, Selectable)]
2374 : #[diesel(table_name = crate::schema::timelines)]
2375 : pub(crate) struct TimelineFromDb {
2376 : pub(crate) tenant_id: String,
2377 : pub(crate) timeline_id: String,
2378 : pub(crate) start_lsn: LsnWrapper,
2379 : pub(crate) generation: i32,
2380 : pub(crate) sk_set: Vec<Option<i64>>,
2381 : pub(crate) new_sk_set: Option<Vec<Option<i64>>>,
2382 : pub(crate) cplane_notified_generation: i32,
2383 : pub(crate) deleted_at: Option<chrono::DateTime<chrono::Utc>>,
2384 : }
2385 :
2386 : impl TimelineFromDb {
2387 0 : fn into_persistence(self) -> TimelinePersistence {
2388 0 : // We should never encounter null entries in the sets, but we need to filter them out.
2389 0 : // There is no way to forbid this in the schema that diesel recognizes (to our knowledge).
2390 0 : let sk_set = self.sk_set.into_iter().flatten().collect::<Vec<_>>();
2391 0 : let new_sk_set = self
2392 0 : .new_sk_set
2393 0 : .map(|s| s.into_iter().flatten().collect::<Vec<_>>());
2394 0 : TimelinePersistence {
2395 0 : tenant_id: self.tenant_id,
2396 0 : timeline_id: self.timeline_id,
2397 0 : start_lsn: self.start_lsn,
2398 0 : generation: self.generation,
2399 0 : sk_set,
2400 0 : new_sk_set,
2401 0 : cplane_notified_generation: self.cplane_notified_generation,
2402 0 : deleted_at: self.deleted_at,
2403 0 : }
2404 0 : }
2405 : }
2406 :
2407 0 : #[derive(Insertable, AsChangeset, Queryable, Selectable, Clone)]
2408 : #[diesel(table_name = crate::schema::safekeeper_timeline_pending_ops)]
2409 : pub(crate) struct TimelinePendingOpPersistence {
2410 : pub(crate) sk_id: i64,
2411 : pub(crate) tenant_id: String,
2412 : pub(crate) timeline_id: String,
2413 : pub(crate) generation: i32,
2414 : pub(crate) op_kind: SafekeeperTimelineOpKind,
2415 : }
2416 :
2417 0 : #[derive(Serialize, Deserialize, FromSqlRow, AsExpression, Eq, PartialEq, Debug, Copy, Clone)]
2418 : #[diesel(sql_type = diesel::sql_types::VarChar)]
2419 : pub(crate) enum SafekeeperTimelineOpKind {
2420 : Pull,
2421 : Exclude,
2422 : Delete,
2423 : }
2424 :
2425 : impl FromSql<diesel::sql_types::VarChar, Pg> for SafekeeperTimelineOpKind {
2426 0 : fn from_sql(
2427 0 : bytes: <Pg as diesel::backend::Backend>::RawValue<'_>,
2428 0 : ) -> diesel::deserialize::Result<Self> {
2429 0 : let bytes = bytes.as_bytes();
2430 0 : match core::str::from_utf8(bytes) {
2431 0 : Ok(s) => match s {
2432 0 : "pull" => Ok(SafekeeperTimelineOpKind::Pull),
2433 0 : "exclude" => Ok(SafekeeperTimelineOpKind::Exclude),
2434 0 : "delete" => Ok(SafekeeperTimelineOpKind::Delete),
2435 0 : _ => Err(format!("can't parse: {s}").into()),
2436 : },
2437 0 : Err(e) => Err(format!("invalid UTF-8 for op_kind: {e}").into()),
2438 : }
2439 0 : }
2440 : }
2441 :
2442 : impl ToSql<diesel::sql_types::VarChar, Pg> for SafekeeperTimelineOpKind {
2443 0 : fn to_sql<'b>(
2444 0 : &'b self,
2445 0 : out: &mut diesel::serialize::Output<'b, '_, Pg>,
2446 0 : ) -> diesel::serialize::Result {
2447 0 : let kind_str = match self {
2448 0 : SafekeeperTimelineOpKind::Pull => "pull",
2449 0 : SafekeeperTimelineOpKind::Exclude => "exclude",
2450 0 : SafekeeperTimelineOpKind::Delete => "delete",
2451 : };
2452 0 : out.write_all(kind_str.as_bytes())
2453 0 : .map(|_| IsNull::No)
2454 0 : .map_err(Into::into)
2455 0 : }
2456 : }
2457 :
2458 0 : #[derive(Serialize, Deserialize, Queryable, Selectable, Insertable, Eq, PartialEq, Clone)]
2459 : #[diesel(table_name = crate::schema::timeline_imports)]
2460 : pub(crate) struct TimelineImportPersistence {
2461 : pub(crate) tenant_id: String,
2462 : pub(crate) timeline_id: String,
2463 : pub(crate) shard_statuses: serde_json::Value,
2464 : }
|
