Line data Source code
1 : pub(crate) mod split_state;
2 : use std::collections::HashMap;
3 : use std::str::FromStr;
4 : use std::time::Duration;
5 : use std::time::Instant;
6 :
7 : use self::split_state::SplitState;
8 : use diesel::prelude::*;
9 : use diesel_async::async_connection_wrapper::AsyncConnectionWrapper;
10 : use diesel_async::pooled_connection::bb8::Pool;
11 : use diesel_async::pooled_connection::AsyncDieselConnectionManager;
12 : use diesel_async::RunQueryDsl;
13 : use diesel_async::{AsyncConnection, AsyncPgConnection};
14 : use itertools::Itertools;
15 : use pageserver_api::controller_api::AvailabilityZone;
16 : use pageserver_api::controller_api::MetadataHealthRecord;
17 : use pageserver_api::controller_api::SafekeeperDescribeResponse;
18 : use pageserver_api::controller_api::ShardSchedulingPolicy;
19 : use pageserver_api::controller_api::SkSchedulingPolicy;
20 : use pageserver_api::controller_api::{NodeSchedulingPolicy, PlacementPolicy};
21 : use pageserver_api::models::TenantConfig;
22 : use pageserver_api::shard::ShardConfigError;
23 : use pageserver_api::shard::ShardIdentity;
24 : use pageserver_api::shard::ShardStripeSize;
25 : use pageserver_api::shard::{ShardCount, ShardNumber, TenantShardId};
26 : use scoped_futures::ScopedBoxFuture;
27 : use serde::{Deserialize, Serialize};
28 : use utils::generation::Generation;
29 : use utils::id::{NodeId, TenantId};
30 :
31 : use crate::metrics::{
32 : DatabaseQueryErrorLabelGroup, DatabaseQueryLatencyLabelGroup, METRICS_REGISTRY,
33 : };
34 : use crate::node::Node;
35 :
36 : use diesel_migrations::{embed_migrations, EmbeddedMigrations};
37 : const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations");
38 :
39 : /// ## What do we store?
40 : ///
41 : /// The storage controller service does not store most of its state durably.
42 : ///
43 : /// The essential things to store durably are:
44 : /// - generation numbers, as these must always advance monotonically to ensure data safety.
45 : /// - Tenant's PlacementPolicy and TenantConfig, as the source of truth for these is something external.
46 : /// - Node's scheduling policies, as the source of truth for these is something external.
47 : ///
48 : /// Other things we store durably as an implementation detail:
49 : /// - Node's host/port: this could be avoided it we made nodes emit a self-registering heartbeat,
50 : /// but it is operationally simpler to make this service the authority for which nodes
51 : /// it talks to.
52 : ///
53 : /// ## Performance/efficiency
54 : ///
55 : /// The storage controller service does not go via the database for most things: there are
56 : /// a couple of places where we must, and where efficiency matters:
57 : /// - Incrementing generation numbers: the Reconciler has to wait for this to complete
58 : /// before it can attach a tenant, so this acts as a bound on how fast things like
59 : /// failover can happen.
60 : /// - Pageserver re-attach: we will increment many shards' generations when this happens,
61 : /// so it is important to avoid e.g. issuing O(N) queries.
62 : ///
63 : /// Database calls relating to nodes have low performance requirements, as they are very rarely
64 : /// updated, and reads of nodes are always from memory, not the database. We only require that
65 : /// we can UPDATE a node's scheduling mode reasonably quickly to mark a bad node offline.
66 : pub struct Persistence {
67 : connection_pool: Pool<AsyncPgConnection>,
68 : }
69 :
70 : /// Legacy format, for use in JSON compat objects in test environment
71 0 : #[derive(Serialize, Deserialize)]
72 : struct JsonPersistence {
73 : tenants: HashMap<TenantShardId, TenantShardPersistence>,
74 : }
75 :
76 : #[derive(thiserror::Error, Debug)]
77 : pub(crate) enum DatabaseError {
78 : #[error(transparent)]
79 : Query(#[from] diesel::result::Error),
80 : #[error(transparent)]
81 : Connection(#[from] diesel::result::ConnectionError),
82 : #[error(transparent)]
83 : ConnectionPool(#[from] diesel_async::pooled_connection::bb8::RunError),
84 : #[error("Logical error: {0}")]
85 : Logical(String),
86 : #[error("Migration error: {0}")]
87 : Migration(String),
88 : }
89 :
90 : #[derive(measured::FixedCardinalityLabel, Copy, Clone)]
91 : pub(crate) enum DatabaseOperation {
92 : InsertNode,
93 : UpdateNode,
94 : DeleteNode,
95 : ListNodes,
96 : BeginShardSplit,
97 : CompleteShardSplit,
98 : AbortShardSplit,
99 : Detach,
100 : ReAttach,
101 : IncrementGeneration,
102 : TenantGenerations,
103 : ShardGenerations,
104 : ListTenantShards,
105 : LoadTenant,
106 : InsertTenantShards,
107 : UpdateTenantShard,
108 : DeleteTenant,
109 : UpdateTenantConfig,
110 : UpdateMetadataHealth,
111 : ListMetadataHealth,
112 : ListMetadataHealthUnhealthy,
113 : ListMetadataHealthOutdated,
114 : ListSafekeepers,
115 : GetLeader,
116 : UpdateLeader,
117 : SetPreferredAzs,
118 : }
119 :
120 : #[must_use]
121 : pub(crate) enum AbortShardSplitStatus {
122 : /// We aborted the split in the database by reverting to the parent shards
123 : Aborted,
124 : /// The split had already been persisted.
125 : Complete,
126 : }
127 :
128 : pub(crate) type DatabaseResult<T> = Result<T, DatabaseError>;
129 :
130 : /// Some methods can operate on either a whole tenant or a single shard
131 : #[derive(Clone)]
132 : pub(crate) enum TenantFilter {
133 : Tenant(TenantId),
134 : Shard(TenantShardId),
135 : }
136 :
137 : /// Represents the results of looking up generation+pageserver for the shards of a tenant
138 : pub(crate) struct ShardGenerationState {
139 : pub(crate) tenant_shard_id: TenantShardId,
140 : pub(crate) generation: Option<Generation>,
141 : pub(crate) generation_pageserver: Option<NodeId>,
142 : }
143 :
144 : // A generous allowance for how many times we may retry serializable transactions
145 : // before giving up. This is not expected to be hit: it is a defensive measure in case we
146 : // somehow engineer a situation where duelling transactions might otherwise live-lock.
147 : const MAX_RETRIES: usize = 128;
148 :
149 : impl Persistence {
150 : // The default postgres connection limit is 100. We use up to 99, to leave one free for a human admin under
151 : // normal circumstances. This assumes we have exclusive use of the database cluster to which we connect.
152 : pub const MAX_CONNECTIONS: u32 = 99;
153 :
154 : // We don't want to keep a lot of connections alive: close them down promptly if they aren't being used.
155 : const IDLE_CONNECTION_TIMEOUT: Duration = Duration::from_secs(10);
156 : const MAX_CONNECTION_LIFETIME: Duration = Duration::from_secs(60);
157 :
158 0 : pub async fn new(database_url: String) -> Self {
159 0 : let manager = AsyncDieselConnectionManager::<AsyncPgConnection>::new(database_url);
160 :
161 : // We will use a connection pool: this is primarily to _limit_ our connection count, rather than to optimize time
162 : // to execute queries (database queries are not generally on latency-sensitive paths).
163 0 : let connection_pool = Pool::builder()
164 0 : .max_size(Self::MAX_CONNECTIONS)
165 0 : .max_lifetime(Some(Self::MAX_CONNECTION_LIFETIME))
166 0 : .idle_timeout(Some(Self::IDLE_CONNECTION_TIMEOUT))
167 0 : // Always keep at least one connection ready to go
168 0 : .min_idle(Some(1))
169 0 : .test_on_check_out(true)
170 0 : .build(manager)
171 0 : .await
172 0 : .expect("Could not build connection pool");
173 0 :
174 0 : Self { connection_pool }
175 0 : }
176 :
177 : /// A helper for use during startup, where we would like to tolerate concurrent restarts of the
178 : /// database and the storage controller, therefore the database might not be available right away
179 0 : pub async fn await_connection(
180 0 : database_url: &str,
181 0 : timeout: Duration,
182 0 : ) -> Result<(), diesel::ConnectionError> {
183 0 : let started_at = Instant::now();
184 : loop {
185 0 : match AsyncPgConnection::establish(database_url).await {
186 : Ok(_) => {
187 0 : tracing::info!("Connected to database.");
188 0 : return Ok(());
189 : }
190 0 : Err(e) => {
191 0 : if started_at.elapsed() > timeout {
192 0 : return Err(e);
193 : } else {
194 0 : tracing::info!("Database not yet available, waiting... ({e})");
195 0 : tokio::time::sleep(Duration::from_millis(100)).await;
196 : }
197 : }
198 : }
199 : }
200 0 : }
201 :
202 : /// Execute the diesel migrations that are built into this binary
203 0 : pub(crate) async fn migration_run(&self) -> DatabaseResult<()> {
204 : use diesel_migrations::{HarnessWithOutput, MigrationHarness};
205 :
206 : // Can't use self.with_conn here as we do spawn_blocking which requires static.
207 0 : let conn = self
208 0 : .connection_pool
209 0 : .dedicated_connection()
210 0 : .await
211 0 : .map_err(|e| DatabaseError::Migration(e.to_string()))?;
212 0 : let mut async_wrapper: AsyncConnectionWrapper<AsyncPgConnection> =
213 0 : AsyncConnectionWrapper::from(conn);
214 0 : tokio::task::spawn_blocking(move || {
215 0 : let mut retry_count = 0;
216 0 : loop {
217 0 : let result = HarnessWithOutput::write_to_stdout(&mut async_wrapper)
218 0 : .run_pending_migrations(MIGRATIONS)
219 0 : .map(|_| ())
220 0 : .map_err(|e| DatabaseError::Migration(e.to_string()));
221 0 : match result {
222 0 : Ok(r) => break Ok(r),
223 : Err(
224 0 : err @ DatabaseError::Query(diesel::result::Error::DatabaseError(
225 0 : diesel::result::DatabaseErrorKind::SerializationFailure,
226 0 : _,
227 0 : )),
228 0 : ) => {
229 0 : retry_count += 1;
230 0 : if retry_count > MAX_RETRIES {
231 0 : tracing::error!(
232 0 : "Exceeded max retries on SerializationFailure errors: {err:?}"
233 : );
234 0 : break Err(err);
235 : } else {
236 : // Retry on serialization errors: these are expected, because even though our
237 : // transactions don't fight for the same rows, they will occasionally collide
238 : // on index pages (e.g. increment_generation for unrelated shards can collide)
239 0 : tracing::debug!(
240 0 : "Retrying transaction on serialization failure {err:?}"
241 : );
242 0 : continue;
243 : }
244 : }
245 0 : Err(e) => break Err(e),
246 : }
247 : }
248 0 : })
249 0 : .await
250 0 : .map_err(|e| DatabaseError::Migration(e.to_string()))??;
251 0 : Ok(())
252 0 : }
253 :
254 : /// Wraps `with_conn` in order to collect latency and error metrics
255 0 : async fn with_measured_conn<'a, 'b, F, R>(
256 0 : &self,
257 0 : op: DatabaseOperation,
258 0 : func: F,
259 0 : ) -> DatabaseResult<R>
260 0 : where
261 0 : F: for<'r> Fn(&'r mut AsyncPgConnection) -> ScopedBoxFuture<'b, 'r, DatabaseResult<R>>
262 0 : + Send
263 0 : + std::marker::Sync
264 0 : + 'a,
265 0 : R: Send + 'b,
266 0 : {
267 0 : let latency = &METRICS_REGISTRY
268 0 : .metrics_group
269 0 : .storage_controller_database_query_latency;
270 0 : let _timer = latency.start_timer(DatabaseQueryLatencyLabelGroup { operation: op });
271 :
272 0 : let res = self.with_conn(func).await;
273 :
274 0 : if let Err(err) = &res {
275 0 : let error_counter = &METRICS_REGISTRY
276 0 : .metrics_group
277 0 : .storage_controller_database_query_error;
278 0 : error_counter.inc(DatabaseQueryErrorLabelGroup {
279 0 : error_type: err.error_label(),
280 0 : operation: op,
281 0 : })
282 0 : }
283 :
284 0 : res
285 0 : }
286 :
287 : /// Call the provided function with a Diesel database connection in a retry loop
288 0 : async fn with_conn<'a, 'b, F, R>(&self, func: F) -> 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 mut retry_count = 0;
297 : loop {
298 0 : let mut conn = self.connection_pool.get().await?;
299 0 : match conn
300 0 : .build_transaction()
301 0 : .serializable()
302 0 : .run(|c| func(c))
303 0 : .await
304 : {
305 0 : Ok(r) => break Ok(r),
306 : Err(
307 0 : err @ DatabaseError::Query(diesel::result::Error::DatabaseError(
308 0 : diesel::result::DatabaseErrorKind::SerializationFailure,
309 0 : _,
310 0 : )),
311 0 : ) => {
312 0 : retry_count += 1;
313 0 : if retry_count > MAX_RETRIES {
314 0 : tracing::error!(
315 0 : "Exceeded max retries on SerializationFailure errors: {err:?}"
316 : );
317 0 : break Err(err);
318 : } else {
319 : // Retry on serialization errors: these are expected, because even though our
320 : // transactions don't fight for the same rows, they will occasionally collide
321 : // on index pages (e.g. increment_generation for unrelated shards can collide)
322 0 : tracing::debug!("Retrying transaction on serialization failure {err:?}");
323 0 : continue;
324 : }
325 : }
326 0 : Err(e) => break Err(e),
327 : }
328 : }
329 0 : }
330 :
331 : /// When a node is first registered, persist it before using it for anything
332 0 : pub(crate) async fn insert_node(&self, node: &Node) -> DatabaseResult<()> {
333 0 : let np = &node.to_persistent();
334 0 : self.with_measured_conn(DatabaseOperation::InsertNode, move |conn| {
335 0 : Box::pin(async move {
336 0 : diesel::insert_into(crate::schema::nodes::table)
337 0 : .values(np)
338 0 : .execute(conn)
339 0 : .await?;
340 0 : Ok(())
341 0 : })
342 0 : })
343 0 : .await
344 0 : }
345 :
346 : /// At startup, populate the list of nodes which our shards may be placed on
347 0 : pub(crate) async fn list_nodes(&self) -> DatabaseResult<Vec<NodePersistence>> {
348 0 : let nodes: Vec<NodePersistence> = self
349 0 : .with_measured_conn(DatabaseOperation::ListNodes, move |conn| {
350 0 : Box::pin(async move {
351 0 : Ok(crate::schema::nodes::table
352 0 : .load::<NodePersistence>(conn)
353 0 : .await?)
354 0 : })
355 0 : })
356 0 : .await?;
357 :
358 0 : tracing::info!("list_nodes: loaded {} nodes", nodes.len());
359 :
360 0 : Ok(nodes)
361 0 : }
362 :
363 0 : pub(crate) async fn update_node(
364 0 : &self,
365 0 : input_node_id: NodeId,
366 0 : input_scheduling: NodeSchedulingPolicy,
367 0 : ) -> DatabaseResult<()> {
368 : use crate::schema::nodes::dsl::*;
369 0 : let updated = self
370 0 : .with_measured_conn(DatabaseOperation::UpdateNode, move |conn| {
371 0 : Box::pin(async move {
372 0 : let updated = diesel::update(nodes)
373 0 : .filter(node_id.eq(input_node_id.0 as i64))
374 0 : .set((scheduling_policy.eq(String::from(input_scheduling)),))
375 0 : .execute(conn)
376 0 : .await?;
377 0 : Ok(updated)
378 0 : })
379 0 : })
380 0 : .await?;
381 :
382 0 : if updated != 1 {
383 0 : Err(DatabaseError::Logical(format!(
384 0 : "Node {node_id:?} not found for update",
385 0 : )))
386 : } else {
387 0 : Ok(())
388 : }
389 0 : }
390 :
391 : /// At startup, load the high level state for shards, such as their config + policy. This will
392 : /// be enriched at runtime with state discovered on pageservers.
393 : ///
394 : /// We exclude shards configured to be detached. During startup, if we see any attached locations
395 : /// for such shards, they will automatically be detached as 'orphans'.
396 0 : pub(crate) async fn load_active_tenant_shards(
397 0 : &self,
398 0 : ) -> DatabaseResult<Vec<TenantShardPersistence>> {
399 : use crate::schema::tenant_shards::dsl::*;
400 0 : self.with_measured_conn(DatabaseOperation::ListTenantShards, move |conn| {
401 0 : Box::pin(async move {
402 0 : let query = tenant_shards.filter(
403 0 : placement_policy.ne(serde_json::to_string(&PlacementPolicy::Detached).unwrap()),
404 0 : );
405 0 : let result = query.load::<TenantShardPersistence>(conn).await?;
406 :
407 0 : Ok(result)
408 0 : })
409 0 : })
410 0 : .await
411 0 : }
412 :
413 : /// When restoring a previously detached tenant into memory, load it from the database
414 0 : pub(crate) async fn load_tenant(
415 0 : &self,
416 0 : filter_tenant_id: TenantId,
417 0 : ) -> DatabaseResult<Vec<TenantShardPersistence>> {
418 : use crate::schema::tenant_shards::dsl::*;
419 0 : self.with_measured_conn(DatabaseOperation::LoadTenant, move |conn| {
420 0 : Box::pin(async move {
421 0 : let query = tenant_shards.filter(tenant_id.eq(filter_tenant_id.to_string()));
422 0 : let result = query.load::<TenantShardPersistence>(conn).await?;
423 :
424 0 : Ok(result)
425 0 : })
426 0 : })
427 0 : .await
428 0 : }
429 :
430 : /// Tenants must be persisted before we schedule them for the first time. This enables us
431 : /// to correctly retain generation monotonicity, and the externally provided placement policy & config.
432 0 : pub(crate) async fn insert_tenant_shards(
433 0 : &self,
434 0 : shards: Vec<TenantShardPersistence>,
435 0 : ) -> DatabaseResult<()> {
436 : use crate::schema::metadata_health;
437 : use crate::schema::tenant_shards;
438 :
439 0 : let now = chrono::Utc::now();
440 0 :
441 0 : let metadata_health_records = shards
442 0 : .iter()
443 0 : .map(|t| MetadataHealthPersistence {
444 0 : tenant_id: t.tenant_id.clone(),
445 0 : shard_number: t.shard_number,
446 0 : shard_count: t.shard_count,
447 0 : healthy: true,
448 0 : last_scrubbed_at: now,
449 0 : })
450 0 : .collect::<Vec<_>>();
451 0 :
452 0 : let shards = &shards;
453 0 : let metadata_health_records = &metadata_health_records;
454 0 : self.with_measured_conn(DatabaseOperation::InsertTenantShards, move |conn| {
455 0 : Box::pin(async move {
456 0 : diesel::insert_into(tenant_shards::table)
457 0 : .values(shards)
458 0 : .execute(conn)
459 0 : .await?;
460 :
461 0 : diesel::insert_into(metadata_health::table)
462 0 : .values(metadata_health_records)
463 0 : .execute(conn)
464 0 : .await?;
465 0 : Ok(())
466 0 : })
467 0 : })
468 0 : .await
469 0 : }
470 :
471 : /// Ordering: call this _after_ deleting the tenant on pageservers, but _before_ dropping state for
472 : /// the tenant from memory on this server.
473 0 : pub(crate) async fn delete_tenant(&self, del_tenant_id: TenantId) -> DatabaseResult<()> {
474 : use crate::schema::tenant_shards::dsl::*;
475 0 : self.with_measured_conn(DatabaseOperation::DeleteTenant, move |conn| {
476 0 : Box::pin(async move {
477 0 : // `metadata_health` status (if exists) is also deleted based on the cascade behavior.
478 0 : diesel::delete(tenant_shards)
479 0 : .filter(tenant_id.eq(del_tenant_id.to_string()))
480 0 : .execute(conn)
481 0 : .await?;
482 0 : Ok(())
483 0 : })
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 0 : diesel::delete(nodes)
493 0 : .filter(node_id.eq(del_node_id.0 as i64))
494 0 : .execute(conn)
495 0 : .await?;
496 :
497 0 : Ok(())
498 0 : })
499 0 : })
500 0 : .await
501 0 : }
502 :
503 : /// When a tenant invokes the /re-attach API, this function is responsible for doing an efficient
504 : /// batched increment of the generations of all tenants whose generation_pageserver is equal to
505 : /// the node that called /re-attach.
506 : #[tracing::instrument(skip_all, fields(node_id))]
507 : pub(crate) async fn re_attach(
508 : &self,
509 : input_node_id: NodeId,
510 : ) -> DatabaseResult<HashMap<TenantShardId, Generation>> {
511 : use crate::schema::nodes::dsl::scheduling_policy;
512 : use crate::schema::nodes::dsl::*;
513 : use crate::schema::tenant_shards::dsl::*;
514 : let updated = self
515 0 : .with_measured_conn(DatabaseOperation::ReAttach, move |conn| {
516 0 : Box::pin(async move {
517 0 : let rows_updated = diesel::update(tenant_shards)
518 0 : .filter(generation_pageserver.eq(input_node_id.0 as i64))
519 0 : .set(generation.eq(generation + 1))
520 0 : .execute(conn)
521 0 : .await?;
522 :
523 0 : tracing::info!("Incremented {} tenants' generations", rows_updated);
524 :
525 : // TODO: UPDATE+SELECT in one query
526 :
527 0 : let updated = tenant_shards
528 0 : .filter(generation_pageserver.eq(input_node_id.0 as i64))
529 0 : .select(TenantShardPersistence::as_select())
530 0 : .load(conn)
531 0 : .await?;
532 :
533 : // If the node went through a drain and restart phase before re-attaching,
534 : // then reset it's node scheduling policy to active.
535 0 : diesel::update(nodes)
536 0 : .filter(node_id.eq(input_node_id.0 as i64))
537 0 : .filter(
538 0 : scheduling_policy
539 0 : .eq(String::from(NodeSchedulingPolicy::PauseForRestart))
540 0 : .or(scheduling_policy
541 0 : .eq(String::from(NodeSchedulingPolicy::Draining)))
542 0 : .or(scheduling_policy
543 0 : .eq(String::from(NodeSchedulingPolicy::Filling))),
544 0 : )
545 0 : .set(scheduling_policy.eq(String::from(NodeSchedulingPolicy::Active)))
546 0 : .execute(conn)
547 0 : .await?;
548 :
549 0 : Ok(updated)
550 0 : })
551 0 : })
552 : .await?;
553 :
554 : let mut result = HashMap::new();
555 : for tsp in updated {
556 : let tenant_shard_id = TenantShardId {
557 : tenant_id: TenantId::from_str(tsp.tenant_id.as_str())
558 0 : .map_err(|e| DatabaseError::Logical(format!("Malformed tenant id: {e}")))?,
559 : shard_number: ShardNumber(tsp.shard_number as u8),
560 : shard_count: ShardCount::new(tsp.shard_count as u8),
561 : };
562 :
563 : let Some(g) = tsp.generation else {
564 : // If the generation_pageserver column was non-NULL, then the generation column should also be non-NULL:
565 : // we only set generation_pageserver when setting generation.
566 : return Err(DatabaseError::Logical(
567 : "Generation should always be set after incrementing".to_string(),
568 : ));
569 : };
570 : result.insert(tenant_shard_id, Generation::new(g as u32));
571 : }
572 :
573 : Ok(result)
574 : }
575 :
576 : /// Reconciler calls this immediately before attaching to a new pageserver, to acquire a unique, monotonically
577 : /// advancing generation number. We also store the NodeId for which the generation was issued, so that in
578 : /// [`Self::re_attach`] we can do a bulk UPDATE on the generations for that node.
579 0 : pub(crate) async fn increment_generation(
580 0 : &self,
581 0 : tenant_shard_id: TenantShardId,
582 0 : node_id: NodeId,
583 0 : ) -> anyhow::Result<Generation> {
584 : use crate::schema::tenant_shards::dsl::*;
585 0 : let updated = self
586 0 : .with_measured_conn(DatabaseOperation::IncrementGeneration, move |conn| {
587 0 : Box::pin(async move {
588 0 : let updated = diesel::update(tenant_shards)
589 0 : .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
590 0 : .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
591 0 : .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
592 0 : .set((
593 0 : generation.eq(generation + 1),
594 0 : generation_pageserver.eq(node_id.0 as i64),
595 0 : ))
596 0 : // TODO: only returning() the generation column
597 0 : .returning(TenantShardPersistence::as_returning())
598 0 : .get_result(conn)
599 0 : .await?;
600 :
601 0 : Ok(updated)
602 0 : })
603 0 : })
604 0 : .await?;
605 :
606 : // Generation is always non-null in the rseult: if the generation column had been NULL, then we
607 : // should have experienced an SQL Confilict error while executing a query that tries to increment it.
608 0 : debug_assert!(updated.generation.is_some());
609 0 : let Some(g) = updated.generation else {
610 0 : return Err(DatabaseError::Logical(
611 0 : "Generation should always be set after incrementing".to_string(),
612 0 : )
613 0 : .into());
614 : };
615 :
616 0 : Ok(Generation::new(g as u32))
617 0 : }
618 :
619 : /// When we want to call out to the running shards for a tenant, e.g. during timeline CRUD operations,
620 : /// we need to know where the shard is attached, _and_ the generation, so that we can re-check the generation
621 : /// afterwards to confirm that our timeline CRUD operation is truly persistent (it must have happened in the
622 : /// latest generation)
623 : ///
624 : /// If the tenant doesn't exist, an empty vector is returned.
625 : ///
626 : /// Output is sorted by shard number
627 0 : pub(crate) async fn tenant_generations(
628 0 : &self,
629 0 : filter_tenant_id: TenantId,
630 0 : ) -> Result<Vec<ShardGenerationState>, DatabaseError> {
631 : use crate::schema::tenant_shards::dsl::*;
632 0 : let rows = self
633 0 : .with_measured_conn(DatabaseOperation::TenantGenerations, move |conn| {
634 0 : Box::pin(async move {
635 0 : let result = tenant_shards
636 0 : .filter(tenant_id.eq(filter_tenant_id.to_string()))
637 0 : .select(TenantShardPersistence::as_select())
638 0 : .order(shard_number)
639 0 : .load(conn)
640 0 : .await?;
641 0 : Ok(result)
642 0 : })
643 0 : })
644 0 : .await?;
645 :
646 0 : Ok(rows
647 0 : .into_iter()
648 0 : .map(|p| ShardGenerationState {
649 0 : tenant_shard_id: p
650 0 : .get_tenant_shard_id()
651 0 : .expect("Corrupt tenant shard id in database"),
652 0 : generation: p.generation.map(|g| Generation::new(g as u32)),
653 0 : generation_pageserver: p.generation_pageserver.map(|n| NodeId(n as u64)),
654 0 : })
655 0 : .collect())
656 0 : }
657 :
658 : /// Read the generation number of specific tenant shards
659 : ///
660 : /// Output is unsorted. Output may not include values for all inputs, if they are missing in the database.
661 0 : pub(crate) async fn shard_generations(
662 0 : &self,
663 0 : mut tenant_shard_ids: impl Iterator<Item = &TenantShardId>,
664 0 : ) -> Result<Vec<(TenantShardId, Option<Generation>)>, DatabaseError> {
665 0 : let mut rows = Vec::with_capacity(tenant_shard_ids.size_hint().0);
666 :
667 : // We will chunk our input to avoid composing arbitrarily long `IN` clauses. Typically we are
668 : // called with a single digit number of IDs, but in principle we could be called with tens
669 : // of thousands (all the shards on one pageserver) from the generation validation API.
670 0 : loop {
671 0 : // A modest hardcoded chunk size to handle typical cases in a single query but never generate particularly
672 0 : // large query strings.
673 0 : let chunk_ids = tenant_shard_ids.by_ref().take(32);
674 0 :
675 0 : // Compose a comma separated list of tuples for matching on (tenant_id, shard_number, shard_count)
676 0 : let in_clause = chunk_ids
677 0 : .map(|tsid| {
678 0 : format!(
679 0 : "('{}', {}, {})",
680 0 : tsid.tenant_id, tsid.shard_number.0, tsid.shard_count.0
681 0 : )
682 0 : })
683 0 : .join(",");
684 0 :
685 0 : // We are done when our iterator gives us nothing to filter on
686 0 : if in_clause.is_empty() {
687 0 : break;
688 0 : }
689 0 :
690 0 : let in_clause = &in_clause;
691 0 : let chunk_rows = self
692 0 : .with_measured_conn(DatabaseOperation::ShardGenerations, move |conn| {
693 0 : Box::pin(async move {
694 : // diesel doesn't support multi-column IN queries, so we compose raw SQL. No escaping is required because
695 : // the inputs are strongly typed and cannot carry any user-supplied raw string content.
696 0 : let result : Vec<TenantShardPersistence> = diesel::sql_query(
697 0 : format!("SELECT * from tenant_shards where (tenant_id, shard_number, shard_count) in ({in_clause});").as_str()
698 0 : ).load(conn).await?;
699 :
700 0 : Ok(result)
701 0 : })
702 0 : })
703 0 : .await?;
704 0 : rows.extend(chunk_rows.into_iter())
705 : }
706 :
707 0 : Ok(rows
708 0 : .into_iter()
709 0 : .map(|tsp| {
710 0 : (
711 0 : tsp.get_tenant_shard_id()
712 0 : .expect("Bad tenant ID in database"),
713 0 : tsp.generation.map(|g| Generation::new(g as u32)),
714 0 : )
715 0 : })
716 0 : .collect())
717 0 : }
718 :
719 : #[allow(non_local_definitions)]
720 : /// For use when updating a persistent property of a tenant, such as its config or placement_policy.
721 : ///
722 : /// Do not use this for settting generation, unless in the special onboarding code path (/location_config)
723 : /// API: use [`Self::increment_generation`] instead. Setting the generation via this route is a one-time thing
724 : /// that we only do the first time a tenant is set to an attached policy via /location_config.
725 0 : pub(crate) async fn update_tenant_shard(
726 0 : &self,
727 0 : tenant: TenantFilter,
728 0 : input_placement_policy: Option<PlacementPolicy>,
729 0 : input_config: Option<TenantConfig>,
730 0 : input_generation: Option<Generation>,
731 0 : input_scheduling_policy: Option<ShardSchedulingPolicy>,
732 0 : ) -> DatabaseResult<()> {
733 : use crate::schema::tenant_shards::dsl::*;
734 :
735 0 : let tenant = &tenant;
736 0 : let input_placement_policy = &input_placement_policy;
737 0 : let input_config = &input_config;
738 0 : let input_generation = &input_generation;
739 0 : let input_scheduling_policy = &input_scheduling_policy;
740 0 : self.with_measured_conn(DatabaseOperation::UpdateTenantShard, move |conn| {
741 0 : Box::pin(async move {
742 0 : let query = match tenant {
743 0 : TenantFilter::Shard(tenant_shard_id) => diesel::update(tenant_shards)
744 0 : .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
745 0 : .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
746 0 : .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
747 0 : .into_boxed(),
748 0 : TenantFilter::Tenant(input_tenant_id) => diesel::update(tenant_shards)
749 0 : .filter(tenant_id.eq(input_tenant_id.to_string()))
750 0 : .into_boxed(),
751 : };
752 :
753 : // Clear generation_pageserver if we are moving into a state where we won't have
754 : // any attached pageservers.
755 0 : let input_generation_pageserver = match input_placement_policy {
756 0 : None | Some(PlacementPolicy::Attached(_)) => None,
757 0 : Some(PlacementPolicy::Detached | PlacementPolicy::Secondary) => Some(None),
758 : };
759 :
760 0 : #[derive(AsChangeset)]
761 : #[diesel(table_name = crate::schema::tenant_shards)]
762 : struct ShardUpdate {
763 : generation: Option<i32>,
764 : placement_policy: Option<String>,
765 : config: Option<String>,
766 : scheduling_policy: Option<String>,
767 : generation_pageserver: Option<Option<i64>>,
768 : }
769 :
770 0 : let update = ShardUpdate {
771 0 : generation: input_generation.map(|g| g.into().unwrap() as i32),
772 0 : placement_policy: input_placement_policy
773 0 : .as_ref()
774 0 : .map(|p| serde_json::to_string(&p).unwrap()),
775 0 : config: input_config
776 0 : .as_ref()
777 0 : .map(|c| serde_json::to_string(&c).unwrap()),
778 0 : scheduling_policy: input_scheduling_policy
779 0 : .map(|p| serde_json::to_string(&p).unwrap()),
780 0 : generation_pageserver: input_generation_pageserver,
781 0 : };
782 0 :
783 0 : query.set(update).execute(conn).await?;
784 :
785 0 : Ok(())
786 0 : })
787 0 : })
788 0 : .await?;
789 :
790 0 : Ok(())
791 0 : }
792 :
793 : /// Note that passing None for a shard clears the preferred AZ (rather than leaving it unmodified)
794 0 : pub(crate) async fn set_tenant_shard_preferred_azs(
795 0 : &self,
796 0 : preferred_azs: Vec<(TenantShardId, Option<AvailabilityZone>)>,
797 0 : ) -> DatabaseResult<Vec<(TenantShardId, Option<AvailabilityZone>)>> {
798 : use crate::schema::tenant_shards::dsl::*;
799 :
800 0 : let preferred_azs = preferred_azs.as_slice();
801 0 : self.with_measured_conn(DatabaseOperation::SetPreferredAzs, move |conn| {
802 0 : Box::pin(async move {
803 0 : let mut shards_updated = Vec::default();
804 :
805 0 : for (tenant_shard_id, preferred_az) in preferred_azs.iter() {
806 0 : let updated = diesel::update(tenant_shards)
807 0 : .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
808 0 : .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
809 0 : .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
810 0 : .set(preferred_az_id.eq(preferred_az.as_ref().map(|az| az.0.clone())))
811 0 : .execute(conn)
812 0 : .await?;
813 :
814 0 : if updated == 1 {
815 0 : shards_updated.push((*tenant_shard_id, preferred_az.clone()));
816 0 : }
817 : }
818 :
819 0 : Ok(shards_updated)
820 0 : })
821 0 : })
822 0 : .await
823 0 : }
824 :
825 0 : pub(crate) async fn detach(&self, tenant_shard_id: TenantShardId) -> anyhow::Result<()> {
826 : use crate::schema::tenant_shards::dsl::*;
827 0 : self.with_measured_conn(DatabaseOperation::Detach, move |conn| {
828 0 : Box::pin(async move {
829 0 : let updated = diesel::update(tenant_shards)
830 0 : .filter(tenant_id.eq(tenant_shard_id.tenant_id.to_string()))
831 0 : .filter(shard_number.eq(tenant_shard_id.shard_number.0 as i32))
832 0 : .filter(shard_count.eq(tenant_shard_id.shard_count.literal() as i32))
833 0 : .set((
834 0 : generation_pageserver.eq(Option::<i64>::None),
835 0 : placement_policy
836 0 : .eq(serde_json::to_string(&PlacementPolicy::Detached).unwrap()),
837 0 : ))
838 0 : .execute(conn)
839 0 : .await?;
840 :
841 0 : Ok(updated)
842 0 : })
843 0 : })
844 0 : .await?;
845 :
846 0 : Ok(())
847 0 : }
848 :
849 : // When we start shard splitting, we must durably mark the tenant so that
850 : // on restart, we know that we must go through recovery.
851 : //
852 : // We create the child shards here, so that they will be available for increment_generation calls
853 : // if some pageserver holding a child shard needs to restart before the overall tenant split is complete.
854 0 : pub(crate) async fn begin_shard_split(
855 0 : &self,
856 0 : old_shard_count: ShardCount,
857 0 : split_tenant_id: TenantId,
858 0 : parent_to_children: Vec<(TenantShardId, Vec<TenantShardPersistence>)>,
859 0 : ) -> DatabaseResult<()> {
860 : use crate::schema::tenant_shards::dsl::*;
861 0 : let parent_to_children = parent_to_children.as_slice();
862 0 : self.with_measured_conn(DatabaseOperation::BeginShardSplit, move |conn| {
863 0 : Box::pin(async move {
864 : // Mark parent shards as splitting
865 :
866 0 : let updated = diesel::update(tenant_shards)
867 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
868 0 : .filter(shard_count.eq(old_shard_count.literal() as i32))
869 0 : .set((splitting.eq(1),))
870 0 : .execute(conn).await?;
871 0 : if u8::try_from(updated)
872 0 : .map_err(|_| DatabaseError::Logical(
873 0 : format!("Overflow existing shard count {} while splitting", updated))
874 0 : )? != old_shard_count.count() {
875 : // Perhaps a deletion or another split raced with this attempt to split, mutating
876 : // the parent shards that we intend to split. In this case the split request should fail.
877 0 : return Err(DatabaseError::Logical(
878 0 : format!("Unexpected existing shard count {updated} when preparing tenant for split (expected {})", old_shard_count.count())
879 0 : ));
880 0 : }
881 0 :
882 0 : // FIXME: spurious clone to sidestep closure move rules
883 0 : let parent_to_children = parent_to_children.to_vec();
884 :
885 : // Insert child shards
886 0 : for (parent_shard_id, children) in parent_to_children {
887 0 : let mut parent = crate::schema::tenant_shards::table
888 0 : .filter(tenant_id.eq(parent_shard_id.tenant_id.to_string()))
889 0 : .filter(shard_number.eq(parent_shard_id.shard_number.0 as i32))
890 0 : .filter(shard_count.eq(parent_shard_id.shard_count.literal() as i32))
891 0 : .load::<TenantShardPersistence>(conn).await?;
892 0 : let parent = if parent.len() != 1 {
893 0 : return Err(DatabaseError::Logical(format!(
894 0 : "Parent shard {parent_shard_id} not found"
895 0 : )));
896 : } else {
897 0 : parent.pop().unwrap()
898 : };
899 0 : for mut shard in children {
900 : // Carry the parent's generation into the child
901 0 : shard.generation = parent.generation;
902 0 :
903 0 : debug_assert!(shard.splitting == SplitState::Splitting);
904 0 : diesel::insert_into(tenant_shards)
905 0 : .values(shard)
906 0 : .execute(conn).await?;
907 : }
908 : }
909 :
910 0 : Ok(())
911 0 : })
912 0 : })
913 0 : .await
914 0 : }
915 :
916 : // When we finish shard splitting, we must atomically clean up the old shards
917 : // and insert the new shards, and clear the splitting marker.
918 0 : pub(crate) async fn complete_shard_split(
919 0 : &self,
920 0 : split_tenant_id: TenantId,
921 0 : old_shard_count: ShardCount,
922 0 : ) -> DatabaseResult<()> {
923 : use crate::schema::tenant_shards::dsl::*;
924 0 : self.with_measured_conn(DatabaseOperation::CompleteShardSplit, move |conn| {
925 0 : Box::pin(async move {
926 0 : // Drop parent shards
927 0 : diesel::delete(tenant_shards)
928 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
929 0 : .filter(shard_count.eq(old_shard_count.literal() as i32))
930 0 : .execute(conn)
931 0 : .await?;
932 :
933 : // Clear sharding flag
934 0 : let updated = diesel::update(tenant_shards)
935 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
936 0 : .set((splitting.eq(0),))
937 0 : .execute(conn)
938 0 : .await?;
939 0 : debug_assert!(updated > 0);
940 :
941 0 : Ok(())
942 0 : })
943 0 : })
944 0 : .await
945 0 : }
946 :
947 : /// Used when the remote part of a shard split failed: we will revert the database state to have only
948 : /// the parent shards, with SplitState::Idle.
949 0 : pub(crate) async fn abort_shard_split(
950 0 : &self,
951 0 : split_tenant_id: TenantId,
952 0 : new_shard_count: ShardCount,
953 0 : ) -> DatabaseResult<AbortShardSplitStatus> {
954 : use crate::schema::tenant_shards::dsl::*;
955 0 : self.with_measured_conn(DatabaseOperation::AbortShardSplit, move |conn| {
956 0 : Box::pin(async move {
957 : // Clear the splitting state on parent shards
958 0 : let updated = diesel::update(tenant_shards)
959 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
960 0 : .filter(shard_count.ne(new_shard_count.literal() as i32))
961 0 : .set((splitting.eq(0),))
962 0 : .execute(conn)
963 0 : .await?;
964 :
965 : // Parent shards are already gone: we cannot abort.
966 0 : if updated == 0 {
967 0 : return Ok(AbortShardSplitStatus::Complete);
968 0 : }
969 0 :
970 0 : // Sanity check: if parent shards were present, their cardinality should
971 0 : // be less than the number of child shards.
972 0 : if updated >= new_shard_count.count() as usize {
973 0 : return Err(DatabaseError::Logical(format!(
974 0 : "Unexpected parent shard count {updated} while aborting split to \
975 0 : count {new_shard_count:?} on tenant {split_tenant_id}"
976 0 : )));
977 0 : }
978 0 :
979 0 : // Erase child shards
980 0 : diesel::delete(tenant_shards)
981 0 : .filter(tenant_id.eq(split_tenant_id.to_string()))
982 0 : .filter(shard_count.eq(new_shard_count.literal() as i32))
983 0 : .execute(conn)
984 0 : .await?;
985 :
986 0 : Ok(AbortShardSplitStatus::Aborted)
987 0 : })
988 0 : })
989 0 : .await
990 0 : }
991 :
992 : /// Stores all the latest metadata health updates durably. Updates existing entry on conflict.
993 : ///
994 : /// **Correctness:** `metadata_health_updates` should all belong the tenant shards managed by the storage controller.
995 : #[allow(dead_code)]
996 0 : pub(crate) async fn update_metadata_health_records(
997 0 : &self,
998 0 : healthy_records: Vec<MetadataHealthPersistence>,
999 0 : unhealthy_records: Vec<MetadataHealthPersistence>,
1000 0 : now: chrono::DateTime<chrono::Utc>,
1001 0 : ) -> DatabaseResult<()> {
1002 : use crate::schema::metadata_health::dsl::*;
1003 :
1004 0 : let healthy_records = healthy_records.as_slice();
1005 0 : let unhealthy_records = unhealthy_records.as_slice();
1006 0 : self.with_measured_conn(DatabaseOperation::UpdateMetadataHealth, move |conn| {
1007 0 : Box::pin(async move {
1008 0 : diesel::insert_into(metadata_health)
1009 0 : .values(healthy_records)
1010 0 : .on_conflict((tenant_id, shard_number, shard_count))
1011 0 : .do_update()
1012 0 : .set((healthy.eq(true), last_scrubbed_at.eq(now)))
1013 0 : .execute(conn)
1014 0 : .await?;
1015 :
1016 0 : diesel::insert_into(metadata_health)
1017 0 : .values(unhealthy_records)
1018 0 : .on_conflict((tenant_id, shard_number, shard_count))
1019 0 : .do_update()
1020 0 : .set((healthy.eq(false), last_scrubbed_at.eq(now)))
1021 0 : .execute(conn)
1022 0 : .await?;
1023 0 : Ok(())
1024 0 : })
1025 0 : })
1026 0 : .await
1027 0 : }
1028 :
1029 : /// Lists all the metadata health records.
1030 : #[allow(dead_code)]
1031 0 : pub(crate) async fn list_metadata_health_records(
1032 0 : &self,
1033 0 : ) -> DatabaseResult<Vec<MetadataHealthPersistence>> {
1034 0 : self.with_measured_conn(DatabaseOperation::ListMetadataHealth, move |conn| {
1035 0 : Box::pin(async {
1036 0 : Ok(crate::schema::metadata_health::table
1037 0 : .load::<MetadataHealthPersistence>(conn)
1038 0 : .await?)
1039 0 : })
1040 0 : })
1041 0 : .await
1042 0 : }
1043 :
1044 : /// Lists all the metadata health records that is unhealthy.
1045 : #[allow(dead_code)]
1046 0 : pub(crate) async fn list_unhealthy_metadata_health_records(
1047 0 : &self,
1048 0 : ) -> DatabaseResult<Vec<MetadataHealthPersistence>> {
1049 : use crate::schema::metadata_health::dsl::*;
1050 0 : self.with_measured_conn(
1051 0 : DatabaseOperation::ListMetadataHealthUnhealthy,
1052 0 : move |conn| {
1053 0 : Box::pin(async {
1054 0 : DatabaseResult::Ok(
1055 0 : crate::schema::metadata_health::table
1056 0 : .filter(healthy.eq(false))
1057 0 : .load::<MetadataHealthPersistence>(conn)
1058 0 : .await?,
1059 : )
1060 0 : })
1061 0 : },
1062 0 : )
1063 0 : .await
1064 0 : }
1065 :
1066 : /// Lists all the metadata health records that have not been updated since an `earlier` time.
1067 : #[allow(dead_code)]
1068 0 : pub(crate) async fn list_outdated_metadata_health_records(
1069 0 : &self,
1070 0 : earlier: chrono::DateTime<chrono::Utc>,
1071 0 : ) -> DatabaseResult<Vec<MetadataHealthPersistence>> {
1072 : use crate::schema::metadata_health::dsl::*;
1073 :
1074 0 : self.with_measured_conn(DatabaseOperation::ListMetadataHealthOutdated, move |conn| {
1075 0 : Box::pin(async move {
1076 0 : let query = metadata_health.filter(last_scrubbed_at.lt(earlier));
1077 0 : let res = query.load::<MetadataHealthPersistence>(conn).await?;
1078 :
1079 0 : Ok(res)
1080 0 : })
1081 0 : })
1082 0 : .await
1083 0 : }
1084 :
1085 : /// Get the current entry from the `leader` table if one exists.
1086 : /// It is an error for the table to contain more than one entry.
1087 0 : pub(crate) async fn get_leader(&self) -> DatabaseResult<Option<ControllerPersistence>> {
1088 0 : let mut leader: Vec<ControllerPersistence> = self
1089 0 : .with_measured_conn(DatabaseOperation::GetLeader, move |conn| {
1090 0 : Box::pin(async move {
1091 0 : Ok(crate::schema::controllers::table
1092 0 : .load::<ControllerPersistence>(conn)
1093 0 : .await?)
1094 0 : })
1095 0 : })
1096 0 : .await?;
1097 :
1098 0 : if leader.len() > 1 {
1099 0 : return Err(DatabaseError::Logical(format!(
1100 0 : "More than one entry present in the leader table: {leader:?}"
1101 0 : )));
1102 0 : }
1103 0 :
1104 0 : Ok(leader.pop())
1105 0 : }
1106 :
1107 : /// Update the new leader with compare-exchange semantics. If `prev` does not
1108 : /// match the current leader entry, then the update is treated as a failure.
1109 : /// When `prev` is not specified, the update is forced.
1110 0 : pub(crate) async fn update_leader(
1111 0 : &self,
1112 0 : prev: Option<ControllerPersistence>,
1113 0 : new: ControllerPersistence,
1114 0 : ) -> DatabaseResult<()> {
1115 : use crate::schema::controllers::dsl::*;
1116 :
1117 0 : let updated = self
1118 0 : .with_measured_conn(DatabaseOperation::UpdateLeader, move |conn| {
1119 0 : let prev = prev.clone();
1120 0 : let new = new.clone();
1121 0 : Box::pin(async move {
1122 0 : let updated = match &prev {
1123 0 : Some(prev) => {
1124 0 : diesel::update(controllers)
1125 0 : .filter(address.eq(prev.address.clone()))
1126 0 : .filter(started_at.eq(prev.started_at))
1127 0 : .set((
1128 0 : address.eq(new.address.clone()),
1129 0 : started_at.eq(new.started_at),
1130 0 : ))
1131 0 : .execute(conn)
1132 0 : .await?
1133 : }
1134 : None => {
1135 0 : diesel::insert_into(controllers)
1136 0 : .values(new.clone())
1137 0 : .execute(conn)
1138 0 : .await?
1139 : }
1140 : };
1141 :
1142 0 : Ok(updated)
1143 0 : })
1144 0 : })
1145 0 : .await?;
1146 :
1147 0 : if updated == 0 {
1148 0 : return Err(DatabaseError::Logical(
1149 0 : "Leader table update failed".to_string(),
1150 0 : ));
1151 0 : }
1152 0 :
1153 0 : Ok(())
1154 0 : }
1155 :
1156 : /// At startup, populate the list of nodes which our shards may be placed on
1157 0 : pub(crate) async fn list_safekeepers(&self) -> DatabaseResult<Vec<SafekeeperPersistence>> {
1158 0 : let safekeepers: Vec<SafekeeperPersistence> = self
1159 0 : .with_measured_conn(DatabaseOperation::ListNodes, move |conn| {
1160 0 : Box::pin(async move {
1161 0 : Ok(crate::schema::safekeepers::table
1162 0 : .load::<SafekeeperPersistence>(conn)
1163 0 : .await?)
1164 0 : })
1165 0 : })
1166 0 : .await?;
1167 :
1168 0 : tracing::info!("list_safekeepers: loaded {} nodes", safekeepers.len());
1169 :
1170 0 : Ok(safekeepers)
1171 0 : }
1172 :
1173 0 : pub(crate) async fn safekeeper_get(
1174 0 : &self,
1175 0 : id: i64,
1176 0 : ) -> Result<SafekeeperPersistence, DatabaseError> {
1177 : use crate::schema::safekeepers::dsl::{id as id_column, safekeepers};
1178 0 : self.with_conn(move |conn| {
1179 0 : Box::pin(async move {
1180 0 : Ok(safekeepers
1181 0 : .filter(id_column.eq(&id))
1182 0 : .select(SafekeeperPersistence::as_select())
1183 0 : .get_result(conn)
1184 0 : .await?)
1185 0 : })
1186 0 : })
1187 0 : .await
1188 0 : }
1189 :
1190 0 : pub(crate) async fn safekeeper_upsert(
1191 0 : &self,
1192 0 : record: SafekeeperUpsert,
1193 0 : ) -> Result<(), DatabaseError> {
1194 : use crate::schema::safekeepers::dsl::*;
1195 :
1196 0 : self.with_conn(move |conn| {
1197 0 : let record = record.clone();
1198 0 : Box::pin(async move {
1199 0 : let bind = record
1200 0 : .as_insert_or_update()
1201 0 : .map_err(|e| DatabaseError::Logical(format!("{e}")))?;
1202 :
1203 0 : let inserted_updated = diesel::insert_into(safekeepers)
1204 0 : .values(&bind)
1205 0 : .on_conflict(id)
1206 0 : .do_update()
1207 0 : .set(&bind)
1208 0 : .execute(conn)
1209 0 : .await?;
1210 :
1211 0 : if inserted_updated != 1 {
1212 0 : return Err(DatabaseError::Logical(format!(
1213 0 : "unexpected number of rows ({})",
1214 0 : inserted_updated
1215 0 : )));
1216 0 : }
1217 0 :
1218 0 : Ok(())
1219 0 : })
1220 0 : })
1221 0 : .await
1222 0 : }
1223 :
1224 0 : pub(crate) async fn set_safekeeper_scheduling_policy(
1225 0 : &self,
1226 0 : id_: i64,
1227 0 : scheduling_policy_: SkSchedulingPolicy,
1228 0 : ) -> Result<(), DatabaseError> {
1229 : use crate::schema::safekeepers::dsl::*;
1230 :
1231 0 : self.with_conn(move |conn| {
1232 0 : Box::pin(async move {
1233 0 : #[derive(Insertable, AsChangeset)]
1234 : #[diesel(table_name = crate::schema::safekeepers)]
1235 : struct UpdateSkSchedulingPolicy<'a> {
1236 : id: i64,
1237 : scheduling_policy: &'a str,
1238 : }
1239 0 : let scheduling_policy_ = String::from(scheduling_policy_);
1240 :
1241 0 : let rows_affected = diesel::update(safekeepers.filter(id.eq(id_)))
1242 0 : .set(scheduling_policy.eq(scheduling_policy_))
1243 0 : .execute(conn)
1244 0 : .await?;
1245 :
1246 0 : if rows_affected != 1 {
1247 0 : return Err(DatabaseError::Logical(format!(
1248 0 : "unexpected number of rows ({rows_affected})",
1249 0 : )));
1250 0 : }
1251 0 :
1252 0 : Ok(())
1253 0 : })
1254 0 : })
1255 0 : .await
1256 0 : }
1257 : }
1258 :
1259 : /// Parts of [`crate::tenant_shard::TenantShard`] that are stored durably
1260 : #[derive(
1261 0 : QueryableByName, Queryable, Selectable, Insertable, Serialize, Deserialize, Clone, Eq, PartialEq,
1262 : )]
1263 : #[diesel(table_name = crate::schema::tenant_shards)]
1264 : pub(crate) struct TenantShardPersistence {
1265 : #[serde(default)]
1266 : pub(crate) tenant_id: String,
1267 : #[serde(default)]
1268 : pub(crate) shard_number: i32,
1269 : #[serde(default)]
1270 : pub(crate) shard_count: i32,
1271 : #[serde(default)]
1272 : pub(crate) shard_stripe_size: i32,
1273 :
1274 : // Latest generation number: next time we attach, increment this
1275 : // and use the incremented number when attaching.
1276 : //
1277 : // Generation is only None when first onboarding a tenant, where it may
1278 : // be in PlacementPolicy::Secondary and therefore have no valid generation state.
1279 : pub(crate) generation: Option<i32>,
1280 :
1281 : // Currently attached pageserver
1282 : #[serde(rename = "pageserver")]
1283 : pub(crate) generation_pageserver: Option<i64>,
1284 :
1285 : #[serde(default)]
1286 : pub(crate) placement_policy: String,
1287 : #[serde(default)]
1288 : pub(crate) splitting: SplitState,
1289 : #[serde(default)]
1290 : pub(crate) config: String,
1291 : #[serde(default)]
1292 : pub(crate) scheduling_policy: String,
1293 :
1294 : // Hint that we should attempt to schedule this tenant shard the given
1295 : // availability zone in order to minimise the chances of cross-AZ communication
1296 : // with compute.
1297 : pub(crate) preferred_az_id: Option<String>,
1298 : }
1299 :
1300 : impl TenantShardPersistence {
1301 0 : pub(crate) fn get_shard_identity(&self) -> Result<ShardIdentity, ShardConfigError> {
1302 0 : if self.shard_count == 0 {
1303 0 : Ok(ShardIdentity::unsharded())
1304 : } else {
1305 0 : Ok(ShardIdentity::new(
1306 0 : ShardNumber(self.shard_number as u8),
1307 0 : ShardCount::new(self.shard_count as u8),
1308 0 : ShardStripeSize(self.shard_stripe_size as u32),
1309 0 : )?)
1310 : }
1311 0 : }
1312 :
1313 0 : pub(crate) fn get_tenant_shard_id(&self) -> Result<TenantShardId, hex::FromHexError> {
1314 0 : Ok(TenantShardId {
1315 0 : tenant_id: TenantId::from_str(self.tenant_id.as_str())?,
1316 0 : shard_number: ShardNumber(self.shard_number as u8),
1317 0 : shard_count: ShardCount::new(self.shard_count as u8),
1318 : })
1319 0 : }
1320 : }
1321 :
1322 : /// Parts of [`crate::node::Node`] that are stored durably
1323 0 : #[derive(Serialize, Deserialize, Queryable, Selectable, Insertable, Eq, PartialEq)]
1324 : #[diesel(table_name = crate::schema::nodes)]
1325 : pub(crate) struct NodePersistence {
1326 : pub(crate) node_id: i64,
1327 : pub(crate) scheduling_policy: String,
1328 : pub(crate) listen_http_addr: String,
1329 : pub(crate) listen_http_port: i32,
1330 : pub(crate) listen_pg_addr: String,
1331 : pub(crate) listen_pg_port: i32,
1332 : pub(crate) availability_zone_id: String,
1333 : }
1334 :
1335 : /// Tenant metadata health status that are stored durably.
1336 0 : #[derive(Queryable, Selectable, Insertable, Serialize, Deserialize, Clone, Eq, PartialEq)]
1337 : #[diesel(table_name = crate::schema::metadata_health)]
1338 : pub(crate) struct MetadataHealthPersistence {
1339 : #[serde(default)]
1340 : pub(crate) tenant_id: String,
1341 : #[serde(default)]
1342 : pub(crate) shard_number: i32,
1343 : #[serde(default)]
1344 : pub(crate) shard_count: i32,
1345 :
1346 : pub(crate) healthy: bool,
1347 : pub(crate) last_scrubbed_at: chrono::DateTime<chrono::Utc>,
1348 : }
1349 :
1350 : impl MetadataHealthPersistence {
1351 0 : pub fn new(
1352 0 : tenant_shard_id: TenantShardId,
1353 0 : healthy: bool,
1354 0 : last_scrubbed_at: chrono::DateTime<chrono::Utc>,
1355 0 : ) -> Self {
1356 0 : let tenant_id = tenant_shard_id.tenant_id.to_string();
1357 0 : let shard_number = tenant_shard_id.shard_number.0 as i32;
1358 0 : let shard_count = tenant_shard_id.shard_count.literal() as i32;
1359 0 :
1360 0 : MetadataHealthPersistence {
1361 0 : tenant_id,
1362 0 : shard_number,
1363 0 : shard_count,
1364 0 : healthy,
1365 0 : last_scrubbed_at,
1366 0 : }
1367 0 : }
1368 :
1369 : #[allow(dead_code)]
1370 0 : pub(crate) fn get_tenant_shard_id(&self) -> Result<TenantShardId, hex::FromHexError> {
1371 0 : Ok(TenantShardId {
1372 0 : tenant_id: TenantId::from_str(self.tenant_id.as_str())?,
1373 0 : shard_number: ShardNumber(self.shard_number as u8),
1374 0 : shard_count: ShardCount::new(self.shard_count as u8),
1375 : })
1376 0 : }
1377 : }
1378 :
1379 : impl From<MetadataHealthPersistence> for MetadataHealthRecord {
1380 0 : fn from(value: MetadataHealthPersistence) -> Self {
1381 0 : MetadataHealthRecord {
1382 0 : tenant_shard_id: value
1383 0 : .get_tenant_shard_id()
1384 0 : .expect("stored tenant id should be valid"),
1385 0 : healthy: value.healthy,
1386 0 : last_scrubbed_at: value.last_scrubbed_at,
1387 0 : }
1388 0 : }
1389 : }
1390 :
1391 : #[derive(
1392 0 : Serialize, Deserialize, Queryable, Selectable, Insertable, Eq, PartialEq, Debug, Clone,
1393 : )]
1394 : #[diesel(table_name = crate::schema::controllers)]
1395 : pub(crate) struct ControllerPersistence {
1396 : pub(crate) address: String,
1397 : pub(crate) started_at: chrono::DateTime<chrono::Utc>,
1398 : }
1399 :
1400 : // What we store in the database
1401 0 : #[derive(Serialize, Deserialize, Queryable, Selectable, Eq, PartialEq, Debug, Clone)]
1402 : #[diesel(table_name = crate::schema::safekeepers)]
1403 : pub(crate) struct SafekeeperPersistence {
1404 : pub(crate) id: i64,
1405 : pub(crate) region_id: String,
1406 : /// 1 is special, it means just created (not currently posted to storcon).
1407 : /// Zero or negative is not really expected.
1408 : /// Otherwise the number from `release-$(number_of_commits_on_branch)` tag.
1409 : pub(crate) version: i64,
1410 : pub(crate) host: String,
1411 : pub(crate) port: i32,
1412 : pub(crate) http_port: i32,
1413 : pub(crate) availability_zone_id: String,
1414 : pub(crate) scheduling_policy: String,
1415 : }
1416 :
1417 : impl SafekeeperPersistence {
1418 0 : pub(crate) fn as_describe_response(&self) -> Result<SafekeeperDescribeResponse, DatabaseError> {
1419 0 : let scheduling_policy =
1420 0 : SkSchedulingPolicy::from_str(&self.scheduling_policy).map_err(|e| {
1421 0 : DatabaseError::Logical(format!("can't construct SkSchedulingPolicy: {e:?}"))
1422 0 : })?;
1423 0 : Ok(SafekeeperDescribeResponse {
1424 0 : id: NodeId(self.id as u64),
1425 0 : region_id: self.region_id.clone(),
1426 0 : version: self.version,
1427 0 : host: self.host.clone(),
1428 0 : port: self.port,
1429 0 : http_port: self.http_port,
1430 0 : availability_zone_id: self.availability_zone_id.clone(),
1431 0 : scheduling_policy,
1432 0 : })
1433 0 : }
1434 : }
1435 :
1436 : /// What we expect from the upsert http api
1437 0 : #[derive(Serialize, Deserialize, Eq, PartialEq, Debug, Clone)]
1438 : pub(crate) struct SafekeeperUpsert {
1439 : pub(crate) id: i64,
1440 : pub(crate) region_id: String,
1441 : /// 1 is special, it means just created (not currently posted to storcon).
1442 : /// Zero or negative is not really expected.
1443 : /// Otherwise the number from `release-$(number_of_commits_on_branch)` tag.
1444 : pub(crate) version: i64,
1445 : pub(crate) host: String,
1446 : pub(crate) port: i32,
1447 : /// The active flag will not be stored in the database and will be ignored.
1448 : pub(crate) active: Option<bool>,
1449 : pub(crate) http_port: i32,
1450 : pub(crate) availability_zone_id: String,
1451 : }
1452 :
1453 : impl SafekeeperUpsert {
1454 0 : fn as_insert_or_update(&self) -> anyhow::Result<InsertUpdateSafekeeper<'_>> {
1455 0 : if self.version < 0 {
1456 0 : anyhow::bail!("negative version: {}", self.version);
1457 0 : }
1458 0 : Ok(InsertUpdateSafekeeper {
1459 0 : id: self.id,
1460 0 : region_id: &self.region_id,
1461 0 : version: self.version,
1462 0 : host: &self.host,
1463 0 : port: self.port,
1464 0 : http_port: self.http_port,
1465 0 : availability_zone_id: &self.availability_zone_id,
1466 0 : // None means a wish to not update this column. We expose abilities to update it via other means.
1467 0 : scheduling_policy: None,
1468 0 : })
1469 0 : }
1470 : }
1471 :
1472 0 : #[derive(Insertable, AsChangeset)]
1473 : #[diesel(table_name = crate::schema::safekeepers)]
1474 : struct InsertUpdateSafekeeper<'a> {
1475 : id: i64,
1476 : region_id: &'a str,
1477 : version: i64,
1478 : host: &'a str,
1479 : port: i32,
1480 : http_port: i32,
1481 : availability_zone_id: &'a str,
1482 : scheduling_policy: Option<&'a str>,
1483 : }
|
